From 13d59beaf64de42d500c387fad92cb88c12ba42f Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:23:03 +0100 Subject: [PATCH 1/6] Point the AgentServer packages at the local preview drop The durable state-store API this branch is built on ships in Core beta.28, which is not on nuget.org yet. The local feed is a stopgap for developing against it and must be removed before this branch ships. --- dotnet/Directory.Packages.props | 6 +++--- dotnet/nuget.config | 4 ++++ .../Microsoft.Agents.AI.Foundry.Hosting.csproj | 2 ++ 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index ac7eb21424..ba7f5601c5 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -23,9 +23,9 @@ - - - + + + diff --git a/dotnet/nuget.config b/dotnet/nuget.config index 128d95e590..3cde656b2a 100644 --- a/dotnet/nuget.config +++ b/dotnet/nuget.config @@ -3,10 +3,14 @@ + + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/Microsoft.Agents.AI.Foundry.Hosting.csproj b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/Microsoft.Agents.AI.Foundry.Hosting.csproj index e1d1cf5a67..0012b388a5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/Microsoft.Agents.AI.Foundry.Hosting.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/Microsoft.Agents.AI.Foundry.Hosting.csproj @@ -31,8 +31,10 @@ + + From f5f4eba3cbaf4384ea14aa4703c6aaa0c1350f50 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:23:29 +0100 Subject: [PATCH 2/6] Keep hosted agent state on the platform instead of the container disk A hosted agent kept its sessions, and a hosted workflow its checkpoints, in files under the container's own directory. That state is lost when the container is replaced and cannot be read by another instance of the same agent, so a conversation could not survive a restart or be served by more than one instance. Both now go to the Foundry durable state store when the process runs in a Foundry container, and stay on disk everywhere else: - FoundryAgentSessionStore holds the agent sessions, partitioned by agent, conversation and end user. - FoundryJsonCheckpointStore holds the workflow checkpoints, one item per checkpoint plus a per-session index that keeps them in commit order. Retrieving a checkpoint deletes the rest of that session's checkpoints, which is the only point at which nothing can still reach them, and is what stops the index growing past the size the platform accepts for one item. A workflow agent is redirected to that checkpoint store when it is resolved for a request, so nothing changes in how a container registers one. An agent built with a checkpoint manager of its own is left alone and reported by the new foundry-workflow-checkpointing readiness check, because its state would go somewhere hosting does not manage. Workflow agents are recognised through a new WorkflowAgentMetadata returned by GetService, which still finds them behind middleware. --- .../AgentFrameworkResponseHandler.cs | 11 +- .../FoundryAgentSessionStore.cs | 223 ++++++++ .../FoundryJsonCheckpointStore.cs | 523 ++++++++++++++++++ .../FoundryStateStoreBinding.cs | 121 ++++ .../HostedStoredOutputHealthCheck.cs | 18 +- .../HostedWorkflowCheckpointingHealthCheck.cs | 95 ++++ .../ServiceCollectionExtensions.cs | 173 +++++- .../WorkflowAgentMetadata.cs | 52 ++ .../WorkflowHostAgent.cs | 42 ++ .../WorkflowHostingExtensions.cs | 39 ++ .../FoundryAgentSessionStoreTests.cs | 312 +++++++++++ .../FoundryJsonCheckpointStoreTests.cs | 434 +++++++++++++++ .../HostedOutboundUserAgentTests.cs | 2 +- ...edWorkflowCheckpointingHealthCheckTests.cs | 135 +++++ .../WorkflowHostingExtensionsTests.cs | 183 ++++++ 15 files changed, 2324 insertions(+), 39 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowAgentMetadata.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index fe2889b444..d9cbfc8d7b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -13,6 +13,7 @@ using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; // The terminal stream events are named the same in two namespaces this file pulls in, and the short // name binds to the one the event objects are not. Naming them here keeps `is` checks against the @@ -53,8 +54,8 @@ public AgentFrameworkResponseHandler( ILogger logger, FoundryToolboxService? toolboxService = null) { - ArgumentNullException.ThrowIfNull(serviceProvider); - ArgumentNullException.ThrowIfNull(logger); + _ = Throw.IfNull(serviceProvider); + _ = Throw.IfNull(logger); this._serviceProvider = serviceProvider; this._logger = logger; @@ -581,7 +582,8 @@ private AIAgent ResolveAgent(CreateResponse request) if (agent is not null) { FoundryHostingExtensions.TryApplyUserAgent(agent); - return FoundryHostingExtensions.ApplyOpenTelemetry(agent); + return FoundryHostingExtensions.ApplyOpenTelemetry( + FoundryHostingExtensions.ApplyWorkflowCheckpointing(agent, this._serviceProvider.GetService())); } if (this._logger.IsEnabled(LogLevel.Warning)) @@ -595,7 +597,8 @@ private AIAgent ResolveAgent(CreateResponse request) if (defaultAgent is not null) { FoundryHostingExtensions.TryApplyUserAgent(defaultAgent); - return FoundryHostingExtensions.ApplyOpenTelemetry(defaultAgent); + return FoundryHostingExtensions.ApplyOpenTelemetry( + FoundryHostingExtensions.ApplyWorkflowCheckpointing(defaultAgent, this._serviceProvider.GetService())); } var errorMessage = string.IsNullOrEmpty(agentName) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs new file mode 100644 index 0000000000..73b826390c --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs @@ -0,0 +1,223 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Security.Cryptography; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Core.Storage; +using Azure.Core; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Foundry.Hosting; + +/// +/// Provides an that persists the agent-framework's serialized +/// state to the Foundry platform's durable state-store API +/// () instead of to the container's local disk. +/// +/// +/// +/// This is the durable counterpart to . The file-system +/// store writes under the container's session directory, which belongs to a single container +/// instance: the state is lost when that container is replaced and cannot be read by another +/// instance. This store writes to the platform instead, so a session survives container restarts +/// and replacement and is visible to every instance of the agent. +/// +/// +/// Layout. All sessions live in one state store, named unless +/// overridden, and each (agent, user, conversation) triple is one item in it. The item key is a +/// hash of the same a-/u-/c- logical key that +/// and use, so +/// all three stores partition sessions identically. Hashing is required because the platform +/// limits an item key to 128 characters, which an agent name plus a user id plus a conversation id +/// can exceed. The readable logical key is stored alongside the session in the item body so a +/// stored item can still be traced back to its conversation. +/// +/// +/// Per-user isolation is expressed through the item key rather than through the state store's own +/// userIsolation option. That option is fixed when the store is created and resolves the +/// user from the calling identity, whereas the user id handled here arrives per request and the +/// container always calls the storage API with its own identity. +/// +/// +/// The bound state store is resolved once, on first use, and reused for the lifetime of this +/// instance. Resolving it costs one round trip (plus one more the very first time, to create the +/// store), so it deliberately does not happen per request. +/// +/// +[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] +public sealed class FoundryAgentSessionStore : AgentSessionStore +{ + /// + /// The default state-store name used to hold every agent session persisted by this store. + /// + public const string DefaultStoreName = "agent-framework/sessions"; + + /// The item-body field holding the serialized session JSON. + private const string SessionField = "session"; + + /// The item-body field holding the readable logical key, for traceability. + private const string KeyField = "key"; + + private readonly FoundryStateStoreBinding _binding; + + /// + /// Initializes a new instance of the class. + /// + /// The credential used to authenticate to the Foundry storage API. + /// + /// The Foundry project endpoint. When , it is read from the + /// FOUNDRY_PROJECT_ENDPOINT environment variable, which the platform sets in a hosted + /// container. + /// + /// The state-store name to hold the sessions. Defaults to . + /// + /// How long a session survives without being written, in seconds. Defaults to the platform + /// default of 30 days; -1 means never expire. A write renews the window, a read does + /// not. The value only takes effect when this store is created for the first time, because the + /// platform fixes it at creation. + /// + public FoundryAgentSessionStore( + TokenCredential credential, + Uri? endpoint = null, + string storeName = DefaultStoreName, + int itemTtlSeconds = FoundryStateStore.DefaultItemTtlSeconds) + { + _ = Throw.IfNull(credential); + _ = Throw.IfNullOrWhitespace(storeName); + + this.StoreName = storeName; + this._binding = new(cancellationToken => FoundryStateStore.GetOrCreateAsync( + storeName, + credential, + endpoint, + description: "Agent Framework hosted agent sessions.", + itemTtlSeconds: itemTtlSeconds, + cancellationToken: cancellationToken)); + } + + /// + /// Initializes a new instance of the class over a + /// caller-supplied state store. Used by tests to substitute the platform client. + /// + /// Resolves the bound state store on first use. + /// The state-store name, for diagnostics. + internal FoundryAgentSessionStore(Func> storeFactory, string storeName = DefaultStoreName) + { + _ = Throw.IfNull(storeFactory); + + this._binding = new(storeFactory); + this.StoreName = storeName; + } + + /// Gets the state-store name that holds the sessions. + public string StoreName { get; } + + /// + public override async ValueTask SaveSessionAsync( + AIAgent agent, + string conversationId, + AgentSession session, + string? userId, + CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(agent); + _ = Throw.IfNullOrWhitespace(conversationId); + _ = Throw.IfNull(session); + + JsonElement serialized = await agent.SerializeSessionAsync(session, cancellationToken: cancellationToken).ConfigureAwait(false); + BinaryData sessionData = ToBinaryData(serialized); + + string logicalKey = BuildLogicalKey(agent, conversationId, userId); + FoundryStateStore store = await this.GetStoreAsync(cancellationToken).ConfigureAwait(false); + + await store.SetItemAsync( + BuildItemKey(logicalKey), + new Dictionary + { + [SessionField] = sessionData, + [KeyField] = ToJsonString(logicalKey), + }, + cancellationToken: cancellationToken).ConfigureAwait(false); + } + + /// + public override async ValueTask GetSessionAsync( + AIAgent agent, + string conversationId, + string? userId, + CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(agent); + _ = Throw.IfNullOrWhitespace(conversationId); + + string logicalKey = BuildLogicalKey(agent, conversationId, userId); + FoundryStateStore store = await this.GetStoreAsync(cancellationToken).ConfigureAwait(false); + + // GetItemAsync already answers null for an item that is not there, which is exactly the + // "nothing stored" result this method contracts to return. + StateStoreItem? item = await store.GetItemAsync(BuildItemKey(logicalKey), cancellationToken).ConfigureAwait(false); + if (!FoundryStateStoreJson.TryGetField(item, SessionField, out BinaryData? sessionData)) + { + return null; + } + + ReadOnlyMemory bytes = sessionData.ToMemory(); + // Parse and clone so the document buffer can be released. + using JsonDocument document = JsonDocument.Parse(bytes); + JsonElement element = document.RootElement.Clone(); + return await agent.DeserializeSessionAsync(element, cancellationToken: cancellationToken).ConfigureAwait(false); + } + + /// + /// Resolves the bound state store, creating it on the platform the first time. See + /// for the caching and failure behaviour. + /// + private ValueTask GetStoreAsync(CancellationToken cancellationToken) + => this._binding.GetAsync(cancellationToken); + + /// + /// Builds the readable partition key. This is the same a-/u-/c- scheme + /// and use, so + /// the three stores partition sessions identically: per hosted agent, then per end user, then + /// per conversation. agent.Id is deliberately not used because it is regenerated on every + /// startup for in-memory-defined agents, which would break session continuity. Each segment is + /// omitted when its value is absent. + /// + internal static string BuildLogicalKey(AIAgent agent, string conversationId, string? userId) + { + StringBuilder builder = new(); + if (!string.IsNullOrEmpty(agent.Name)) + { + builder.Append("a-").Append(agent.Name).Append(':'); + } + + if (!string.IsNullOrWhiteSpace(userId)) + { + builder.Append("u-").Append(userId).Append(':'); + } + + return builder.Append("c-").Append(conversationId).ToString(); + } + + /// + /// Reduces a logical key to a fixed-length item key. The platform limits an item key to 128 + /// characters, which an agent name plus a user id plus a conversation id can exceed, so the + /// logical key is hashed rather than truncated: truncation would let two different conversations + /// share a key and therefore overwrite each other's session. + /// + internal static string BuildItemKey(string logicalKey) + { + byte[] hash = SHA256.HashData(Encoding.UTF8.GetBytes(logicalKey)); + return $"s-{Convert.ToBase64String(hash).TrimEnd('=').Replace('+', '-').Replace('/', '_')}"; + } + + private static BinaryData ToBinaryData(JsonElement element) => FoundryStateStoreJson.ToBinaryData(element); + + private static BinaryData ToJsonString(string value) => FoundryStateStoreJson.ToJsonString(value); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs new file mode 100644 index 0000000000..55809bd32f --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs @@ -0,0 +1,523 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Security.Cryptography; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Core.Storage; +using Azure.Core; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Extensions.Logging; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Foundry.Hosting; + +/// +/// Provides a that persists workflow checkpoints to the Foundry +/// platform's durable state-store API (). +/// +/// +/// +/// Item keys are hashes of the session identifier and the checkpoint identifier, because the +/// platform limits an item key to 128 characters and neither identifier is bounded. Hashing rather +/// than truncating means two different checkpoints can never end up sharing a key and overwriting +/// each other. +/// +/// +/// Retention. A workflow session writes one checkpoint per superstep but only ever resumes from +/// the most recent one. Retrieving a checkpoint happens when a workflow is resuming from it, and at +/// that point every other checkpoint of that session is deleted. A conversation therefore holds one +/// turn's worth of checkpoints rather than growing for as long as it lasts to avoid storing +/// limitations. Note that this makes the point where old +/// checkpoints are collected, so retrieving one is not a read-only operation on this store. +/// +/// +/// Concurrency. Adding a checkpoint writes the checkpoint item and then updates the session's index +/// item using the platform's optimistic concurrency check, retrying a bounded number of times when +/// another writer got there first. Two instances committing checkpoints for the same session at the +/// same time therefore do not lose entries. +/// +/// +/// This store partitions only by workflow session identifier, which is the only partition the +/// contract carries. It does not partition by end +/// user. Callers that serve several end users from one workflow session must keep user separation +/// in the session identifier itself. +/// +/// +[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] +public sealed class FoundryJsonCheckpointStore : JsonCheckpointStore +{ + /// + /// The default state-store name used to hold every workflow checkpoint persisted by this store. + /// + public const string DefaultStoreName = "agent-framework/checkpoints"; + + /// + /// How many times a losing index update is retried before giving up. Each attempt re-reads the + /// index, so a retry only happens when another writer committed a checkpoint in between. + /// + private const int MaxIndexUpdateAttempts = 8; + + /// The item-body field holding the serialized checkpoint JSON. + private const string CheckpointField = "checkpoint"; + + /// The item-body field holding the owning session identifier, for traceability. + private const string SessionField = "session"; + + /// The item-body field of an index item holding the ordered checkpoint entries. + private const string EntriesField = "entries"; + + private const string EntryIdProperty = "id"; + private const string EntryParentProperty = "parent"; + private const string EntryHasParentProperty = "hasParent"; + + private readonly FoundryStateStoreBinding _binding; + private readonly ILogger? _logger; + + /// + /// Initializes a new instance of the class. + /// + /// The credential used to authenticate to the Foundry storage API. + /// + /// The Foundry project endpoint. When , it is read from the + /// FOUNDRY_PROJECT_ENDPOINT environment variable, which the platform sets in a hosted + /// container. + /// + /// The state-store name to hold the checkpoints. Defaults to . + /// + /// How long a checkpoint survives without being written, in seconds. Defaults to the platform + /// default of 30 days; -1 means never expire. A write renews the window, a read does + /// not. The value only takes effect when this store is created for the first time, because the + /// platform fixes it at creation. + /// + /// + /// Creates the logger this store reports through. Optional, but without one a failure to clean + /// up old checkpoints leaves no trace, since it is deliberately not allowed to fail the call it + /// happens in. + /// + public FoundryJsonCheckpointStore( + TokenCredential credential, + Uri? endpoint = null, + string storeName = DefaultStoreName, + int itemTtlSeconds = FoundryStateStore.DefaultItemTtlSeconds, + ILoggerFactory? loggerFactory = null) + { + _ = Throw.IfNull(credential); + _ = Throw.IfNullOrWhitespace(storeName); + + this.StoreName = storeName; + this._logger = loggerFactory?.CreateLogger(); + this._binding = new(cancellationToken => FoundryStateStore.GetOrCreateAsync( + storeName, + credential, + endpoint, + description: "Agent Framework hosted workflow checkpoints.", + itemTtlSeconds: itemTtlSeconds, + cancellationToken: cancellationToken)); + } + + /// + /// Initializes a new instance of the class over a + /// caller-supplied state store. Used by tests to substitute the platform client. + /// + /// Resolves the bound state store on first use. + /// The state-store name, for diagnostics. + /// Creates the logger this store reports through. + internal FoundryJsonCheckpointStore( + Func> storeFactory, + string storeName = DefaultStoreName, + ILoggerFactory? loggerFactory = null) + { + _ = Throw.IfNull(storeFactory); + + this._binding = new(storeFactory); + this.StoreName = storeName; + this._logger = loggerFactory?.CreateLogger(); + } + + /// Gets the state-store name that holds the checkpoints. + public string StoreName { get; } + + /// + public override async ValueTask CreateCheckpointAsync(string sessionId, JsonElement value, CheckpointInfo? parent = null) + { + _ = Throw.IfNullOrWhitespace(sessionId); + + BinaryData checkpointData = FoundryStateStoreJson.ToBinaryData(value); + + FoundryStateStore store = await this._binding.GetAsync(CancellationToken.None).ConfigureAwait(false); + string sessionIndexKey = BuildIndexKey(sessionId); + + // The identifier is chosen once, so a retried index update does not leave behind an orphan + // checkpoint item under a discarded identifier. + CheckpointInfo checkpointInfo = new(sessionId, Guid.NewGuid().ToString("N")); + + // Store the checkpoint itself, once. Only the index update below is ever retried. + await store.SetItemAsync( + BuildCheckpointKey(sessionId, checkpointInfo.CheckpointId), + new Dictionary + { + [CheckpointField] = checkpointData, + [SessionField] = FoundryStateStoreJson.ToJsonString(sessionId), + }, + cancellationToken: CancellationToken.None).ConfigureAwait(false); + + // Announce the stored checkpoint by appending its identifier to the session's index, giving + // way and reading again whenever another instance updated that same index first. + for (int attempt = 0; attempt < MaxIndexUpdateAttempts; attempt++) + { + StateStoreItem? indexItem = await store.GetItemAsync(sessionIndexKey, CancellationToken.None).ConfigureAwait(false); + List entries = ReadEntries(indexItem); + + if (Contains(entries, checkpointInfo.CheckpointId)) + { + // Two random identifiers colliding is not realistic, but the file-system and + // in-memory stores both guard against it and this store keeps the same guarantee. + throw new InvalidOperationException( + $"The generated checkpoint identifier '{checkpointInfo.CheckpointId}' is already in use for session '{sessionId}'."); + } + + entries.Add(new IndexEntry(checkpointInfo.CheckpointId, parent?.CheckpointId, HasParentMetadata: true)); + + try + { + await WriteEntriesAsync(store, sessionIndexKey, sessionId, entries, indexItem?.Etag).ConfigureAwait(false); + return checkpointInfo; + } + catch (FoundryStorageException ex) when (IsLostRace(ex)) + { + // Another writer added a checkpoint to the same session between the read and the + // write. The checkpoint item is already stored under its own key, so the next + // attempt simply re-reads the index and appends to the newer list. + if (this._logger?.IsEnabled(LogLevel.Debug) is true) + { + this._logger.LogDebug( + ex, + "Attempt {Attempt} of {MaxAttempts} to index checkpoint '{CheckpointId}' for session '{SessionId}' lost to another writer. Retrying.", + attempt + 1, + MaxIndexUpdateAttempts, + checkpointInfo.CheckpointId, + sessionId); + } + + continue; + } + } + + throw new InvalidOperationException( + $"Could not add a checkpoint for session '{sessionId}' to the Foundry state store after {MaxIndexUpdateAttempts} attempts because other writers kept updating the same session index."); + } + + /// + /// Returns a stored checkpoint and, in the same call, deletes every other checkpoint of that + /// session. + /// + /// + /// + /// The method also deletes every other checkpoint of that session except the one that has just been retrieved, which is + /// the one the workflow is resuming from. The deletion is not incidental, it is how this store keeps a session's checkpoints from + /// piling up. + /// + /// + /// A workflow writes one checkpoint per superstep and only ever resumes from the most + /// recent one, so a conversation that ran for a long time would otherwise leave behind every + /// checkpoint it ever wrote, and the index listing them would grow past the size the platform + /// accepts for a single item. + /// + /// + /// The workflow session that owns the checkpoint. + /// Identifies the checkpoint to return, and the one checkpoint left in place. + /// The stored checkpoint. + /// No such checkpoint is stored for that session. + public override async ValueTask RetrieveCheckpointAsync(string sessionId, CheckpointInfo key) + { + _ = Throw.IfNullOrWhitespace(sessionId); + _ = Throw.IfNull(key); + + FoundryStateStore store = await this._binding.GetAsync(CancellationToken.None).ConfigureAwait(false); + StateStoreItem? item = await store.GetItemAsync(BuildCheckpointKey(sessionId, key.CheckpointId), CancellationToken.None).ConfigureAwait(false); + if (!FoundryStateStoreJson.TryGetField(item, CheckpointField, out BinaryData? checkpointData)) + { + throw new KeyNotFoundException( + $"Checkpoint '{key.CheckpointId}' was not found for session '{sessionId}' in the Foundry state store '{this.StoreName}'."); + } + + JsonElement checkpoint = ParseCheckpoint(checkpointData); + + // Keeps a session's checkpoints from piling up. + await this.PruneObsoleteCheckpointsAsync(store, sessionId, key.CheckpointId).ConfigureAwait(false); + + return checkpoint; + } + + /// + /// Reads the stored bytes into a standalone . + /// + /// + /// returns an element that owns its own + /// memory, so there is no document to dispose and no copy to take. The reader lives in this + /// method rather than in the caller because it is a ref struct, which cannot be held + /// across an await. + /// + private static JsonElement ParseCheckpoint(BinaryData checkpointData) + { + Utf8JsonReader reader = new(checkpointData.ToMemory().Span); + return JsonElement.ParseValue(ref reader); + } + + /// + /// Deletes every checkpoint of a session except the one that has just been retrieved, which is + /// the one the workflow is resuming from. + /// + /// + /// + /// A session accumulates one checkpoint per superstep, and only the most recent one is ever + /// resumed from. Without this, a long conversation leaves behind every checkpoint it ever wrote + /// and the session's index item grows until it can no longer be saved. + /// + /// + private async Task PruneObsoleteCheckpointsAsync(FoundryStateStore store, string sessionId, string resumedCheckpointId) + { + string sessionIndexKey = BuildIndexKey(sessionId); + + try + { + StateStoreItem? indexItem = await store.GetItemAsync(sessionIndexKey, CancellationToken.None).ConfigureAwait(false); + + List obsolete = []; + IndexEntry? resumed = null; + foreach (IndexEntry entry in ReadEntries(indexItem)) + { + if (entry.CheckpointId == resumedCheckpointId) + { + resumed = entry; + } + else + { + obsolete.Add(entry); + } + } + + if (resumed is null || obsolete.Count == 0) + { + return; + } + + // The index is shortened first. A checkpoint item that is still listed but already gone + // would be read as a missing checkpoint, whereas one that is listed nowhere is simply + // never asked for. + await WriteEntriesAsync(store, sessionIndexKey, sessionId, [resumed], indexItem?.Etag).ConfigureAwait(false); + + foreach (IndexEntry entry in obsolete) + { + try + { + await store.DeleteItemAsync(BuildCheckpointKey(sessionId, entry.CheckpointId), cancellationToken: CancellationToken.None).ConfigureAwait(false); + } + catch (FoundryStorageNotFoundException ex) + { + // Already deleted, by an earlier attempt or another instance. + if (this._logger?.IsEnabled(LogLevel.Debug) is true) + { + this._logger.LogDebug( + ex, + "Obsolete checkpoint '{CheckpointId}' of session '{SessionId}' was already gone.", + entry.CheckpointId, + sessionId); + } + } + } + } + catch (FoundryStorageException ex) when (IsLostRace(ex)) + { + // Another instance updated the same session index first. Its own resume prunes whatever + // this one left behind, so nothing is leaked and there is nothing to report. + if (this._logger?.IsEnabled(LogLevel.Debug) is true) + { + this._logger.LogDebug( + ex, + "Pruning obsolete checkpoints of session '{SessionId}' lost to another writer. The winning writer prunes them instead.", + sessionId); + } + } + catch (FoundryStorageException ex) + { + // Not a lost race: the store refused the call for a reason of its own, a credential or a + // network problem for instance. The checkpoint has already been retrieved by this point, + // so failing the resume over housekeeping would break a conversation that was about to + // carry on. It is reported instead, because left unreported this is how a session's + // checkpoints would silently pile up until the index no longer fits. + this._logger?.LogWarning( + ex, + "Could not prune obsolete checkpoints of session '{SessionId}' in the Foundry state store '{StoreName}'. The resume itself succeeded; the leftovers stay until the store's own expiry removes them.", + sessionId, + this.StoreName); + } + } + + /// + public override async ValueTask> RetrieveIndexAsync(string sessionId, CheckpointInfo? withParent = null) + { + _ = Throw.IfNullOrWhitespace(sessionId); + + FoundryStateStore store = await this._binding.GetAsync(CancellationToken.None).ConfigureAwait(false); + StateStoreItem? indexItem = await store.GetItemAsync(BuildIndexKey(sessionId), CancellationToken.None).ConfigureAwait(false); + + List result = []; + foreach (IndexEntry entry in ReadEntries(indexItem)) + { + // Same filter the file-system store applies: an entry written before parents were + // recorded is always included, because its parent is unknown rather than different. + if (withParent is null || !entry.HasParentMetadata || entry.ParentCheckpointId == withParent.CheckpointId) + { + result.Add(new CheckpointInfo(sessionId, entry.CheckpointId)); + } + } + + return result; + } + + /// Reports whether the index already lists the given checkpoint identifier. + private static bool Contains(List entries, string checkpointId) + { + foreach (IndexEntry entry in entries) + { + if (entry.CheckpointId == checkpointId) + { + return true; + } + } + + return false; + } + + /// + /// Reports whether a storage failure means "someone else wrote this item first", which is the + /// only failure this store retries. A 412 says the optimistic concurrency check failed; a 409 + /// says an item this store expected to be absent had already been created. + /// + private static bool IsLostRace(FoundryStorageException exception) + => exception is FoundryStoragePreconditionException or FoundryStorageConflictException; + + private static List ReadEntries(StateStoreItem? indexItem) + { + List entries = []; + + if (!FoundryStateStoreJson.TryGetField(indexItem, EntriesField, out BinaryData? entriesData)) + { + return entries; + } + + using JsonDocument document = JsonDocument.Parse(entriesData.ToMemory()); + if (document.RootElement.ValueKind != JsonValueKind.Array) + { + return entries; + } + + foreach (JsonElement element in document.RootElement.EnumerateArray()) + { + if (element.ValueKind != JsonValueKind.Object || + !element.TryGetProperty(EntryIdProperty, out JsonElement idElement) || + idElement.GetString() is not string checkpointId || + checkpointId.Length == 0) + { + continue; + } + + string? parentId = element.TryGetProperty(EntryParentProperty, out JsonElement parentElement) && parentElement.ValueKind == JsonValueKind.String + ? parentElement.GetString() + : null; + + bool hasParentMetadata = element.TryGetProperty(EntryHasParentProperty, out JsonElement hasParentElement) && + hasParentElement.ValueKind == JsonValueKind.True; + + entries.Add(new IndexEntry(checkpointId, parentId, hasParentMetadata)); + } + + return entries; + } + + private static async Task WriteEntriesAsync(FoundryStateStore store, string sessionIndexKey, string sessionId, List entries, string? ifMatch) + { + Dictionary value = new() + { + [EntriesField] = WriteEntries(entries), + [SessionField] = FoundryStateStoreJson.ToJsonString(sessionId), + }; + + if (ifMatch is null) + { + // The index did not exist a moment ago. CreateItemAsync fails with a conflict if another + // writer created it in the meantime, which the caller retries. + await store.CreateItemAsync(sessionIndexKey, value, cancellationToken: CancellationToken.None).ConfigureAwait(false); + return; + } + + await store.SetItemAsync(sessionIndexKey, value, ifMatch: ifMatch, cancellationToken: CancellationToken.None).ConfigureAwait(false); + } + + private static BinaryData WriteEntries(List entries) + { + System.Buffers.ArrayBufferWriter buffer = new(); + using (Utf8JsonWriter writer = new(buffer)) + { + writer.WriteStartArray(); + foreach (IndexEntry entry in entries) + { + writer.WriteStartObject(); + writer.WriteString(EntryIdProperty, entry.CheckpointId); + if (entry.ParentCheckpointId is not null) + { + writer.WriteString(EntryParentProperty, entry.ParentCheckpointId); + } + + writer.WriteBoolean(EntryHasParentProperty, entry.HasParentMetadata); + writer.WriteEndObject(); + } + + writer.WriteEndArray(); + } + + return BinaryData.FromBytes(buffer.WrittenMemory); + } + + /// + /// Builds the key of the item holding a session's ordered checkpoint index. The wi- + /// prefix separates index items from checkpoint items, which share one state store. + /// + internal static string BuildIndexKey(string sessionId) => $"wi-{HashKeyParts(sessionId)}"; + + /// + /// Builds the key of the item holding one checkpoint's serialized state. The wc- prefix + /// separates checkpoint items from index items, which share one state store. + /// + internal static string BuildCheckpointKey(string sessionId, string checkpointId) => $"wc-{HashKeyParts(sessionId, checkpointId)}"; + + /// + /// Folds the identifiers an item key is made of into a fixed-length string. + /// + /// + /// The platform caps an item key at 128 characters, and neither a workflow session identifier + /// nor a checkpoint identifier has a bounded length, so they are hashed. Hashing rather than + /// truncating matters: two sessions whose identifiers share a long prefix would otherwise be cut + /// down to the same key and overwrite each other's checkpoints. The parts are joined with a NUL + /// character, which cannot appear inside either identifier, so no two different combinations can + /// produce the same input. The result is written in the URL-safe base64 alphabet because the + /// key becomes a segment of the request path the platform client builds. + /// + /// The identifiers that make this key unique, in a fixed order. + /// The hashed key body, without the prefix that says what kind of item it is. + private static string HashKeyParts(params string[] parts) + { + byte[] hash = SHA256.HashData(Encoding.UTF8.GetBytes(string.Join("\u0000", parts))); + return Convert.ToBase64String(hash).TrimEnd('=').Replace('+', '-').Replace('/', '_'); + } + + private sealed record IndexEntry(string CheckpointId, string? ParentCheckpointId, bool HasParentMetadata); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs new file mode 100644 index 0000000000..2e538b9a7f --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Buffers; +using System.Diagnostics.CodeAnalysis; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Core.Storage; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Foundry.Hosting; + +/// +/// Resolves a once and hands the same instance to every later +/// caller. Resolving costs a network round trip, plus one more the very first time to create the +/// store on the platform, so it deliberately does not happen per request. +/// +/// +/// A failed attempt is not kept: the next call starts a fresh one, so a transient network or +/// permission failure at startup does not leave the store unusable for the life of the process. +/// +internal sealed class FoundryStateStoreBinding +{ + private readonly Func> _factory; + private readonly object _gate = new(); + private Task? _pending; + + public FoundryStateStoreBinding(Func> factory) + { + this._factory = Throw.IfNull(factory); + } + + public async ValueTask GetAsync(CancellationToken cancellationToken) + { + Task binding; + lock (this._gate) + { + // The shared work is started without the caller's cancellation token so one cancelled + // request cannot cancel the binding for every other request. + binding = this._pending ??= this._factory(CancellationToken.None); + } + + try + { + // WaitAsync applies the caller's token to this caller's wait only. + return await binding.WaitAsync(cancellationToken).ConfigureAwait(false); + } + catch (Exception ex) when (ex is not OperationCanceledException) + { + lock (this._gate) + { + if (ReferenceEquals(this._pending, binding)) + { + this._pending = null; + } + } + + throw; + } + } +} + +/// +/// Small JSON helpers shared by the Foundry-backed stores. They are written with +/// rather than a serializer so the callers stay trimming and +/// ahead-of-time compilation safe. +/// +internal static class FoundryStateStoreJson +{ + /// Writes a out as UTF-8 bytes. + public static BinaryData ToBinaryData(JsonElement element) + { + ArrayBufferWriter buffer = new(); + using (Utf8JsonWriter writer = new(buffer)) + { + element.WriteTo(writer); + } + + return BinaryData.FromBytes(buffer.WrittenMemory); + } + + /// Encodes a plain string as a JSON string value. + public static BinaryData ToJsonString(string value) + { + ArrayBufferWriter buffer = new(); + using (Utf8JsonWriter writer = new(buffer)) + { + writer.WriteStringValue(value); + } + + return BinaryData.FromBytes(buffer.WrittenMemory); + } + + /// + /// Reads one field out of a state-store item body, treating a missing item, a missing field and + /// an empty value all as "nothing stored". + /// + public static bool TryGetField(StateStoreItem? item, string field, [NotNullWhen(true)] out BinaryData? data) + { + data = null; + + if (item is null) + { + return false; + } + + if (!item.Value.TryGetValue(field, out BinaryData? value) || value is null) + { + return false; + } + + if (value.ToMemory().IsEmpty) + { + return false; + } + + data = value; + return true; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs index 38d3c67380..4b9429b617 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs @@ -6,11 +6,11 @@ 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.Extensions.Options; using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Foundry.Hosting; @@ -46,7 +46,7 @@ public HostedStoredOutputHealthCheck( IOptions? hostingOptions = null, ILogger? logger = null) { - ArgumentNullException.ThrowIfNull(serviceProvider); + _ = Throw.IfNull(serviceProvider); this._serviceProvider = serviceProvider; this._hostingOptions = hostingOptions?.Value ?? new FoundryResponsesOptions(); @@ -55,7 +55,7 @@ public HostedStoredOutputHealthCheck( public async Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default) { - ArgumentNullException.ThrowIfNull(context); + _ = Throw.IfNull(context); if (this._hostingOptions.AllowStoredOutputEnabled) { @@ -149,15 +149,5 @@ private async Task StoresItsOwnResponsesAsync(AIAgent agent, CancellationT /// /// Every agent this container can serve: the ones registered under a name, plus the default. /// - private List ResolveAgents() - { - var agents = new List(this._serviceProvider.GetKeyedServices(KeyedService.AnyKey)); - - if (this._serviceProvider.GetService() is { } defaultAgent && !agents.Contains(defaultAgent)) - { - agents.Add(defaultAgent); - } - - return agents; - } + private List ResolveAgents() => FoundryHostingExtensions.ResolveRegisteredAgents(this._serviceProvider); } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs new file mode 100644 index 0000000000..c9c1c49dd5 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Globalization; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Extensions.Diagnostics.HealthChecks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Foundry.Hosting; + +/// +/// Reports, on the GET /readiness probe, a registered workflow agent that was built with a +/// checkpoint manager of its own, so a container whose workflow state would be written somewhere +/// hosting cannot manage is caught before it takes any traffic. +/// +/// +/// +/// A hosted workflow has its checkpoints redirected to the Foundry durable state store, so that the +/// state a conversation builds up survives the container being restarted or replaced and is readable +/// by every instance of the agent. An agent that already names its own checkpoint manager is left +/// alone, because overriding an explicit choice silently would be worse. The result is a container +/// whose workflow state goes somewhere hosting does not manage, which is reported here rather than +/// discovered later. +/// +/// +/// Only an agent that runs a workflow is considered. Everything else, a +/// or an agent written by the container author, has no checkpoints and is passed over. +/// +/// +[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] +internal sealed class HostedWorkflowCheckpointingHealthCheck : IHealthCheck +{ + private readonly IServiceProvider _serviceProvider; + + public HostedWorkflowCheckpointingHealthCheck(IServiceProvider serviceProvider) + { + _ = Throw.IfNull(serviceProvider); + + this._serviceProvider = serviceProvider; + } + + /// + /// Whether the process is running inside a Foundry container. Settable so a test does not depend + /// on the process-wide, statically-cached value. + /// + internal bool IsHosted { get; set; } = FoundryEnvironment.IsHosted; + + public Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(context); + + if (!this.IsHosted) + { + // 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.")); + } + + List agentsWithOwnCheckpointing = []; + var checkedAgents = 0; + + foreach (var agent in FoundryHostingExtensions.ResolveRegisteredAgents(this._serviceProvider)) + { + if (agent.GetService() is not { } metadata) + { + continue; + } + + checkedAgents++; + if (metadata.UsesOwnCheckpointStorage) + { + agentsWithOwnCheckpointing.Add(agent.Name ?? agent.Id); + } + } + + if (agentsWithOwnCheckpointing.Count > 0) + { + return Task.FromResult(new HealthCheckResult( + status: context.Registration.FailureStatus, + description: string.Create( + CultureInfo.InvariantCulture, + $"Workflow checkpointing: {agentsWithOwnCheckpointing.Count} registered workflow agent(s) were built with a checkpoint manager of their own. A hosted workflow has its checkpoints written to the Foundry state store so they survive the container being replaced and can be read by every instance; an agent that names its own manager keeps that state somewhere this container does not manage. Build the agent without passing an execution environment configured with WithCheckpointing, and let hosting supply the store."), + data: new Dictionary(StringComparer.Ordinal) { ["agentsWithOwnCheckpointing"] = agentsWithOwnCheckpointing })); + } + + return Task.FromResult(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 28152362cb..2f0d2a4ce1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs @@ -2,12 +2,15 @@ using System; using System.ClientModel.Primitives; +using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Globalization; using System.Linq; using System.Runtime.CompilerServices; using Azure.AI.AgentServer.Responses; using Azure.Core; +using Azure.Identity; +using Microsoft.Agents.AI.Workflows; using Microsoft.AspNetCore.Builder; using Microsoft.AspNetCore.Routing; using Microsoft.AspNetCore.Server.Kestrel.Core; @@ -19,6 +22,7 @@ using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Foundry.Hosting; @@ -58,12 +62,12 @@ public static class FoundryHostingExtensions /// The service collection for chaining. public static IServiceCollection AddFoundryResponses(this IServiceCollection services, Action? configure = null) { - ArgumentNullException.ThrowIfNull(services); - services.AddResponsesServer(); + _ = Throw.IfNull(services); + AddResponsesServerOnce(services); services.AddHealthChecks(); ConfigureFoundryListenPort(services); ConfigureFoundryResponsesOptions(services, configure); - services.TryAddSingleton(_ => FileSystemAgentSessionStore.CreateDefault()); + services.TryAddSingleton(_ => CreateDefaultAgentSessionStore()); services.TryAddSingleton(); return services; } @@ -90,7 +94,7 @@ 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, a file-system session store is used, rooted at /.checkpoints when running in a Foundry hosted environment and {cwd}/.checkpoints locally. + /// The agent session store to use for managing agent sessions server-side. If null, the default store is chosen for the environment: the Foundry durable state store when the platform has supplied a project endpoint (FOUNDRY_PROJECT_ENDPOINT), otherwise a file-system store rooted at {$HOME}/.checkpoints when hosted and {cwd}/.checkpoints locally. /// /// Optional callback to configure , for example to allow the /// agent's own service to store the responses it produces. @@ -102,14 +106,14 @@ public static IServiceCollection AddFoundryResponses( AgentSessionStore? agentSessionStore = null, Action? configure = null) { - ArgumentNullException.ThrowIfNull(services); - ArgumentNullException.ThrowIfNull(agent); + _ = Throw.IfNull(services); + _ = Throw.IfNull(agent); - services.AddResponsesServer(); + AddResponsesServerOnce(services); services.AddHealthChecks(); ConfigureFoundryListenPort(services); ConfigureFoundryResponsesOptions(services, configure); - agentSessionStore ??= FileSystemAgentSessionStore.CreateDefault(); + agentSessionStore ??= CreateDefaultAgentSessionStore(); if (!string.IsNullOrWhiteSpace(agent.Name)) { @@ -127,12 +131,13 @@ public static IServiceCollection AddFoundryResponses( } /// - /// Applies the caller's and registers the readiness check that - /// reports an agent configured to have its own service store the responses it produces. + /// Applies the caller's and registers the readiness checks + /// that report a misconfigured agent: one having its own service store the responses it produces, + /// and a workflow agent writing its checkpoints somewhere hosting does not manage. /// /// - /// The check is registered on the same /readiness pipeline that - /// maps, so a container that would record the conversation twice never takes traffic. + /// The checks are registered on the same /readiness pipeline that + /// maps, so such a container never takes traffic. /// AddCheck does not dedupe by name, so a repeated registration is guarded here. /// private static void ConfigureFoundryResponsesOptions(IServiceCollection services, Action? configure) @@ -142,24 +147,32 @@ private static void ConfigureFoundryResponsesOptions(IServiceCollection services services.Configure(configure); } - const string HealthCheckName = "foundry-stored-output"; + AddReadinessCheckOnce(services, "foundry-stored-output", sp => ActivatorUtilities.CreateInstance(sp)); + AddReadinessCheckOnce(services, "foundry-workflow-checkpointing", sp => ActivatorUtilities.CreateInstance(sp)); + } + + /// + /// Registers a readiness check under a name, skipping the registration when that name is already + /// taken, because AddCheck does not dedupe and both AddFoundryResponses overloads + /// are documented as safe to call more than once. + /// + private static void AddReadinessCheckOnce(IServiceCollection services, string name, Func factory) => services.Configure(opts => { foreach (var existing in opts.Registrations) { - if (string.Equals(existing.Name, HealthCheckName, StringComparison.Ordinal)) + if (string.Equals(existing.Name, name, StringComparison.Ordinal)) { return; } } opts.Registrations.Add(new HealthCheckRegistration( - name: HealthCheckName, - factory: sp => ActivatorUtilities.CreateInstance(sp), + name: name, + factory: factory, failureStatus: HealthStatus.Unhealthy, tags: ["foundry", "responses", "readiness"])); }); - } /// /// Registers the Foundry Toolbox service, which eagerly connects to the Foundry Toolboxes @@ -206,8 +219,8 @@ public static IServiceCollection AddFoundryToolboxes( Action? configureOptions, params string[] toolboxNames) { - ArgumentNullException.ThrowIfNull(services); - ArgumentNullException.ThrowIfNull(credential); + _ = Throw.IfNull(services); + _ = Throw.IfNull(credential); if (services.Any(d => d.ServiceType == typeof(FoundryToolboxService))) { @@ -297,7 +310,7 @@ public static IServiceCollection AddFoundryToolboxes( /// The endpoint route builder for chaining. public static IEndpointRouteBuilder MapFoundryResponses(this IEndpointRouteBuilder endpoints, string prefix = "") { - ArgumentNullException.ThrowIfNull(endpoints); + _ = Throw.IfNull(endpoints); endpoints.MapResponsesServer(prefix); MapReadinessIfMissing(endpoints); return endpoints; @@ -320,12 +333,70 @@ public static IEndpointRouteBuilder MapFoundryResponses(this IEndpointRouteBuild /// internal const int DefaultListenPort = 8088; + /// + /// Registers the Responses Server SDK exactly once per service collection. + /// + /// + /// AddResponsesServer registers a resilient task under a fixed name and throws when that + /// name is already taken, so calling it a second time on the same service collection fails. + /// Both AddFoundryResponses overloads are documented as safe to call more than once, and + /// a host that registers several agents naturally does, so the second and later calls are + /// skipped here. + /// + private static void AddResponsesServerOnce(IServiceCollection services) + { + if (services.Any(static d => d.ServiceType == typeof(FoundryResponsesServerMarker))) + { + return; + } + + services.AddSingleton(); + services.AddResponsesServer(); + } + + /// + /// Creates the used when the caller did not supply one. + /// + /// + /// Inside a Foundry container, sessions are held by the platform's durable state store, because + /// that state survives the container being restarted or replaced and is readable by every + /// instance of the agent. Anywhere else there is no such service to call, so the container falls + /// back to writing session files under its own session directory, which is what a local run does. + /// + private static AgentSessionStore CreateDefaultAgentSessionStore() => + FoundryEnvironment.IsHosted + ? new FoundryAgentSessionStore(new DefaultAzureCredential()) + : FileSystemAgentSessionStore.CreateDefault(); + + /// + /// Every agent a container can serve: the ones registered under a name, plus the default. + /// + /// The provider the agents were registered with. + /// The registered agents, without duplicates. + internal static List ResolveRegisteredAgents(IServiceProvider serviceProvider) + { + var agents = new List(serviceProvider.GetKeyedServices(KeyedService.AnyKey)); + + if (serviceProvider.GetService() is { } defaultAgent && !agents.Contains(defaultAgent)) + { + agents.Add(defaultAgent); + } + + return agents; + } + /// /// Marker registered once per so the Foundry listen-port /// configuration is applied at most once, even across multiple AddFoundryResponses calls. /// private sealed class FoundryListenPortMarker; + /// + /// Marker registered once per so the Responses Server SDK is + /// registered at most once, even across multiple AddFoundryResponses calls. + /// + private sealed class FoundryResponsesServerMarker; + /// /// Binds Kestrel to the port the Foundry hosted runtime probes and routes to, so a plain /// WebApplication.CreateBuilder host (Tier 3) works with no Dockerfile. Mirrors @@ -451,6 +522,68 @@ internal static AIAgent ApplyOpenTelemetry(AIAgent agent) .Build(); } + /// + /// Points a workflow-hosting agent at the Foundry durable state store for its checkpoints, + /// when running inside a Foundry container. + /// + /// + /// + /// This runs when the agent is resolved for a request rather than when it is registered, + /// because a host can register agents as factories that are only built later, and because a + /// registered agent is a finished object whose checkpoint storage is fixed at construction. + /// + /// + /// Without this, a hosted workflow keeps every checkpoint of a session inside the saved session + /// record, and the platform limits a single record to 1 MB, so a long workflow eventually stops + /// being able to save. With it, each checkpoint becomes its own record and the session keeps + /// only the pointer to the last one. + /// + /// + /// The method is a no-op when the agent does not host a workflow, when the workflow was built + /// with an explicit checkpoint manager, and when the process is not running on the platform. + /// The redirected agent is cached against the agent it came from, so the substitution happens + /// once rather than on every request. + /// + /// + /// The resolved agent. + /// Creates the logger the checkpoint store reports through. + /// The agent to serve the request with. + internal static AIAgent ApplyWorkflowCheckpointing(AIAgent agent, ILoggerFactory? loggerFactory = null) + { + if (!FoundryEnvironment.IsHosted) + { + return agent; + } + + return s_workflowCheckpointingAgents.GetValue( + agent, + source => source.WithCheckpointing(GetFoundryWorkflowCheckpointManager(loggerFactory))); + } + + /// + /// The single checkpoint manager shared by every hosted workflow in this process. It is created + /// on first use so that no credential is built and no platform call is made when the process is + /// not running on the platform, which also means the first caller supplies its logger. + /// + private static CheckpointManager GetFoundryWorkflowCheckpointManager(ILoggerFactory? loggerFactory) + { + lock (s_checkpointManagerGate) + { + return s_foundryWorkflowCheckpointManager ??= CheckpointManager.CreateJson( + new FoundryJsonCheckpointStore(new DefaultAzureCredential(), loggerFactory: loggerFactory)); + } + } + + private static readonly object s_checkpointManagerGate = new(); + private static CheckpointManager? s_foundryWorkflowCheckpointManager; + + /// + /// Caches the redirected copy of each agent. Rebuilding it per request would restart the + /// agent's protocol validation and throw away the session identifiers it tracks, so the copy + /// has to live as long as the agent it was made from. + /// + private static readonly ConditionalWeakTable s_workflowCheckpointingAgents = new(); + /// /// Registers the hosted-agent User-Agent supplement policy /// () on the agent's underlying chat client via the diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowAgentMetadata.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowAgentMetadata.cs new file mode 100644 index 0000000000..f6ea2135a4 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowAgentMetadata.cs @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.Workflows; + +/// +/// Describes an that runs a , for hosts that receive a +/// finished agent and have to treat a workflow differently from other agents. +/// +/// +/// +/// Retrieve it with agent.GetService<WorkflowAgentMetadata>(). Getting an instance back +/// is what identifies the agent as running a workflow; means it does not. +/// Going through means the answer is still +/// found when the agent has been wrapped, by middleware for example, which a test on the type of the +/// agent would miss. +/// +/// +/// This is separate from rather than a specialization of it, because +/// that type is sealed. It also carries nothing that belongs there: the provider name it holds names +/// the inference service behind an agent, and a workflow has none of its own. +/// +/// +public sealed class WorkflowAgentMetadata +{ + /// + /// Initializes a new instance of the class. + /// + /// + /// Whether the agent was built with an execution environment that already names a checkpoint + /// manager. + /// + public WorkflowAgentMetadata(bool usesOwnCheckpointStorage) + { + this.UsesOwnCheckpointStorage = usesOwnCheckpointStorage; + } + + /// + /// Gets a value indicating whether the agent already writes its checkpoints to a + /// named when the agent was built. + /// + /// + /// + /// When this is , the agent keeps its checkpoints in memory and they are + /// carried inside the serialized agent session. When it is , the caller + /// passed an execution environment built with + /// , and + /// leaves + /// such an agent alone rather than overriding that choice. + /// + /// + public bool UsesOwnCheckpointStorage { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index 01f9637967..d498d13b65 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs @@ -55,6 +55,48 @@ public WorkflowHostAgent(Workflow workflow, string? id = null, string? name = nu public override string? Name { get; } public override string? Description { get; } + /// + /// Reports whether this agent was built with an execution environment that already names a + /// checkpoint manager, meaning the caller chose where its checkpoints are written. + /// + internal bool UsesOwnCheckpointStorage => this._executionEnvironment.IsCheckpointingEnabled; + + /// + public override object? GetService(Type serviceType, object? serviceKey = null) + { + Throw.IfNull(serviceType); + + return base.GetService(serviceType, serviceKey) + ?? (serviceKey is null && serviceType == typeof(WorkflowAgentMetadata) + ? this._metadata ??= new WorkflowAgentMetadata(this.UsesOwnCheckpointStorage) + : null); + } + + private WorkflowAgentMetadata? _metadata; + + /// + /// Builds a copy of this agent that writes its checkpoints to . + /// Returns this same instance when the execution environment already names a checkpoint manager, + /// because that means the caller made an explicit choice. + /// + internal AIAgent WithCheckpointing(CheckpointManager checkpointManager) + { + if (this._executionEnvironment.IsCheckpointingEnabled || + this._executionEnvironment is not InProcessExecutionEnvironment inProcEnvironment) + { + return this; + } + + return new WorkflowHostAgent( + this._workflow, + this._id, + this.Name, + this.Description, + inProcEnvironment.WithCheckpointing(checkpointManager), + this._includeExceptionDetails, + this._includeWorkflowOutputsInResponse); + } + private string GenerateNewId() { string result; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostingExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostingExtensions.cs index 210537588b..0a0237cc66 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostingExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostingExtensions.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Workflows; @@ -37,6 +38,44 @@ public static AIAgent AsAIAgent( return new WorkflowHostAgent(workflow, id, name, description, executionEnvironment, includeExceptionDetails, includeWorkflowOutputsInResponse); } + /// + /// Returns a copy of a workflow-hosting agent that writes its checkpoints to the supplied + /// , so a host can redirect checkpoint storage on an agent that + /// has already been built. + /// + /// + /// + /// This exists for hosts that receive a finished and need to decide where + /// its checkpoints live, which is something only the host knows. The agent itself cannot be + /// changed in place, because the execution environment is fixed when the agent is constructed, + /// so a copy is returned instead. Everything else about the agent is preserved, including its + /// identifier, name, description and the workflow it runs. + /// + /// + /// The call leaves the agent untouched and returns it as-is in three cases: when it does not + /// host a workflow, when it was built with an execution environment that already names a + /// checkpoint manager, and when the workflow-hosting agent sits behind a wrapper such as + /// middleware. The last case is a limitation rather than a choice: only the innermost agent can + /// be copied, and returning it alone would silently throw the wrapper away. + /// + /// + /// The returned copy is a distinct agent, so a host that calls this on every request should + /// keep the result rather than rebuilding it each time. + /// + /// + /// The agent whose checkpoint storage should be redirected. + /// The checkpoint manager the copy should write to. + /// The redirected copy, or itself when nothing was changed. + public static AIAgent WithCheckpointing(this AIAgent agent, CheckpointManager checkpointManager) + { + Throw.IfNull(agent); + Throw.IfNull(checkpointManager); + + return agent is WorkflowHostAgent workflowAgent + ? workflowAgent.WithCheckpointing(checkpointManager) + : agent; + } + internal static FunctionCallContent ToFunctionCall(this ExternalRequest request) { Dictionary parameters = new() diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs new file mode 100644 index 0000000000..23ded1512c --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs @@ -0,0 +1,312 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Core.Storage; +using Microsoft.Agents.AI.Foundry.Hosting; + +namespace Microsoft.Agents.AI.Foundry.UnitTests.Hosting; + +public sealed class FoundryAgentSessionStoreTests +{ + [Fact] + public async Task SaveSessionAsync_ThenGetSessionAsync_RoundTripsAsync() + { + // Arrange + var backing = new FakeStateStore(); + var store = NewStore(backing); + var agent = new TestAgent("{\"foo\":7}", name: "Concierge"); + + // Act + await store.SaveSessionAsync(agent, "round-trip", new TestSession(), userId: "alice"); + var session = await store.GetSessionAsync(agent, "round-trip", userId: "alice"); + + // Assert + Assert.NotNull(session); + Assert.Equal(1, agent.SerializeCalls); + Assert.Equal(1, agent.DeserializeCalls); + Assert.Equal(7, agent.LastDeserialized!.Value.GetProperty("foo").GetInt32()); + } + + [Fact] + public async Task SaveSessionAsync_StoresReadableLogicalKeyAlongsideTheSessionAsync() + { + // Arrange + var backing = new FakeStateStore(); + var store = NewStore(backing); + var agent = new TestAgent(name: "Concierge"); + + // Act + await store.SaveSessionAsync(agent, "conv-1", new TestSession(), userId: "alice"); + + // Assert: the item body keeps the readable key so a stored item can be traced back. + var item = Assert.Single(backing.Items); + Assert.Equal("\"a-Concierge:u-alice:c-conv-1\"", item["key"].ToString()); + } + + [Fact] + public async Task GetSessionAsync_NothingStored_ReturnsNullAsync() + { + // Arrange + var store = NewStore(new FakeStateStore()); + var agent = new TestAgent(); + + // Act + var session = await store.GetSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.Null(session); + Assert.Equal(0, agent.CreateCalls); + Assert.Equal(0, agent.DeserializeCalls); + } + + [Fact] + public async Task GetOrCreateSessionAsync_NothingStored_ReturnsFreshSessionFromAgentAsync() + { + // Arrange + var store = NewStore(new FakeStateStore()); + var agent = new TestAgent(); + + // Act + var session = await store.GetOrCreateSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.NotNull(session); + Assert.Equal(1, agent.CreateCalls); + Assert.Equal(0, agent.DeserializeCalls); + } + + [Fact] + public async Task GetSessionAsync_DifferentUser_DoesNotReadAnotherUsersSessionAsync() + { + // Arrange: Alice saves under the conversation id Bob will forge. + var store = NewStore(new FakeStateStore()); + var agent = new TestAgent("{\"secret\":\"alice-only\"}", name: "Concierge"); + await store.SaveSessionAsync(agent, "shared-conv", new TestSession(), userId: "alice"); + + // Act + var bobSession = await store.GetSessionAsync(agent, "shared-conv", userId: "bob"); + + // Assert + Assert.Null(bobSession); + Assert.Equal(0, agent.DeserializeCalls); + } + + [Fact] + public async Task GetSessionAsync_DifferentAgent_DoesNotReadAnotherAgentsSessionAsync() + { + // Arrange: one container hosts several keyed agents that must not collide on a shared id. + var backing = new FakeStateStore(); + var store = NewStore(backing); + var concierge = new TestAgent("{\"owner\":\"concierge\"}", name: "Concierge"); + var researcher = new TestAgent(name: "Researcher"); + await store.SaveSessionAsync(concierge, "shared-conv", new TestSession(), userId: "alice"); + + // Act + var otherSession = await store.GetSessionAsync(researcher, "shared-conv", userId: "alice"); + + // Assert + Assert.Null(otherSession); + Assert.Equal(0, researcher.DeserializeCalls); + } + + [Fact] + public async Task GetStoreAsync_ResolvesTheStoreOnceAcrossManyCallsAsync() + { + // Arrange: binding the store costs a round trip, so it must not happen per request. + var backing = new FakeStateStore(); + var bindCount = 0; + var store = new FoundryAgentSessionStore(_ => + { + Interlocked.Increment(ref bindCount); + return Task.FromResult(backing); + }); + var agent = new TestAgent(name: "Concierge"); + + // Act + await store.SaveSessionAsync(agent, "conv-1", new TestSession(), userId: null); + await store.GetSessionAsync(agent, "conv-1", userId: null); + await store.GetSessionAsync(agent, "conv-2", userId: null); + + // Assert + Assert.Equal(1, bindCount); + } + + [Fact] + public async Task GetStoreAsync_FailedBinding_IsRetriedOnTheNextCallAsync() + { + // Arrange: a transient failure while binding must not disable the store for the process. + var backing = new FakeStateStore(); + var attempts = 0; + var store = new FoundryAgentSessionStore(_ => + { + attempts++; + return attempts == 1 + ? Task.FromException(new FoundryStorageApiException(503, "transient")) + : Task.FromResult(backing); + }); + var agent = new TestAgent(name: "Concierge"); + + // Act + await Assert.ThrowsAsync( + async () => await store.GetSessionAsync(agent, "conv-1", userId: null)); + var session = await store.GetSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.Null(session); + Assert.Equal(2, attempts); + } + + [Theory] + [InlineData("Concierge", "alice", "conv-1", "a-Concierge:u-alice:c-conv-1")] + [InlineData("Concierge", null, "conv-1", "a-Concierge:c-conv-1")] + [InlineData(null, "alice", "conv-1", "u-alice:c-conv-1")] + [InlineData(null, null, "conv-1", "c-conv-1")] + [InlineData("x", "x", "conv-1", "a-x:u-x:c-conv-1")] + public void BuildLogicalKey_UsesTheSamePrefixSchemeAsTheOtherStores(string? agentName, string? userId, string conversationId, string expected) + { + // Arrange + var agent = new TestAgent(name: agentName); + + // Act + var key = FoundryAgentSessionStore.BuildLogicalKey(agent, conversationId, userId); + + // Assert + Assert.Equal(expected, key); + } + + [Fact] + public void BuildItemKey_StaysWithinThePlatformKeyLimitForAnyInput() + { + // Arrange: an agent name plus a user id plus a conversation id can easily pass 128 chars. + var logicalKey = FoundryAgentSessionStore.BuildLogicalKey( + new TestAgent(name: new string('a', 200)), + new string('c', 200), + new string('u', 200)); + + // Act + var itemKey = FoundryAgentSessionStore.BuildItemKey(logicalKey); + + // Assert + Assert.InRange(itemKey.Length, 1, 128); + } + + [Fact] + public void BuildItemKey_IsStableAndDistinctPerLogicalKey() + { + // Arrange / Act + var first = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-alice:c-conv-1"); + var same = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-alice:c-conv-1"); + var other = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-bob:c-conv-1"); + + // Assert + Assert.Equal(first, same); + Assert.NotEqual(first, other); + } + + [Fact] + public void Constructor_NullOrWhitespaceStoreName_Throws() + { + // Arrange + var credential = new FakeCredential(); + + // Act / Assert + Assert.Throws(() => new FoundryAgentSessionStore(credential, storeName: null!)); + Assert.Throws(() => new FoundryAgentSessionStore(credential, storeName: " ")); + } + + private static FoundryAgentSessionStore NewStore(FakeStateStore backing) + => new(_ => Task.FromResult(backing)); + + private sealed class FakeCredential : Azure.Core.TokenCredential + { + public override Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, CancellationToken cancellationToken) + => new("token", DateTimeOffset.MaxValue); + + public override ValueTask GetTokenAsync(Azure.Core.TokenRequestContext requestContext, CancellationToken cancellationToken) + => new(new Azure.Core.AccessToken("token", DateTimeOffset.MaxValue)); + } + + /// + /// An in-memory stand-in for the platform state store. exposes a + /// protected constructor and virtual members precisely so it can be substituted like this. + /// + private sealed class FakeStateStore : FoundryStateStore + { + private readonly ConcurrentDictionary> _items = new(StringComparer.Ordinal); + + public IReadOnlyCollection> Items => (IReadOnlyCollection>)this._items.Values; + + public override string Name => FoundryAgentSessionStore.DefaultStoreName; + + public override Task SetItemAsync( + string key, + IDictionary value, + IReadOnlyDictionary? tags = null, + string? ifMatch = null, + bool requireExists = false, + CancellationToken cancellationToken = default) + { + this._items[key] = value; + return Task.FromResult(AzureAIAgentServerCoreStorageModelFactory.StateStoreItemRef(id: key, key: key, etag: "etag")); + } + + public override Task GetItemAsync(string key, CancellationToken cancellationToken = default) + => Task.FromResult(this._items.TryGetValue(key, out var value) + ? AzureAIAgentServerCoreStorageModelFactory.StateStoreItem(id: key, key: key, value: value, etag: "etag") + : null); + } + + private sealed class TestSession : AgentSession + { + } + + private sealed class TestAgent : AIAgent + { + private readonly string _serializedJson; + private readonly string? _name; + + public TestAgent(string serializedJson = "{}", string? name = null) + { + this._serializedJson = serializedJson; + this._name = name; + } + + public override string? Name => this._name; + + public int CreateCalls { get; private set; } + public int SerializeCalls { get; private set; } + public int DeserializeCalls { get; private set; } + public JsonElement? LastDeserialized { get; private set; } + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) + { + this.CreateCalls++; + return new ValueTask(new TestSession()); + } + + protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) + { + this.SerializeCalls++; + using var doc = JsonDocument.Parse(this._serializedJson); + return new ValueTask(doc.RootElement.Clone()); + } + + protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) + { + this.DeserializeCalls++; + this.LastDeserialized = serializedState.Clone(); + return new ValueTask(new TestSession()); + } + + protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) + => throw new NotSupportedException(); + + protected override IAsyncEnumerable RunCoreStreamingAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) + => throw new NotSupportedException(); + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs new file mode 100644 index 0000000000..a3f1982d5a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs @@ -0,0 +1,434 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Core.Storage; +using Microsoft.Agents.AI.Foundry.Hosting; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Extensions.Logging; + +namespace Microsoft.Agents.AI.Foundry.UnitTests.Hosting; + +public sealed class FoundryJsonCheckpointStoreTests +{ + [Fact] + public async Task CreateCheckpointAsync_ThenRetrieveCheckpointAsync_RoundTripsAsync() + { + // Arrange + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + + // Act + var key = await store.CreateCheckpointAsync("session-1", Json("{\"step\":3}")); + var value = await store.RetrieveCheckpointAsync("session-1", key); + + // Assert + Assert.Equal("session-1", key.SessionId); + Assert.Equal(3, value.GetProperty("step").GetInt32()); + } + + [Fact] + public async Task RetrieveIndexAsync_ReturnsCheckpointsInCommitOrderAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + + // Act + var first = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var second = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + var third = await store.CreateCheckpointAsync("session-1", Json("{\"step\":3}")); + var index = (await store.RetrieveIndexAsync("session-1")).ToList(); + + // Assert: the contract is oldest first, most recently committed last. + Assert.Equal([first, second, third], index); + } + + [Fact] + public async Task RetrieveIndexAsync_UnknownSession_ReturnsEmptyAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + + // Act + var index = await store.RetrieveIndexAsync("never-written"); + + // Assert + Assert.Empty(index); + } + + [Fact] + public async Task RetrieveIndexAsync_WithParent_ReturnsOnlyThatParentsChildrenAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + var root = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var childOfRoot = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}"), parent: root); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":3}"), parent: childOfRoot); + + // Act + var index = (await store.RetrieveIndexAsync("session-1", withParent: root)).ToList(); + + // Assert + Assert.Equal([childOfRoot], index); + } + + [Fact] + public async Task RetrieveIndexAsync_PartitionsBySessionAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + var forFirst = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var forSecond = await store.CreateCheckpointAsync("session-2", Json("{\"step\":1}")); + + // Act + var firstIndex = (await store.RetrieveIndexAsync("session-1")).ToList(); + var secondIndex = (await store.RetrieveIndexAsync("session-2")).ToList(); + + // Assert + Assert.Equal([forFirst], firstIndex); + Assert.Equal([forSecond], secondIndex); + } + + [Fact] + public async Task RetrieveCheckpointAsync_UnknownCheckpoint_ThrowsAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + + // Act & Assert + await Assert.ThrowsAsync( + async () => await store.RetrieveCheckpointAsync("session-1", new CheckpointInfo("session-1", "missing"))); + } + + [Fact] + public async Task RetrieveCheckpointAsync_CheckpointOfAnotherSession_ThrowsAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + var key = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + + // Act & Assert: the item key mixes in the session, so the same checkpoint id does not leak + // across sessions. + await Assert.ThrowsAsync( + async () => await store.RetrieveCheckpointAsync("session-2", new CheckpointInfo("session-2", key.CheckpointId))); + } + + [Fact] + public async Task CreateCheckpointAsync_LosesTheIndexRace_RetriesAndKeepsBothEntriesAsync() + { + // Arrange: the first index write is rejected as if another instance had committed a + // checkpoint for the same session in between the read and the write. + var backing = new FakeCheckpointStateStore { FailNextIndexWrites = 1 }; + var store = NewStore(backing); + + // Act + var first = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var second = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + var index = (await store.RetrieveIndexAsync("session-1")).ToList(); + + // Assert + Assert.Equal([first, second], index); + } + + [Fact] + public async Task CreateCheckpointAsync_UsesConditionalWriteOnAnExistingIndexAsync() + { + // Arrange + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + + // Act + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + + // Assert: the first index write creates the item, later ones carry the concurrency token. + Assert.Equal(1, backing.IndexCreateCalls); + Assert.Equal(["etag-2"], backing.ObservedIndexIfMatch); + } + + [Fact] + public async Task GetLatestCheckpointAsync_ReturnsTheMostRecentlyCommittedCheckpointAsync() + { + // Arrange + var store = NewStore(new FakeCheckpointStateStore()); + var manager = CheckpointManager.CreateJson(store); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var last = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + + // Act + var latest = await manager.GetLatestCheckpointAsync("session-1"); + + // Assert + Assert.Equal(last, latest); + } + + [Fact] + public void BuildCheckpointKey_StaysWithinThePlatformKeyLimit() + { + // Arrange + var longSessionId = new string('s', 4096); + var longCheckpointId = new string('c', 4096); + + // Act + var key = FoundryJsonCheckpointStore.BuildCheckpointKey(longSessionId, longCheckpointId); + var indexKey = FoundryJsonCheckpointStore.BuildIndexKey(longSessionId); + + // Assert + Assert.True(key.Length <= 128, $"Checkpoint key was {key.Length} characters."); + Assert.True(indexKey.Length <= 128, $"Index key was {indexKey.Length} characters."); + } + + [Fact] + public async Task RetrieveCheckpointAsync_DeletesEveryOtherCheckpointOfTheSessionAsync() + { + // Arrange: a session that ran three supersteps, so it holds three checkpoints. + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + var first = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var second = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + var third = await store.CreateCheckpointAsync("session-1", Json("{\"step\":3}")); + + // Act: resuming reads the latest checkpoint back. + var resumed = await store.RetrieveCheckpointAsync("session-1", third); + + // Assert: the resumed checkpoint is returned, and it is the only one left. + Assert.Equal("{\"step\":3}", resumed.GetRawText()); + Assert.Equal([third], (await store.RetrieveIndexAsync("session-1")).ToList()); + Assert.False(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", first.CheckpointId))); + Assert.False(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", second.CheckpointId))); + Assert.True(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", third.CheckpointId))); + } + + [Fact] + public async Task RetrieveCheckpointAsync_LeavesOtherSessionsAloneAsync() + { + // Arrange: two sessions, each holding checkpoints of its own. + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + var otherFirst = await store.CreateCheckpointAsync("session-2", Json("{\"step\":1}")); + var otherSecond = await store.CreateCheckpointAsync("session-2", Json("{\"step\":2}")); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var resumeTarget = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + + // Act + await store.RetrieveCheckpointAsync("session-1", resumeTarget); + + // Assert: pruning is scoped to the session that resumed. + Assert.Equal([otherFirst, otherSecond], (await store.RetrieveIndexAsync("session-2")).ToList()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_PruningFails_StillReturnsTheCheckpointAsync() + { + // Arrange: housekeeping is refused, which must not break a conversation that is resuming. + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var resumeTarget = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + backing.FailNextIndexWrites = 1; + + // Act + var resumed = await store.RetrieveCheckpointAsync("session-1", resumeTarget); + + // Assert + Assert.Equal("{\"step\":2}", resumed.GetRawText()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_PruningFailsForARealReason_ReportsItAndStillReturnsTheCheckpointAsync() + { + // Arrange: the store refuses the index write for a reason that is not a lost race, which is + // how a credential or network problem would show up. That must be traceable, because + // unreported it is how a session's checkpoints silently pile up. + var backing = new FakeCheckpointStateStore(); + var logs = new RecordingLoggerFactory(); + var store = NewStore(backing, logs); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var resumeTarget = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + backing.FailNextIndexWritesAuthentically = 1; + + // Act + var resumed = await store.RetrieveCheckpointAsync("session-1", resumeTarget); + + // Assert: the resume succeeds, and the failure is reported rather than swallowed. + Assert.Equal("{\"step\":2}", resumed.GetRawText()); + var warning = Assert.Single(logs.Entries, entry => entry.Level == LogLevel.Warning); + Assert.Contains("session-1", warning.Message, StringComparison.Ordinal); + Assert.NotNull(warning.Exception); + } + + [Fact] + public async Task RetrieveCheckpointAsync_PruningLosesARace_DoesNotWarnAsync() + { + // Arrange: losing to another writer is expected under concurrency and is not a problem, so + // it must not be reported as one. + var backing = new FakeCheckpointStateStore(); + var logs = new RecordingLoggerFactory(); + var store = NewStore(backing, logs); + await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var resumeTarget = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + backing.FailNextIndexWrites = 1; + + // Act + await store.RetrieveCheckpointAsync("session-1", resumeTarget); + + // Assert + Assert.DoesNotContain(logs.Entries, entry => entry.Level >= LogLevel.Warning); + } + + [Fact] + public void BuildCheckpointKey_DifferentSessionsNeverShareAKey() + { + // Act + var first = FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", "abc"); + var second = FoundryJsonCheckpointStore.BuildCheckpointKey("session-2", "abc"); + + // Assert + Assert.NotEqual(first, second); + } + + private static FoundryJsonCheckpointStore NewStore(FoundryStateStore backing, ILoggerFactory? loggerFactory = null) + => new(_ => Task.FromResult(backing), loggerFactory: loggerFactory); + + private static JsonElement Json(string json) + { + using var document = JsonDocument.Parse(json); + return document.RootElement.Clone(); + } + + /// + /// An in-memory stand-in for the platform store. It models the parts this store depends on: + /// item bodies keyed by item key, a concurrency token that changes on every write, and the + /// rejection of a conditional write whose token is stale. + /// + private sealed class FakeCheckpointStateStore : FoundryStateStore + { + private int _etagCounter; + + public override string Name => FoundryJsonCheckpointStore.DefaultStoreName; + + /// How many of the next index writes are rejected as having lost a race. + public int FailNextIndexWrites { get; set; } + + /// + /// How many of the next index writes are rejected for a reason that is not a lost race, + /// which is how a credential or network problem would reach this store. + /// + public int FailNextIndexWritesAuthentically { get; set; } + + public int IndexCreateCalls { get; private set; } + + public List ObservedIndexIfMatch { get; } = []; + + public override Task CreateItemAsync( + string key, + IDictionary value, + IReadOnlyDictionary? tags = null, + CancellationToken cancellationToken = default) + { + if (key.StartsWith("wi-", StringComparison.Ordinal)) + { + this.IndexCreateCalls++; + } + + if (this.Items.ContainsKey(key)) + { + throw new FoundryStorageConflictException("conflict"); + } + + return Task.FromResult(this.Write(key, value)); + } + + public override Task SetItemAsync( + string key, + IDictionary value, + IReadOnlyDictionary? tags = null, + string? ifMatch = null, + bool requireExists = false, + CancellationToken cancellationToken = default) + { + if (key.StartsWith("wi-", StringComparison.Ordinal)) + { + if (ifMatch is not null) + { + this.ObservedIndexIfMatch.Add(ifMatch); + } + + if (this.FailNextIndexWrites > 0) + { + this.FailNextIndexWrites--; + throw new FoundryStoragePreconditionException("precondition failed"); + } + + if (this.FailNextIndexWritesAuthentically > 0) + { + this.FailNextIndexWritesAuthentically--; + throw new FoundryStorageException(503, "service unavailable"); + } + } + + if (ifMatch is not null && (!this.Items.TryGetValue(key, out var existing) || existing.Etag != ifMatch)) + { + throw new FoundryStoragePreconditionException("precondition failed"); + } + + return Task.FromResult(this.Write(key, value)); + } + + public override Task GetItemAsync(string key, CancellationToken cancellationToken = default) + => Task.FromResult(this.Items.TryGetValue(key, out var entry) + ? AzureAIAgentServerCoreStorageModelFactory.StateStoreItem(id: key, key: key, value: entry.Value, etag: entry.Etag) + : null); + + public override Task DeleteItemAsync(string key, string? ifMatch = null, CancellationToken cancellationToken = default) + { + if (!this.Items.TryRemove(key, out _)) + { + throw new FoundryStorageNotFoundException("not found"); + } + + return Task.FromResult(AzureAIAgentServerCoreStorageModelFactory.DeletedStateStoreItem(id: key, deleted: true)); + } + + /// The item bodies currently held, so a test can assert what was deleted. + public ConcurrentDictionary Value, string Etag)> Items { get; } = new(StringComparer.Ordinal); + + private StateStoreItemRef Write(string key, IDictionary value) + { + string etag = string.Create(System.Globalization.CultureInfo.InvariantCulture, $"etag-{++this._etagCounter}"); + this.Items[key] = (value, etag); + return AzureAIAgentServerCoreStorageModelFactory.StateStoreItemRef(id: key, key: key, etag: etag); + } + } + + /// Captures what the store reported, so a test can assert on it. + private sealed class RecordingLoggerFactory : ILoggerFactory + { + public List<(LogLevel Level, string Message, Exception? Exception)> Entries { get; } = []; + + public ILogger CreateLogger(string categoryName) => new RecordingLogger(this.Entries); + + public void AddProvider(ILoggerProvider provider) + { + } + + public void Dispose() + { + } + + private sealed class RecordingLogger(List<(LogLevel Level, string Message, Exception? Exception)> entries) : ILogger + { + public IDisposable? BeginScope(TState state) where TState : notnull => null; + + public bool IsEnabled(LogLevel logLevel) => true; + + public void Log(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func formatter) + => entries.Add((logLevel, formatter(state, exception), exception)); + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedOutboundUserAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedOutboundUserAgentTests.cs index 090633f7d5..6e9838ab12 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedOutboundUserAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedOutboundUserAgentTests.cs @@ -114,7 +114,7 @@ private async Task StartHostedServerAsync() IChatClient chatClient = projectResponsesClient.AsIChatClient(Deployment); AIAgent agent = new ChatClientAgent(chatClient); - builder.Services.AddFoundryResponses(agent); + builder.Services.AddFoundryResponses(agent, new InMemoryAgentSessionStore()); builder.Services.AddSingleton(new FakeHostedSessionIsolationKeyProvider()); builder.Services.AddLogging(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs new file mode 100644 index 0000000000..197b9f1f42 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs @@ -0,0 +1,135 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Agents.AI.Workflows.InProc; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Diagnostics.HealthChecks; +using Moq; + +namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; + +/// +/// Covers the readiness check that reports a workflow agent built with a checkpoint manager of its +/// own, whose workflow state would be written somewhere hosting does not manage. +/// +public class HostedWorkflowCheckpointingHealthCheckTests +{ + [Fact] + public async Task CheckHealthAsync_WorkflowAgentLeavingStorageToHosting_IsHealthyAsync() + { + // Arrange: a workflow agent built the way a hosted container should build one. + var check = BuildCheckFor(BuildWorkflowAgent(executionEnvironment: null)); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Healthy, result.Status); + } + + [Fact] + public async Task CheckHealthAsync_WorkflowAgentWithItsOwnCheckpointManager_IsUnhealthyAsync() + { + // Arrange: the caller named where checkpoints go, so hosting leaves the agent alone and its + // workflow state never reaches the durable store. + var environment = InProcessExecution.OffThread.WithCheckpointing(CheckpointManager.CreateInMemory()); + var check = BuildCheckFor(BuildWorkflowAgent(environment)); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Unhealthy, result.Status); + var reported = Assert.IsType>(result.Data["agentsWithOwnCheckpointing"]); + Assert.Equal(["WorkflowAgent"], reported); + } + + [Fact] + public async Task CheckHealthAsync_WorkflowAgentBehindAWrapper_IsStillReportedAsync() + { + // Arrange: middleware around the agent must not hide the misconfiguration. + var environment = InProcessExecution.OffThread.WithCheckpointing(CheckpointManager.CreateInMemory()); + var check = BuildCheckFor(new PassThroughAgent(BuildWorkflowAgent(environment))); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Unhealthy, result.Status); + } + + [Fact] + public async Task CheckHealthAsync_AgentThatDoesNotRunAWorkflow_IsPassedOverAsync() + { + // Arrange: an agent with no checkpoints at all has nothing to misconfigure. + var check = BuildCheckFor(new ChatClientAgent(NewSilentChatClient(), new ChatClientAgentOptions { Name = "plain" })); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Healthy, result.Status); + } + + [Fact] + public async Task CheckHealthAsync_NotInAFoundryContainer_IsHealthyAsync() + { + // Arrange: outside a Foundry container nothing redirects checkpoints, so an agent bringing + // its own manager is not competing with anything. + var environment = InProcessExecution.OffThread.WithCheckpointing(CheckpointManager.CreateInMemory()); + var check = BuildCheckFor(BuildWorkflowAgent(environment)); + check.IsHosted = false; + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Healthy, result.Status); + } + + private static HostedWorkflowCheckpointingHealthCheck BuildCheckFor(AIAgent agent) + { + var services = new ServiceCollection(); + services.AddSingleton(agent); + + return new HostedWorkflowCheckpointingHealthCheck(services.BuildServiceProvider()) + { + IsHosted = true, + }; + } + + private static AIAgent BuildWorkflowAgent(InProcessExecutionEnvironment? executionEnvironment) + { + var inner = new ChatClientAgent(NewSilentChatClient(), new ChatClientAgentOptions { Name = "inner" }); + var workflow = new ConcurrentWorkflowBuilder(inner).WithOutputFrom(inner).Build(); + + return workflow.AsAIAgent( + id: "workflow-agent", + name: "WorkflowAgent", + executionEnvironment: executionEnvironment); + } + + private static HealthCheckContext NewContext() => new() + { + Registration = new HealthCheckRegistration( + "foundry-workflow-checkpointing", + _ => new Mock().Object, + HealthStatus.Unhealthy, + tags: null), + }; + + private static IChatClient NewSilentChatClient() + { + var client = new Mock(); + client.Setup(c => c.GetResponseAsync( + It.IsAny>(), It.IsAny(), It.IsAny())) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok"))); + return client.Object; + } + + private sealed class PassThroughAgent(AIAgent innerAgent) : DelegatingAIAgent(innerAgent); +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs new file mode 100644 index 0000000000..2655b49c73 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs @@ -0,0 +1,183 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.Agents.AI.Workflows.InProc; + +namespace Microsoft.Agents.AI.Workflows.UnitTests; + +/// +/// Covers , which lets a host redirect +/// where a already-built workflow agent writes its checkpoints. +/// +public class WorkflowHostingExtensionsTests +{ + [Fact] + public void WithCheckpointing_AgentThatDoesNotHostAWorkflow_IsLeftAlone() + { + // Arrange + AIAgent agent = new OrchestrationTestHelpers.DoubleEchoAgent("plain"); + + // Act + AIAgent result = agent.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Assert + Assert.Same(agent, result); + } + + [Fact] + public void WithCheckpointing_WorkflowAgentWithoutAnExplicitStore_IsRedirected() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act + AIAgent result = agent.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Assert: a copy is produced, and the original object is not modified. + Assert.NotSame(agent, result); + Assert.Equal(agent.Id, result.Id); + Assert.Equal(agent.Name, result.Name); + Assert.Equal(agent.Description, result.Description); + } + + [Fact] + public void WithCheckpointing_AppliedTwice_StopsAfterTheFirstRedirection() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + AIAgent redirected = agent.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Act + AIAgent again = redirected.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Assert: the copy already names a checkpoint manager, so it is now the explicit choice. + Assert.Same(redirected, again); + } + + [Fact] + public void WithCheckpointing_CallerAlreadyChoseAStore_DoesNotOverrideIt() + { + // Arrange + InProcessExecutionEnvironment callerChoice = InProcessExecution.Lockstep.WithCheckpointing(CheckpointManager.CreateInMemory()); + AIAgent agent = BuildWorkflowAgent(callerChoice); + + // Act + AIAgent result = agent.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Assert + Assert.Same(agent, result); + } + + [Fact] + public void WithCheckpointing_WorkflowAgentBehindAWrapper_IsLeftAlone() + { + // Arrange: only the innermost agent could be copied, and returning it alone would throw the + // wrapper away, so a wrapped workflow agent is deliberately not redirected. + AIAgent wrapper = new PassThroughAgent(BuildWorkflowAgent(executionEnvironment: null)); + + // Act + AIAgent result = wrapper.WithCheckpointing(CheckpointManager.CreateInMemory()); + + // Assert + Assert.Same(wrapper, result); + } + + [Fact] + public void WithCheckpointing_NullArguments_Throw() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act & Assert + Assert.Throws(() => ((AIAgent)null!).WithCheckpointing(CheckpointManager.CreateInMemory())); + Assert.Throws(() => agent.WithCheckpointing(null!)); + } + + [Fact] + public void GetService_WorkflowAgent_ExposesItsMetadata() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act + var metadata = agent.GetService(); + + // Assert: getting an instance back is what identifies a workflow agent. + Assert.NotNull(metadata); + Assert.False(metadata.UsesOwnCheckpointStorage); + } + + [Fact] + public void GetService_WorkflowAgent_DoesNotAnswerTheBaseAgentMetadataType() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act & Assert: AIAgentMetadata names the inference service behind an agent, which a + // workflow does not have, so the agent leaves that question unanswered as it did before. + Assert.Null(agent.GetService()); + } + + [Fact] + public void GetService_WorkflowAgentWithItsOwnStore_SaysSo() + { + // Arrange + InProcessExecutionEnvironment callerChoice = InProcessExecution.Lockstep.WithCheckpointing(CheckpointManager.CreateInMemory()); + AIAgent agent = BuildWorkflowAgent(callerChoice); + + // Act + var metadata = agent.GetService(); + + // Assert + Assert.NotNull(metadata); + Assert.True(metadata.UsesOwnCheckpointStorage); + } + + [Fact] + public void GetService_WorkflowAgentBehindAWrapper_IsStillFound() + { + // Arrange: detection has to see through middleware, which is why it goes through GetService + // rather than testing the type of the agent. + AIAgent wrapper = new PassThroughAgent(BuildWorkflowAgent(executionEnvironment: null)); + + // Act + var metadata = wrapper.GetService(); + + // Assert + Assert.NotNull(metadata); + } + + [Fact] + public void GetService_AgentThatDoesNotHostAWorkflow_HasNoWorkflowMetadata() + { + // Arrange + AIAgent agent = new OrchestrationTestHelpers.DoubleEchoAgent("plain"); + + // Act & Assert + Assert.Null(agent.GetService()); + } + + [Fact] + public void GetService_WithAServiceKey_ReturnsNothing() + { + // Arrange: a key means something this agent knows nothing about, so it must not answer. + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act & Assert + Assert.Null(agent.GetService(typeof(WorkflowAgentMetadata), serviceKey: "some-key")); + } + + private static AIAgent BuildWorkflowAgent(InProcessExecutionEnvironment? executionEnvironment) + { + OrchestrationTestHelpers.DoubleEchoAgent inner = new("inner"); + Workflow workflow = new ConcurrentWorkflowBuilder(inner).WithOutputFrom(inner).Build(); + + return workflow.AsAIAgent( + id: "workflow-agent", + name: "WorkflowAgent", + description: "A workflow hosted as an agent.", + executionEnvironment: executionEnvironment); + } + + private sealed class PassThroughAgent(AIAgent innerAgent) : DelegatingAIAgent(innerAgent); +} From f73cc27d5ace74e8a9b57c034bdcf0939274c7ca Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:52:03 +0100 Subject: [PATCH 3/6] Keep the readiness probe from running the agent's providers The stored-output probe ran the registered agent with its chat client replaced, which still set the agent's chat history provider and context providers running. Those are the parts most likely to reach outside the container and to write state, so every readiness probe could make external calls and add its own empty turn to real conversations. The probe now runs a stand-in built from the agent's own options with both kinds of provider dropped. It keeps what decides the setting, the chat options and the raw request factory, and cannot see a decorator wrapped around the agent, which is accepted for a readiness check. --- .../HostedStoredOutputHealthCheck.cs | 30 ++++++--- .../HostedStoredOutputHealthCheckTests.cs | 67 +++++++++++++++++++ 2 files changed, 87 insertions(+), 10 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs index 4b9429b617..b81d23f7f5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs @@ -97,30 +97,40 @@ public async Task CheckHealthAsync(HealthCheckContext context } /// - /// Runs the agent with its chat client replaced by one that calls nothing, and reports whether the - /// request the agent built asks for the response to be stored. + /// Runs a stand-in built from the agent's own configuration, with its chat client replaced by one + /// that calls nothing, and reports whether the request that configuration produces asks for the + /// response to be stored. /// /// /// The run carries no chat options of its own, so the agent's own configuration is what reaches the /// probe. Overriding the setting here, the way the request handler does per turn, would only show /// the override back. /// - /// The agent's chat history provider is stood down for this run, because it would otherwise read - /// and write its own store on every readiness probe. A provider backed by a database would then be - /// doing external calls, and adding this probe's empty turn to a real conversation, for a run that - /// asks the agent nothing. + /// A stand-in is built rather than running the registered agent because that agent's chat history + /// provider and context providers would run with it. Those are the parts most likely to reach + /// outside the container, a memory or search provider for instance, and to write state: a readiness + /// probe would then make external calls and add its own empty turn to real conversations, on every + /// probe, for a run that asks the agent nothing. The stand-in keeps everything that decides the + /// stored output setting, the chat options and the raw request factory among them, and drops both + /// kinds of provider, so the probe stays free of side effects. + /// + /// + /// What the stand-in cannot see is a decorator wrapped around the agent. One that changed this + /// setting would go unreported, which is accepted here: the check exists to catch how a container + /// configured its agent, and a silent probe is worth more than a complete one. /// /// private async Task StoresItsOwnResponsesAsync(AIAgent agent, CancellationToken cancellationToken) { var probe = new StoredOutputProbeChatClient(); - var runOptions = new ChatClientAgentRunOptions { ChatClientFactory = _ => probe }; - runOptions.AdditionalProperties ??= []; - runOptions.AdditionalProperties.Add(new VolatileChatHistoryProvider()); + var probeOptions = agent.GetService()?.Clone() ?? new ChatClientAgentOptions(); + probeOptions.ChatHistoryProvider = null; + probeOptions.AIContextProviders = null; try { - await agent.RunAsync([], options: runOptions, cancellationToken: cancellationToken).ConfigureAwait(false); + var probeAgent = new ChatClientAgent(probe, probeOptions); + await probeAgent.RunAsync([], cancellationToken: cancellationToken).ConfigureAwait(false); } catch (Exception ex) when (ex is not OperationCanceledException || !cancellationToken.IsCancellationRequested) { 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 b44b5f96cc..a6f2eb0349 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs @@ -124,6 +124,37 @@ public async Task CheckHealthAsync_AgentThatIsNotAChatClientAgent_IsHealthyAsync Assert.Equal(HealthStatus.Healthy, result.Status); } + [Fact] + public async Task CheckHealthAsync_AgentWithProviders_LeavesThemUntouchedAsync() + { + // Arrange: a context provider and a chat history provider are the parts most likely to call + // outside the container and to write state, so a readiness probe must not set them running. + var contextProvider = new RecordingContextProvider(); + var historyProvider = new RecordingChatHistoryProvider(); + var agent = new ChatClientAgent( + NewSilentChatClient(), + new ChatClientAgentOptions + { + Name = "has-providers", + ChatHistoryProvider = historyProvider, + AIContextProviders = [contextProvider], + ChatOptions = new ChatOptions + { + RawRepresentationFactory = _ => new CreateResponseOptions { StoredOutputEnabled = true }, + }, + }); + + var check = BuildCheckFor(agent); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert: the setting is still read, and neither provider was asked to do anything. + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.False(contextProvider.WasInvoked); + Assert.False(historyProvider.WasInvoked); + } + private static HostedStoredOutputHealthCheck BuildCheckFor(AIAgent agent, FoundryResponsesOptions? hostingOptions = null) { var services = new ServiceCollection(); @@ -160,4 +191,40 @@ private static async IAsyncEnumerable OneUpdateAsync() await Task.CompletedTask; yield return new ChatResponseUpdate(ChatRole.Assistant, "ok"); } + + /// Records whether the agent ever set it running. Stands in for a memory or search provider. + private sealed class RecordingContextProvider : AIContextProvider + { + public bool WasInvoked { get; private set; } + + protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + this.WasInvoked = true; + return new(new AIContext()); + } + + protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.WasInvoked = true; + return default; + } + } + + /// Records whether the agent ever set it running. Stands in for a database-backed history store. + private sealed class RecordingChatHistoryProvider : ChatHistoryProvider + { + public bool WasInvoked { get; private set; } + + protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + this.WasInvoked = true; + return new([]); + } + + protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.WasInvoked = true; + return default; + } + } } From 0f0fa8f1143c9f6698c98c237a55512fb73c1431 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:26:11 +0100 Subject: [PATCH 4/6] build: bump AgentServer preview packages Core beta.29 adds the shared local state-store fallback used by hosted sessions and workflow checkpoints. Align its Azure Core and System package dependencies to avoid assembly and downgrade conflicts. --- dotnet/Directory.Packages.props | 14 +++++++------- .../Microsoft.Agents.AI.Foundry.UnitTests.csproj | 1 + 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index ba7f5601c5..a45e2d7cd9 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -23,14 +23,14 @@ - - - + + + - + @@ -42,17 +42,17 @@ - + - + - + diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/Microsoft.Agents.AI.Foundry.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/Microsoft.Agents.AI.Foundry.UnitTests.csproj index 7862225653..6ee4bf83a4 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/Microsoft.Agents.AI.Foundry.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/Microsoft.Agents.AI.Foundry.UnitTests.csproj @@ -6,6 +6,7 @@ + From 9da7d9d1e7eebda4dcc959fa4441e2a8c74314bc Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:26:31 +0100 Subject: [PATCH 5/6] feat(foundry): use AgentServer state fallback Use FoundryStateStore for sessions and workflow checkpoints in every environment. Core beta.29 selects Foundry Storage when hosted and a file-backed local store otherwise, so local runs exercise the production storage shape without requiring Azure credentials. Give the hosted workflow sample stable inner-agent identities so its checkpoints remain compatible after container replacement. --- .../Hosted-Workflow-Simple/Program.cs | 29 ++++- .../FoundryAgentSessionStore.cs | 26 ++-- .../FoundryJsonCheckpointStore.cs | 23 ++-- .../ServiceCollectionExtensions.cs | 39 +++--- .../FoundryAgentSessionStoreTests.cs | 24 ++-- .../FoundryJsonCheckpointStoreTests.cs | 10 ++ .../FoundryStateStoreLocalFallbackTests.cs | 113 ++++++++++++++++++ ...edWorkflowCheckpointingHealthCheckTests.cs | 14 +++ .../ServiceCollectionExtensionsTests.cs | 15 +++ 9 files changed, 234 insertions(+), 59 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryStateStoreLocalFallbackTests.cs diff --git a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-Workflow-Simple/Program.cs b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-Workflow-Simple/Program.cs index 6ed5f33601..d8bc908573 100644 --- a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-Workflow-Simple/Program.cs +++ b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-Workflow-Simple/Program.cs @@ -20,7 +20,10 @@ string endpoint = Environment.GetEnvironmentVariable("FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("FOUNDRY_PROJECT_ENDPOINT is not set."); -string deploymentName = Environment.GetEnvironmentVariable("FOUNDRY_MODEL") ?? "gpt-4o"; +string deploymentName = + Environment.GetEnvironmentVariable("AZURE_AI_MODEL_DEPLOYMENT_NAME") + ?? Environment.GetEnvironmentVariable("FOUNDRY_MODEL") + ?? "gpt-4o"; // WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production. // In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid @@ -37,10 +40,26 @@ .GetChatClient(deploymentName) .AsIChatClient(); -// Create translation agents -AIAgent frenchAgent = chatClient.AsAIAgent("You are a translation assistant that translates the provided text to French."); -AIAgent spanishAgent = chatClient.AsAIAgent("You are a translation assistant that translates the provided text to Spanish."); -AIAgent englishAgent = chatClient.AsAIAgent("You are a translation assistant that translates the provided text to English."); +// A workflow checkpoint records each executor identity. Keep both Id and Name stable so a new +// container instance reconstructs the same workflow and can resume checkpoints written earlier. +AIAgent frenchAgent = chatClient.AsAIAgent(new ChatClientAgentOptions +{ + Id = "french-translator", + Name = "French Translator", + ChatOptions = new() { Instructions = "You are a translation assistant that translates the provided text to French." }, +}); +AIAgent spanishAgent = chatClient.AsAIAgent(new ChatClientAgentOptions +{ + Id = "spanish-translator", + Name = "Spanish Translator", + ChatOptions = new() { Instructions = "You are a translation assistant that translates the provided text to Spanish." }, +}); +AIAgent englishAgent = chatClient.AsAIAgent(new ChatClientAgentOptions +{ + Id = "english-translator", + Name = "English Translator", + ChatOptions = new() { Instructions = "You are a translation assistant that translates the provided text to English." }, +}); // Build the sequential workflow: French → Spanish → English AIAgent agent = new WorkflowBuilder(frenchAgent) diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs index 73b826390c..f6240eb88b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs @@ -17,16 +17,14 @@ namespace Microsoft.Agents.AI.Foundry.Hosting; /// /// Provides an that persists the agent-framework's serialized -/// state to the Foundry platform's durable state-store API -/// () instead of to the container's local disk. +/// state through . /// /// /// -/// This is the durable counterpart to . The file-system -/// store writes under the container's session directory, which belongs to a single container -/// instance: the state is lost when that container is replaced and cannot be read by another -/// instance. This store writes to the platform instead, so a session survives container restarts -/// and replacement and is visible to every instance of the agent. +/// The AgentServer SDK selects the backend. In Foundry hosting it writes to the platform's durable +/// state store, so a session survives container replacement and is visible to every instance of the +/// agent. Outside Foundry hosting it uses the SDK's local state-store fallback under +/// ~/.agentserver/state_stores. /// /// /// Layout. All sessions live in one state store, named unless @@ -69,11 +67,14 @@ public sealed class FoundryAgentSessionStore : AgentSessionStore /// /// Initializes a new instance of the class. /// - /// The credential used to authenticate to the Foundry storage API. + /// + /// The credential used to authenticate to the Foundry storage API. May be + /// outside Foundry hosting, where the AgentServer SDK uses its local state-store fallback. + /// /// - /// The Foundry project endpoint. When , it is read from the - /// FOUNDRY_PROJECT_ENDPOINT environment variable, which the platform sets in a hosted - /// container. + /// The Foundry project endpoint. Used only in Foundry hosting. When , + /// it is read from the FOUNDRY_PROJECT_ENDPOINT environment variable. Outside Foundry + /// hosting, the AgentServer SDK ignores it and uses its local state-store fallback. /// /// The state-store name to hold the sessions. Defaults to . /// @@ -83,12 +84,11 @@ public sealed class FoundryAgentSessionStore : AgentSessionStore /// platform fixes it at creation. /// public FoundryAgentSessionStore( - TokenCredential credential, + TokenCredential? credential = null, Uri? endpoint = null, string storeName = DefaultStoreName, int itemTtlSeconds = FoundryStateStore.DefaultItemTtlSeconds) { - _ = Throw.IfNull(credential); _ = Throw.IfNullOrWhitespace(storeName); this.StoreName = storeName; diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs index 55809bd32f..bfc00c4a40 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs @@ -19,11 +19,16 @@ namespace Microsoft.Agents.AI.Foundry.Hosting; /// -/// Provides a that persists workflow checkpoints to the Foundry -/// platform's durable state-store API (). +/// Provides a that persists workflow checkpoints through +/// . /// /// /// +/// The AgentServer SDK selects the backend. In Foundry hosting it writes to the platform's durable +/// state store. Outside Foundry hosting it uses the SDK's local state-store fallback under +/// ~/.agentserver/state_stores. +/// +/// /// Item keys are hashes of the session identifier and the checkpoint identifier, because the /// platform limits an item key to 128 characters and neither identifier is bounded. Hashing rather /// than truncating means two different checkpoints can never end up sharing a key and overwriting @@ -83,11 +88,14 @@ public sealed class FoundryJsonCheckpointStore : JsonCheckpointStore /// /// Initializes a new instance of the class. /// - /// The credential used to authenticate to the Foundry storage API. + /// + /// The credential used to authenticate to the Foundry storage API. May be + /// outside Foundry hosting, where the AgentServer SDK uses its local state-store fallback. + /// /// - /// The Foundry project endpoint. When , it is read from the - /// FOUNDRY_PROJECT_ENDPOINT environment variable, which the platform sets in a hosted - /// container. + /// The Foundry project endpoint. Used only in Foundry hosting. When , + /// it is read from the FOUNDRY_PROJECT_ENDPOINT environment variable. Outside Foundry + /// hosting, the AgentServer SDK ignores it and uses its local state-store fallback. /// /// The state-store name to hold the checkpoints. Defaults to . /// @@ -102,13 +110,12 @@ public sealed class FoundryJsonCheckpointStore : JsonCheckpointStore /// happens in. /// public FoundryJsonCheckpointStore( - TokenCredential credential, + TokenCredential? credential = null, Uri? endpoint = null, string storeName = DefaultStoreName, int itemTtlSeconds = FoundryStateStore.DefaultItemTtlSeconds, ILoggerFactory? loggerFactory = null) { - _ = Throw.IfNull(credential); _ = Throw.IfNullOrWhitespace(storeName); this.StoreName = storeName; diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs index 2f0d2a4ce1..91cb878688 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs @@ -67,7 +67,7 @@ public static IServiceCollection AddFoundryResponses(this IServiceCollection ser services.AddHealthChecks(); ConfigureFoundryListenPort(services); ConfigureFoundryResponsesOptions(services, configure); - services.TryAddSingleton(_ => CreateDefaultAgentSessionStore()); + services.TryAddSingleton(_ => CreateDefaultAgentSessionStore()); services.TryAddSingleton(); return services; } @@ -94,7 +94,7 @@ 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 is chosen for the environment: the Foundry durable state store when the platform has supplied a project endpoint (FOUNDRY_PROJECT_ENDPOINT), otherwise a file-system store rooted at {$HOME}/.checkpoints when hosted and {cwd}/.checkpoints locally. + /// The agent session store to use for managing agent sessions server-side. If null, is used: the Foundry durable state store when hosted, and the AgentServer SDK's local state-store fallback otherwise. /// /// Optional callback to configure , for example to allow the /// agent's own service to store the responses it produces. @@ -358,15 +358,12 @@ private static void AddResponsesServerOnce(IServiceCollection services) /// Creates the used when the caller did not supply one. /// /// - /// Inside a Foundry container, sessions are held by the platform's durable state store, because - /// that state survives the container being restarted or replaced and is readable by every - /// instance of the agent. Anywhere else there is no such service to call, so the container falls - /// back to writing session files under its own session directory, which is what a local run does. + /// The AgentServer SDK selects the backend. Inside a Foundry container it uses the platform's + /// durable state store, which survives replacement and is readable by every instance. Anywhere + /// else it uses the SDK's local state-store fallback under ~/.agentserver/state_stores. /// - private static AgentSessionStore CreateDefaultAgentSessionStore() => - FoundryEnvironment.IsHosted - ? new FoundryAgentSessionStore(new DefaultAzureCredential()) - : FileSystemAgentSessionStore.CreateDefault(); + private static FoundryAgentSessionStore CreateDefaultAgentSessionStore() => + new(CreateStateStoreCredential()); /// /// Every agent a container can serve: the ones registered under a name, plus the default. @@ -539,10 +536,10 @@ internal static AIAgent ApplyOpenTelemetry(AIAgent agent) /// only the pointer to the last one. /// /// - /// The method is a no-op when the agent does not host a workflow, when the workflow was built - /// with an explicit checkpoint manager, and when the process is not running on the platform. - /// The redirected agent is cached against the agent it came from, so the substitution happens - /// once rather than on every request. + /// The method is a no-op when the agent does not host a workflow or when the workflow was built + /// with an explicit checkpoint manager. The AgentServer SDK selects the hosted or local + /// state-store backend. The redirected agent is cached against the agent it came from, so the + /// substitution happens once rather than on every request. /// /// /// The resolved agent. @@ -550,11 +547,6 @@ internal static AIAgent ApplyOpenTelemetry(AIAgent agent) /// The agent to serve the request with. internal static AIAgent ApplyWorkflowCheckpointing(AIAgent agent, ILoggerFactory? loggerFactory = null) { - if (!FoundryEnvironment.IsHosted) - { - return agent; - } - return s_workflowCheckpointingAgents.GetValue( agent, source => source.WithCheckpointing(GetFoundryWorkflowCheckpointManager(loggerFactory))); @@ -570,10 +562,17 @@ private static CheckpointManager GetFoundryWorkflowCheckpointManager(ILoggerFact lock (s_checkpointManagerGate) { return s_foundryWorkflowCheckpointManager ??= CheckpointManager.CreateJson( - new FoundryJsonCheckpointStore(new DefaultAzureCredential(), loggerFactory: loggerFactory)); + new FoundryJsonCheckpointStore(CreateStateStoreCredential(), loggerFactory: loggerFactory)); } } + /// + /// Creates the credential required by the hosted state-store backend. The beta.29 SDK requires + /// no credential for its local fallback, so local development does not construct one. + /// + private static DefaultAzureCredential? CreateStateStoreCredential() => + FoundryEnvironment.IsHosted ? new DefaultAzureCredential() : null; + private static readonly object s_checkpointManagerGate = new(); private static CheckpointManager? s_foundryWorkflowCheckpointManager; diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs index 23ded1512c..d6299ff209 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs @@ -211,26 +211,24 @@ public void BuildItemKey_IsStableAndDistinctPerLogicalKey() [Fact] public void Constructor_NullOrWhitespaceStoreName_Throws() { - // Arrange - var credential = new FakeCredential(); - // Act / Assert - Assert.Throws(() => new FoundryAgentSessionStore(credential, storeName: null!)); - Assert.Throws(() => new FoundryAgentSessionStore(credential, storeName: " ")); + Assert.Throws(() => new FoundryAgentSessionStore(storeName: null!)); + Assert.Throws(() => new FoundryAgentSessionStore(storeName: " ")); } - private static FoundryAgentSessionStore NewStore(FakeStateStore backing) - => new(_ => Task.FromResult(backing)); - - private sealed class FakeCredential : Azure.Core.TokenCredential + [Fact] + public void Constructor_WithoutCredential_IsAllowedForTheSdkLocalFallback() { - public override Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, CancellationToken cancellationToken) - => new("token", DateTimeOffset.MaxValue); + // Act + var store = new FoundryAgentSessionStore(); - public override ValueTask GetTokenAsync(Azure.Core.TokenRequestContext requestContext, CancellationToken cancellationToken) - => new(new Azure.Core.AccessToken("token", DateTimeOffset.MaxValue)); + // Assert + Assert.Equal(FoundryAgentSessionStore.DefaultStoreName, store.StoreName); } + private static FoundryAgentSessionStore NewStore(FakeStateStore backing) + => new(_ => Task.FromResult(backing)); + /// /// An in-memory stand-in for the platform state store. exposes a /// protected constructor and virtual members precisely so it can be substituted like this. diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs index a3f1982d5a..1e7621889e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs @@ -16,6 +16,16 @@ namespace Microsoft.Agents.AI.Foundry.UnitTests.Hosting; public sealed class FoundryJsonCheckpointStoreTests { + [Fact] + public void Constructor_WithoutCredential_IsAllowedForTheSdkLocalFallback() + { + // Act + var store = new FoundryJsonCheckpointStore(); + + // Assert + Assert.Equal(FoundryJsonCheckpointStore.DefaultStoreName, store.StoreName); + } + [Fact] public async Task CreateCheckpointAsync_ThenRetrieveCheckpointAsync_RoundTripsAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryStateStoreLocalFallbackTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryStateStoreLocalFallbackTests.cs new file mode 100644 index 0000000000..68e7fb1798 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryStateStoreLocalFallbackTests.cs @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.IO; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; + +[CollectionDefinition(CollectionName, DisableParallelization = true)] +public sealed class FoundryStateStoreLocalFallbackCollectionDefinition +{ + public const string CollectionName = "Foundry state-store local fallback"; +} + +[Collection(FoundryStateStoreLocalFallbackCollectionDefinition.CollectionName)] +public sealed class FoundryStateStoreLocalFallbackTests +{ + [Fact] + public async Task StoresWithoutCredential_RoundTripThroughTheSdkLocalFallbackAsync() + { + // Arrange + string root = Path.Combine(Path.GetTempPath(), $"foundry-state-store-local-{Guid.NewGuid():N}"); + string? previousStateRoot = Environment.GetEnvironmentVariable("AGENTSERVER_STATE_ROOT"); + string? previousHostingEnvironment = Environment.GetEnvironmentVariable("FOUNDRY_HOSTING_ENVIRONMENT"); + + try + { + Environment.SetEnvironmentVariable("AGENTSERVER_STATE_ROOT", root); + Environment.SetEnvironmentVariable("FOUNDRY_HOSTING_ENVIRONMENT", null); + + var agent = new TestAgent(); + var sessionStore = new FoundryAgentSessionStore(); + var checkpointStore = new FoundryJsonCheckpointStore(); + + // Act + await sessionStore.SaveSessionAsync(agent, "conversation-1", new TestSession(), userId: "user-1"); + AgentSession? session = await sessionStore.GetSessionAsync(agent, "conversation-1", userId: "user-1"); + + using JsonDocument document = JsonDocument.Parse("""{"step":1}"""); + CheckpointInfo checkpointInfo = await checkpointStore.CreateCheckpointAsync( + "workflow-session-1", + document.RootElement.Clone()); + JsonElement checkpoint = await checkpointStore.RetrieveCheckpointAsync( + "workflow-session-1", + checkpointInfo); + + // Assert + Assert.NotNull(session); + Assert.Equal("saved", agent.LastDeserialized?.GetProperty("session").GetString()); + Assert.Equal(1, checkpoint.GetProperty("step").GetInt32()); + Assert.True(Directory.Exists(Path.Combine(root, "state_stores"))); + } + finally + { + Environment.SetEnvironmentVariable("AGENTSERVER_STATE_ROOT", previousStateRoot); + Environment.SetEnvironmentVariable("FOUNDRY_HOSTING_ENVIRONMENT", previousHostingEnvironment); + + if (Directory.Exists(root)) + { + Directory.Delete(root, recursive: true); + } + } + } + + private sealed class TestSession : AgentSession; + + private sealed class TestAgent : AIAgent + { + public override string? Name => "local-fallback-agent"; + + public JsonElement? LastDeserialized { get; private set; } + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => + new(new TestSession()); + + protected override ValueTask SerializeSessionCoreAsync( + AgentSession session, + JsonSerializerOptions? jsonSerializerOptions = null, + CancellationToken cancellationToken = default) + { + using JsonDocument document = JsonDocument.Parse("""{"session":"saved"}"""); + return new(document.RootElement.Clone()); + } + + protected override ValueTask DeserializeSessionCoreAsync( + JsonElement serializedState, + JsonSerializerOptions? jsonSerializerOptions = null, + CancellationToken cancellationToken = default) + { + this.LastDeserialized = serializedState.Clone(); + return new(new TestSession()); + } + + protected override Task RunCoreAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + protected override IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + } +} 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 197b9f1f42..fc33c4af41 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs @@ -91,6 +91,20 @@ public async Task CheckHealthAsync_NotInAFoundryContainer_IsHealthyAsync() Assert.Equal(HealthStatus.Healthy, result.Status); } + [Fact] + public void ApplyWorkflowCheckpointing_NotInAFoundryContainer_UsesTheSdkLocalFallback() + { + // Arrange + AIAgent agent = BuildWorkflowAgent(executionEnvironment: null); + + // Act + AIAgent result = FoundryHostingExtensions.ApplyWorkflowCheckpointing(agent); + + // Assert: a redirected copy is returned even locally; the SDK chooses its local backend. + Assert.NotSame(agent, result); + Assert.True(result.GetService()?.UsesOwnCheckpointStorage); + } + private static HostedWorkflowCheckpointingHealthCheck BuildCheckFor(AIAgent agent) { var services = new ServiceCollection(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ServiceCollectionExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ServiceCollectionExtensionsTests.cs index 472f7f2b4e..d3162ef8df 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ServiceCollectionExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ServiceCollectionExtensionsTests.cs @@ -32,6 +32,21 @@ public void AddFoundryResponses_RegistersResponseHandler() Assert.Equal(typeof(AgentFrameworkResponseHandler), descriptor.ImplementationType); } + [Fact] + public void AddFoundryResponses_UsesTheStateStoreAdapterByDefault() + { + // Arrange + var services = new ServiceCollection(); + services.AddLogging(); + + // Act + services.AddFoundryResponses(); + using var provider = services.BuildServiceProvider(); + + // Assert: the AgentServer SDK behind this adapter chooses the hosted or local backend. + Assert.IsType(provider.GetRequiredService()); + } + [Fact] public void AddFoundryResponses_CalledTwice_RegistersOnce() { From 08b57eb9a00ee4659621fe6b63f9e3ca61ece153 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Fri, 14 Aug 2026 16:02:43 +0100 Subject: [PATCH 6/6] fix(hosting): harden durable state storage Use published AgentServer packages so CI no longer depends on a local package source. --- dotnet/Directory.Packages.props | 6 +- dotnet/nuget.config | 4 - .../AgentFrameworkResponseHandler.cs | 73 +++++++-- .../FoundryAgentSessionStore.cs | 80 +++++++--- .../FoundryJsonCheckpointStore.cs | 50 +++--- .../FoundryStateStoreBinding.cs | 7 +- .../HostedStoredOutputHealthCheck.cs | 27 ++-- .../HostedWorkflowCheckpointingHealthCheck.cs | 18 ++- .../FoundryAgentSessionStoreTests.cs | 147 ++++++++++++++++-- .../FoundryJsonCheckpointStoreTests.cs | 23 ++- .../HostedStoredOutputHealthCheckTests.cs | 45 ++++++ ...edWorkflowCheckpointingHealthCheckTests.cs | 17 +- .../WorkflowHostingExtensionsTests.cs | 2 +- 13 files changed, 392 insertions(+), 107 deletions(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index a45e2d7cd9..f2f4ae89c5 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -23,9 +23,9 @@ - - - + + + diff --git a/dotnet/nuget.config b/dotnet/nuget.config index 3cde656b2a..128d95e590 100644 --- a/dotnet/nuget.config +++ b/dotnet/nuget.config @@ -3,14 +3,10 @@ - - - - diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index d9cbfc8d7b..b7a8d4d61e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -69,7 +69,7 @@ public override async IAsyncEnumerable CreateAsync( [EnumeratorCancellation] CancellationToken cancellationToken) { // 1. Resolve agent - var agent = this.ResolveAgent(request); + var agent = this.ResolveAgent(request, out string agentStorageIdentity); var sessionStore = this.ResolveSessionStore(request); // Fail fast with a clear, actionable error when this 2.0.0-only image is served container @@ -131,9 +131,28 @@ public override async IAsyncEnumerable CreateAsync( // Load the session for this conversation, or start a new one. The store returns null when // nothing is persisted for the key, so a fresh conversation and a resumed one both end up with // a session to run against. - AgentSession? session = !string.IsNullOrWhiteSpace(agentSessionId) - ? await sessionStore.GetOrCreateSessionAsync(agent, agentSessionId, resolvedUserId, cancellationToken).ConfigureAwait(false) - : await agent.CreateSessionAsync(cancellationToken).ConfigureAwait(false); + AgentSession? session; + if (string.IsNullOrWhiteSpace(agentSessionId)) + { + session = await agent.CreateSessionAsync(cancellationToken).ConfigureAwait(false); + } + else + { + session = sessionStore is FoundryAgentSessionStore foundrySessionStore + ? await foundrySessionStore.GetSessionAsync( + agent, + agentStorageIdentity, + agentSessionId, + resolvedUserId, + cancellationToken).ConfigureAwait(false) + : await sessionStore.GetSessionAsync( + agent, + agentSessionId, + resolvedUserId, + cancellationToken).ConfigureAwait(false); + + session ??= await agent.CreateSessionAsync(cancellationToken).ConfigureAwait(false); + } // Capture the platform per-request call id (x-agent-foundry-call-id, protocol 2.0.0 only). // It is re-applied to the ambient HostedCallContext immediately before each outbound egress @@ -509,7 +528,25 @@ bool CheckNotAllowedStoreUsage() => // Persist the session for the next turn of this conversation, unless this one is being failed. if (session is not null && !turnFailed) { - await sessionStore.SaveSessionAsync(agent, agentSessionId!, session, resolvedUserId, cancellationToken).ConfigureAwait(false); + if (sessionStore is FoundryAgentSessionStore foundrySessionStore) + { + await foundrySessionStore.SaveSessionAsync( + agent, + agentStorageIdentity, + agentSessionId!, + session, + resolvedUserId, + cancellationToken).ConfigureAwait(false); + } + else + { + await sessionStore.SaveSessionAsync( + agent, + agentSessionId!, + session, + resolvedUserId, + cancellationToken).ConfigureAwait(false); + } } } @@ -572,7 +609,7 @@ private static string NewOAuthConsentItemId() /// Tries agent.name first, then falls back to metadata["entity_id"]. /// If neither is present, attempts to resolve a default (non-keyed) . /// - private AIAgent ResolveAgent(CreateResponse request) + private AIAgent ResolveAgent(CreateResponse request, out string storageIdentity) { var agentName = GetAgentName(request); @@ -581,9 +618,8 @@ private AIAgent ResolveAgent(CreateResponse request) var agent = this._serviceProvider.GetKeyedService(agentName); if (agent is not null) { - FoundryHostingExtensions.TryApplyUserAgent(agent); - return FoundryHostingExtensions.ApplyOpenTelemetry( - FoundryHostingExtensions.ApplyWorkflowCheckpointing(agent, this._serviceProvider.GetService())); + storageIdentity = $"key:{agentName}"; + return this.PrepareResolvedAgent(agent); } if (this._logger.IsEnabled(LogLevel.Warning)) @@ -596,9 +632,11 @@ private AIAgent ResolveAgent(CreateResponse request) var defaultAgent = this._serviceProvider.GetService(); if (defaultAgent is not null) { - FoundryHostingExtensions.TryApplyUserAgent(defaultAgent); - return FoundryHostingExtensions.ApplyOpenTelemetry( - FoundryHostingExtensions.ApplyWorkflowCheckpointing(defaultAgent, this._serviceProvider.GetService())); + storageIdentity = !string.IsNullOrWhiteSpace(defaultAgent.Name) + ? $"name:{defaultAgent.Name}" + : "default"; + + return this.PrepareResolvedAgent(defaultAgent); } var errorMessage = string.IsNullOrEmpty(agentName) @@ -608,6 +646,17 @@ private AIAgent ResolveAgent(CreateResponse request) throw new InvalidOperationException(errorMessage); } + private AIAgent PrepareResolvedAgent(AIAgent agent) + { + FoundryHostingExtensions.TryApplyUserAgent(agent); + + AIAgent prepared = FoundryHostingExtensions.ApplyWorkflowCheckpointing( + agent, + this._serviceProvider.GetService()); + + return FoundryHostingExtensions.ApplyOpenTelemetry(prepared); + } + /// /// Resolves an from the request. /// Tries agent.name first, then falls back to metadata["entity_id"]. diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs index f6240eb88b..7f37f0c5ae 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs @@ -29,12 +29,10 @@ namespace Microsoft.Agents.AI.Foundry.Hosting; /// /// Layout. All sessions live in one state store, named unless /// overridden, and each (agent, user, conversation) triple is one item in it. The item key is a -/// hash of the same a-/u-/c- logical key that -/// and use, so -/// all three stores partition sessions identically. Hashing is required because the platform -/// limits an item key to 128 characters, which an agent name plus a user id plus a conversation id -/// can exceed. The readable logical key is stored alongside the session in the item body so a -/// stored item can still be traced back to its conversation. +/// hash of an unambiguous, length-prefixed encoding of the hosted registration identity, user id, +/// and conversation id. Hashing is required because the platform limits an item key to 128 +/// characters. The readable encoding is stored alongside the session so an item can still be traced +/// back to its partition. /// /// /// Per-user isolation is expressed through the item key rather than through the state store's own @@ -119,21 +117,37 @@ internal FoundryAgentSessionStore(Func - public override async ValueTask SaveSessionAsync( + public override ValueTask SaveSessionAsync( AIAgent agent, string conversationId, AgentSession session, string? userId, CancellationToken cancellationToken = default) + => this.SaveSessionAsync( + agent, + ResolveAgentIdentity(agent), + conversationId, + session, + userId, + cancellationToken); + + internal async ValueTask SaveSessionAsync( + AIAgent agent, + string agentIdentity, + string conversationId, + AgentSession session, + string? userId, + CancellationToken cancellationToken = default) { _ = Throw.IfNull(agent); + _ = Throw.IfNullOrWhitespace(agentIdentity); _ = Throw.IfNullOrWhitespace(conversationId); _ = Throw.IfNull(session); JsonElement serialized = await agent.SerializeSessionAsync(session, cancellationToken: cancellationToken).ConfigureAwait(false); BinaryData sessionData = ToBinaryData(serialized); - string logicalKey = BuildLogicalKey(agent, conversationId, userId); + string logicalKey = BuildLogicalKey(agentIdentity, conversationId, userId); FoundryStateStore store = await this.GetStoreAsync(cancellationToken).ConfigureAwait(false); await store.SetItemAsync( @@ -147,16 +161,30 @@ await store.SetItemAsync( } /// - public override async ValueTask GetSessionAsync( + public override ValueTask GetSessionAsync( + AIAgent agent, + string conversationId, + string? userId, + CancellationToken cancellationToken = default) + => this.GetSessionAsync( + agent, + ResolveAgentIdentity(agent), + conversationId, + userId, + cancellationToken); + + internal async ValueTask GetSessionAsync( AIAgent agent, + string agentIdentity, string conversationId, string? userId, CancellationToken cancellationToken = default) { _ = Throw.IfNull(agent); + _ = Throw.IfNullOrWhitespace(agentIdentity); _ = Throw.IfNullOrWhitespace(conversationId); - string logicalKey = BuildLogicalKey(agent, conversationId, userId); + string logicalKey = BuildLogicalKey(agentIdentity, conversationId, userId); FoundryStateStore store = await this.GetStoreAsync(cancellationToken).ConfigureAwait(false); // GetItemAsync already answers null for an item that is not there, which is exactly the @@ -182,27 +210,31 @@ private ValueTask GetStoreAsync(CancellationToken cancellatio => this._binding.GetAsync(cancellationToken); /// - /// Builds the readable partition key. This is the same a-/u-/c- scheme - /// and use, so - /// the three stores partition sessions identically: per hosted agent, then per end user, then - /// per conversation. agent.Id is deliberately not used because it is regenerated on every - /// startup for in-memory-defined agents, which would break session continuity. Each segment is - /// omitted when its value is absent. + /// Builds an unambiguous readable partition key from the hosted agent identity, end user, and + /// conversation. Each component carries its length so delimiters inside values cannot collide. /// - internal static string BuildLogicalKey(AIAgent agent, string conversationId, string? userId) + internal static string BuildLogicalKey(string agentIdentity, string conversationId, string? userId) { StringBuilder builder = new(); - if (!string.IsNullOrEmpty(agent.Name)) - { - builder.Append("a-").Append(agent.Name).Append(':'); - } + AppendComponent(builder, 'a', Throw.IfNullOrWhitespace(agentIdentity)); + AppendComponent(builder, 'u', string.IsNullOrWhiteSpace(userId) ? null : userId); + AppendComponent(builder, 'c', Throw.IfNullOrWhitespace(conversationId)); + builder.Length--; + return builder.ToString(); + } + + private static string ResolveAgentIdentity(AIAgent agent) => + !string.IsNullOrWhiteSpace(agent.Name) ? $"name:{agent.Name}" : $"id:{agent.Id}"; - if (!string.IsNullOrWhiteSpace(userId)) + private static void AppendComponent(StringBuilder builder, char prefix, string? value) + { + builder.Append(prefix).Append(value?.Length ?? -1).Append(':'); + if (value is not null) { - builder.Append("u-").Append(userId).Append(':'); + builder.Append(value); } - return builder.Append("c-").Append(conversationId).ToString(); + builder.Append('|'); } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs index bfc00c4a40..7b676f344d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryJsonCheckpointStore.cs @@ -37,10 +37,9 @@ namespace Microsoft.Agents.AI.Foundry.Hosting; /// /// Retention. A workflow session writes one checkpoint per superstep but only ever resumes from /// the most recent one. Retrieving a checkpoint happens when a workflow is resuming from it, and at -/// that point every other checkpoint of that session is deleted. A conversation therefore holds one -/// turn's worth of checkpoints rather than growing for as long as it lasts to avoid storing -/// limitations. Note that this makes the point where old -/// checkpoints are collected, so retrieving one is not a read-only operation on this store. +/// that point checkpoints committed before it are deleted. Checkpoints committed later are retained +/// so a concurrent run cannot lose its live state. Note that this makes +/// a write operation as well as a read. /// /// /// Concurrency. Adding a checkpoint writes the checkpoint item and then updates the session's index @@ -222,14 +221,12 @@ await store.SetItemAsync( } /// - /// Returns a stored checkpoint and, in the same call, deletes every other checkpoint of that - /// session. + /// Returns a stored checkpoint and deletes checkpoints committed before it. /// /// /// - /// The method also deletes every other checkpoint of that session except the one that has just been retrieved, which is - /// the one the workflow is resuming from. The deletion is not incidental, it is how this store keeps a session's checkpoints from - /// piling up. + /// The deletion keeps old checkpoints from piling up. Checkpoints committed after the retrieved + /// checkpoint are retained because they may belong to a concurrent run. /// /// /// A workflow writes one checkpoint per superstep and only ever resumes from the most @@ -239,7 +236,7 @@ await store.SetItemAsync( /// /// /// The workflow session that owns the checkpoint. - /// Identifies the checkpoint to return, and the one checkpoint left in place. + /// Identifies the checkpoint to return and the oldest checkpoint retained. /// The stored checkpoint. /// No such checkpoint is stored for that session. public override async ValueTask RetrieveCheckpointAsync(string sessionId, CheckpointInfo key) @@ -279,8 +276,8 @@ private static JsonElement ParseCheckpoint(BinaryData checkpointData) } /// - /// Deletes every checkpoint of a session except the one that has just been retrieved, which is - /// the one the workflow is resuming from. + /// Deletes every checkpoint committed before the one that has just been retrieved, which is the + /// one the workflow is resuming from. /// /// /// @@ -297,29 +294,20 @@ private async Task PruneObsoleteCheckpointsAsync(FoundryStateStore store, string { StateStoreItem? indexItem = await store.GetItemAsync(sessionIndexKey, CancellationToken.None).ConfigureAwait(false); - List obsolete = []; - IndexEntry? resumed = null; - foreach (IndexEntry entry in ReadEntries(indexItem)) - { - if (entry.CheckpointId == resumedCheckpointId) - { - resumed = entry; - } - else - { - obsolete.Add(entry); - } - } - - if (resumed is null || obsolete.Count == 0) + List entries = ReadEntries(indexItem); + int resumedIndex = entries.FindIndex(entry => entry.CheckpointId == resumedCheckpointId); + if (resumedIndex <= 0) { return; } + List obsolete = entries.GetRange(0, resumedIndex); + List retained = entries.GetRange(resumedIndex, entries.Count - resumedIndex); + // The index is shortened first. A checkpoint item that is still listed but already gone // would be read as a missing checkpoint, whereas one that is listed nowhere is simply // never asked for. - await WriteEntriesAsync(store, sessionIndexKey, sessionId, [resumed], indexItem?.Etag).ConfigureAwait(false); + await WriteEntriesAsync(store, sessionIndexKey, sessionId, retained, indexItem?.Etag).ConfigureAwait(false); foreach (IndexEntry entry in obsolete) { @@ -343,13 +331,13 @@ private async Task PruneObsoleteCheckpointsAsync(FoundryStateStore store, string } catch (FoundryStorageException ex) when (IsLostRace(ex)) { - // Another instance updated the same session index first. Its own resume prunes whatever - // this one left behind, so nothing is leaked and there is nothing to report. + // Another instance updated the same session index first. Leaving the old items in place + // is safer than deleting against a stale index; a later resume can prune them. if (this._logger?.IsEnabled(LogLevel.Debug) is true) { this._logger.LogDebug( ex, - "Pruning obsolete checkpoints of session '{SessionId}' lost to another writer. The winning writer prunes them instead.", + "Pruning obsolete checkpoints of session '{SessionId}' lost to another writer. The old checkpoints remain until a later resume or expiry.", sessionId); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs index 2e538b9a7f..ff603975aa 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryStateStoreBinding.cs @@ -46,7 +46,12 @@ public async ValueTask GetAsync(CancellationToken cancellatio // WaitAsync applies the caller's token to this caller's wait only. return await binding.WaitAsync(cancellationToken).ConfigureAwait(false); } - catch (Exception ex) when (ex is not OperationCanceledException) + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested && !binding.IsCompleted) + { + // This caller stopped waiting, but the shared binding is still usable by everyone else. + throw; + } + catch { lock (this._gate) { diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs index b81d23f7f5..d992f61ae5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs @@ -64,11 +64,12 @@ public async Task CheckHealthAsync(HealthCheckContext context } List storingAgents = []; + List unprobeableAgents = []; var checkedAgents = 0; foreach (var agent in this.ResolveAgents()) { - if (agent.GetService() is null) + 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. @@ -76,20 +77,30 @@ public async Task CheckHealthAsync(HealthCheckContext context } checkedAgents++; + if (!ReferenceEquals(agent, chatClientAgent) && agent is not FoundryAgent) + { + unprobeableAgents.Add(agent.Name ?? agent.Id); + continue; + } + if (await this.StoresItsOwnResponsesAsync(agent, cancellationToken).ConfigureAwait(false)) { storingAgents.Add(agent.Name ?? agent.Id); } } - if (storingAgents.Count > 0) + if (storingAgents.Count > 0 || unprobeableAgents.Count > 0) { return new HealthCheckResult( status: context.Registration.FailureStatus, description: string.Create( CultureInfo.InvariantCulture, - $"Stored output: {storingAgents.Count} registered agent(s) should not have server side storage enabled. This will produce a new untracked conversation/response in the server while the hosted agent will also generate a conversation for the request of the agent. This setting is only allowed when enabling the FoundryResponsesOptions.AllowStoredOutputEnabled flag, which leaves the agent's own storage setting untouched and keeps that second recording on purpose."), - data: new Dictionary(StringComparer.Ordinal) { ["storingAgents"] = storingAgents }); + $"Stored output: {storingAgents.Count} registered agent(s) ask their service to store responses, and {unprobeableAgents.Count} wrapped agent(s) cannot be inspected without running their middleware. Disable server-side storage, register the ChatClientAgent directly, or explicitly allow stored output."), + data: new Dictionary(StringComparer.Ordinal) + { + ["storingAgents"] = storingAgents, + ["unprobeableAgents"] = unprobeableAgents, + }); } return HealthCheckResult.Healthy( @@ -114,11 +125,9 @@ public async Task CheckHealthAsync(HealthCheckContext context /// stored output setting, the chat options and the raw request factory among them, and drops both /// kinds of provider, so the probe stays free of side effects. /// - /// - /// What the stand-in cannot see is a decorator wrapped around the agent. One that changed this - /// setting would go unreported, which is accepted here: the check exists to catch how a container - /// configured its agent, and a silent probe is worth more than a complete one. - /// + /// A is safe to inspect because it transparently delegates runs to its + /// inner . Other wrappers are rejected before this method is called + /// because rebuilding only the leaf would miss middleware that changes the effective run options. /// private async Task StoresItsOwnResponsesAsync(AIAgent agent, CancellationToken cancellationToken) { diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs index c9c1c49dd5..9c57beb781 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs @@ -7,7 +7,9 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Workflows; +using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Diagnostics.HealthChecks; +using Microsoft.Extensions.Logging; using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; @@ -62,7 +64,7 @@ public Task CheckHealthAsync(HealthCheckContext context, Canc "Workflow checkpointing: not running in a Foundry container, so workflow checkpoints are left where each agent puts them.")); } - List agentsWithOwnCheckpointing = []; + List incompatibleAgents = []; var checkedAgents = 0; foreach (var agent in FoundryHostingExtensions.ResolveRegisteredAgents(this._serviceProvider)) @@ -73,20 +75,24 @@ public Task CheckHealthAsync(HealthCheckContext context, Canc } checkedAgents++; - if (metadata.UsesOwnCheckpointStorage) + AIAgent redirected = FoundryHostingExtensions.ApplyWorkflowCheckpointing( + agent, + this._serviceProvider.GetService()); + + if (metadata.UsesOwnCheckpointStorage || ReferenceEquals(redirected, agent)) { - agentsWithOwnCheckpointing.Add(agent.Name ?? agent.Id); + incompatibleAgents.Add(agent.Name ?? agent.Id); } } - if (agentsWithOwnCheckpointing.Count > 0) + if (incompatibleAgents.Count > 0) { return Task.FromResult(new HealthCheckResult( status: context.Registration.FailureStatus, description: string.Create( CultureInfo.InvariantCulture, - $"Workflow checkpointing: {agentsWithOwnCheckpointing.Count} registered workflow agent(s) were built with a checkpoint manager of their own. A hosted workflow has its checkpoints written to the Foundry state store so they survive the container being replaced and can be read by every instance; an agent that names its own manager keeps that state somewhere this container does not manage. Build the agent without passing an execution environment configured with WithCheckpointing, and let hosting supply the store."), - data: new Dictionary(StringComparer.Ordinal) { ["agentsWithOwnCheckpointing"] = agentsWithOwnCheckpointing })); + $"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 })); } return Task.FromResult(HealthCheckResult.Healthy( diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs index d6299ff209..9cdad4c75a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentSessionStoreTests.cs @@ -45,7 +45,7 @@ public async Task SaveSessionAsync_StoresReadableLogicalKeyAlongsideTheSessionAs // Assert: the item body keeps the readable key so a stored item can be traced back. var item = Assert.Single(backing.Items); - Assert.Equal("\"a-Concierge:u-alice:c-conv-1\"", item["key"].ToString()); + Assert.Equal("\"a14:name:Concierge|u5:alice|c6:conv-1\"", item["key"].ToString()); } [Fact] @@ -114,6 +114,34 @@ public async Task GetSessionAsync_DifferentAgent_DoesNotReadAnotherAgentsSession Assert.Equal(0, researcher.DeserializeCalls); } + [Fact] + public async Task GetSessionAsync_DifferentKeyedRegistration_DoesNotReadAnotherAgentsSessionAsync() + { + // Arrange: both agents are unnamed, so their keyed DI registrations are the only stable + // identities that can separate their sessions. + var backing = new FakeStateStore(); + var store = NewStore(backing); + var billing = new TestAgent("{\"owner\":\"billing\"}"); + var support = new TestAgent(); + await store.SaveSessionAsync( + billing, + "key:billing", + "shared-conv", + new TestSession(), + userId: "alice"); + + // Act + var supportSession = await store.GetSessionAsync( + support, + "key:support", + "shared-conv", + userId: "alice"); + + // Assert + Assert.Null(supportSession); + Assert.Equal(0, support.DeserializeCalls); + } + [Fact] public async Task GetStoreAsync_ResolvesTheStoreOnceAcrossManyCallsAsync() { @@ -161,30 +189,121 @@ await Assert.ThrowsAsync( Assert.Equal(2, attempts); } - [Theory] - [InlineData("Concierge", "alice", "conv-1", "a-Concierge:u-alice:c-conv-1")] - [InlineData("Concierge", null, "conv-1", "a-Concierge:c-conv-1")] - [InlineData(null, "alice", "conv-1", "u-alice:c-conv-1")] - [InlineData(null, null, "conv-1", "c-conv-1")] - [InlineData("x", "x", "conv-1", "a-x:u-x:c-conv-1")] - public void BuildLogicalKey_UsesTheSamePrefixSchemeAsTheOtherStores(string? agentName, string? userId, string conversationId, string expected) + [Fact] + public async Task GetStoreAsync_CanceledBinding_IsRetriedOnTheNextCallAsync() + { + // Arrange: the shared binding task itself was canceled, rather than one caller choosing to + // stop waiting for an otherwise healthy shared task. + var backing = new FakeStateStore(); + var attempts = 0; + var store = new FoundryAgentSessionStore(_ => + { + attempts++; + return attempts == 1 + ? Task.FromCanceled(new CancellationToken(canceled: true)) + : Task.FromResult(backing); + }); + var agent = new TestAgent(name: "Concierge"); + + // Act + await Assert.ThrowsAnyAsync( + async () => await store.GetSessionAsync(agent, "conv-1", userId: null)); + var session = await store.GetSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.Null(session); + Assert.Equal(2, attempts); + } + + [Fact] + public async Task GetStoreAsync_BindingFaultedWithCancellation_IsRetriedOnTheNextCallAsync() + { + // Arrange: some async APIs fault with OperationCanceledException instead of returning a + // canceled task. That completed shared failure must not remain cached. + var backing = new FakeStateStore(); + var attempts = 0; + var store = new FoundryAgentSessionStore(_ => + { + attempts++; + return attempts == 1 + ? Task.FromException(new OperationCanceledException()) + : Task.FromResult(backing); + }); + var agent = new TestAgent(name: "Concierge"); + + // Act + await Assert.ThrowsAnyAsync( + async () => await store.GetSessionAsync(agent, "conv-1", userId: null)); + var session = await store.GetSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.Null(session); + Assert.Equal(2, attempts); + } + + [Fact] + public async Task GetStoreAsync_CallerCancellation_DoesNotDiscardTheSharedBindingAsync() { // Arrange - var agent = new TestAgent(name: agentName); + var backing = new FakeStateStore(); + var binding = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var attempts = 0; + var store = new FoundryAgentSessionStore(_ => + { + attempts++; + return binding.Task; + }); + var agent = new TestAgent(name: "Concierge"); + using var cancellation = new CancellationTokenSource(); + cancellation.Cancel(); // Act - var key = FoundryAgentSessionStore.BuildLogicalKey(agent, conversationId, userId); + await Assert.ThrowsAnyAsync( + async () => await store.GetSessionAsync(agent, "conv-1", userId: null, cancellation.Token)); + binding.SetResult(backing); + var session = await store.GetSessionAsync(agent, "conv-1", userId: null); + + // Assert + Assert.Null(session); + Assert.Equal(1, attempts); + } + + [Theory] + [InlineData("name:Concierge", "alice", "conv-1", "a14:name:Concierge|u5:alice|c6:conv-1")] + [InlineData("name:Concierge", null, "conv-1", "a14:name:Concierge|u-1:|c6:conv-1")] + [InlineData("default", "alice", "conv-1", "a7:default|u5:alice|c6:conv-1")] + [InlineData("default", null, "conv-1", "a7:default|u-1:|c6:conv-1")] + [InlineData("name:x", "x", "conv-1", "a6:name:x|u1:x|c6:conv-1")] + public void BuildLogicalKey_UsesLengthPrefixedComponents(string agentIdentity, string? userId, string conversationId, string expected) + { + // Act + var key = FoundryAgentSessionStore.BuildLogicalKey(agentIdentity, conversationId, userId); // Assert Assert.Equal(expected, key); } + [Fact] + public void BuildLogicalKey_DelimitersInsideComponents_DoNotCollide() + { + // Act: these tuples produced the same delimiter-joined string before components carried + // their lengths. + string first = FoundryAgentSessionStore.BuildLogicalKey("name:Concierge", "x:c-y", "alice"); + string second = FoundryAgentSessionStore.BuildLogicalKey("name:Concierge", "y", "alice:c-x"); + + // Assert + Assert.NotEqual(first, second); + Assert.NotEqual( + FoundryAgentSessionStore.BuildItemKey(first), + FoundryAgentSessionStore.BuildItemKey(second)); + } + [Fact] public void BuildItemKey_StaysWithinThePlatformKeyLimitForAnyInput() { // Arrange: an agent name plus a user id plus a conversation id can easily pass 128 chars. var logicalKey = FoundryAgentSessionStore.BuildLogicalKey( - new TestAgent(name: new string('a', 200)), + $"name:{new string('a', 200)}", new string('c', 200), new string('u', 200)); @@ -199,9 +318,9 @@ public void BuildItemKey_StaysWithinThePlatformKeyLimitForAnyInput() public void BuildItemKey_IsStableAndDistinctPerLogicalKey() { // Arrange / Act - var first = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-alice:c-conv-1"); - var same = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-alice:c-conv-1"); - var other = FoundryAgentSessionStore.BuildItemKey("a-Concierge:u-bob:c-conv-1"); + var first = FoundryAgentSessionStore.BuildItemKey("a14:name:Concierge|u5:alice|c6:conv-1"); + var same = FoundryAgentSessionStore.BuildItemKey("a14:name:Concierge|u5:alice|c6:conv-1"); + var other = FoundryAgentSessionStore.BuildItemKey("a14:name:Concierge|u3:bob|c6:conv-1"); // Assert Assert.Equal(first, same); diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs index 1e7621889e..3dfb91fef0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryJsonCheckpointStoreTests.cs @@ -194,7 +194,7 @@ public void BuildCheckpointKey_StaysWithinThePlatformKeyLimit() } [Fact] - public async Task RetrieveCheckpointAsync_DeletesEveryOtherCheckpointOfTheSessionAsync() + public async Task RetrieveCheckpointAsync_DeletesPredecessorsOfTheResumeTargetAsync() { // Arrange: a session that ran three supersteps, so it holds three checkpoints. var backing = new FakeCheckpointStateStore(); @@ -214,6 +214,27 @@ public async Task RetrieveCheckpointAsync_DeletesEveryOtherCheckpointOfTheSessio Assert.True(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", third.CheckpointId))); } + [Fact] + public async Task RetrieveCheckpointAsync_RetainsCheckpointsCommittedAfterTheResumeTargetAsync() + { + // Arrange: the third checkpoint models another request committing after this request chose + // the second checkpoint as its resume target. + var backing = new FakeCheckpointStateStore(); + var store = NewStore(backing); + var first = await store.CreateCheckpointAsync("session-1", Json("{\"step\":1}")); + var second = await store.CreateCheckpointAsync("session-1", Json("{\"step\":2}")); + var third = await store.CreateCheckpointAsync("session-1", Json("{\"step\":3}")); + + // Act + await store.RetrieveCheckpointAsync("session-1", second); + + // Assert: only predecessors are obsolete. The concurrent request's later checkpoint remains. + Assert.Equal([second, third], (await store.RetrieveIndexAsync("session-1")).ToList()); + Assert.False(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", first.CheckpointId))); + Assert.True(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", second.CheckpointId))); + Assert.True(backing.Items.ContainsKey(FoundryJsonCheckpointStore.BuildCheckpointKey("session-1", third.CheckpointId))); + } + [Fact] public async Task RetrieveCheckpointAsync_LeavesOtherSessionsAloneAsync() { 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 a6f2eb0349..ecad36248d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs @@ -155,6 +155,49 @@ public async Task CheckHealthAsync_AgentWithProviders_LeavesThemUntouchedAsync() Assert.False(historyProvider.WasInvoked); } + [Fact] + public async Task CheckHealthAsync_WrappedChatClientAgent_IsUnhealthyWithoutRunningItAsync() + { + // Arrange: rebuilding only the leaf would miss any option changes made by this wrapper. + var inner = new ChatClientAgent( + NewSilentChatClient(), + new ChatClientAgentOptions { Name = "wrapped" }); + AIAgent wrapper = new PassThroughAgent(inner); + var check = BuildCheckFor(wrapper); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.Equal(["wrapped"], Assert.IsType>(result.Data["unprobeableAgents"])); + } + + [Fact] + public async Task CheckHealthAsync_FoundryAgent_IsProbeableAsync() + { + // Arrange: FoundryAgent is a transparent wrapper around its ChatClientAgent. + var inner = new ChatClientAgent( + NewSilentChatClient(), + new ChatClientAgentOptions + { + Name = "foundry-agent", + ChatOptions = new ChatOptions + { + RawRepresentationFactory = _ => new CreateResponseOptions { StoredOutputEnabled = true }, + }, + }); + var check = BuildCheckFor(new FoundryAgent(inner)); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.Contains("foundry-agent", (List)result.Data["storingAgents"]); + Assert.Empty((List)result.Data["unprobeableAgents"]); + } + private static HostedStoredOutputHealthCheck BuildCheckFor(AIAgent agent, FoundryResponsesOptions? hostingOptions = null) { var services = new ServiceCollection(); @@ -227,4 +270,6 @@ protected override ValueTask InvokedCoreAsync(InvokedContext context, Cancellati return default; } } + + private sealed class PassThroughAgent(AIAgent innerAgent) : DelegatingAIAgent(innerAgent); } 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 fc33c4af41..f52e71729b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs @@ -44,7 +44,7 @@ public async Task CheckHealthAsync_WorkflowAgentWithItsOwnCheckpointManager_IsUn // Assert Assert.Equal(HealthStatus.Unhealthy, result.Status); - var reported = Assert.IsType>(result.Data["agentsWithOwnCheckpointing"]); + var reported = Assert.IsType>(result.Data["incompatibleAgents"]); Assert.Equal(["WorkflowAgent"], reported); } @@ -62,6 +62,21 @@ public async Task CheckHealthAsync_WorkflowAgentBehindAWrapper_IsStillReportedAs Assert.Equal(HealthStatus.Unhealthy, result.Status); } + [Fact] + public async Task CheckHealthAsync_WrappedWorkflowWithoutOwnManager_IsUnhealthyAsync() + { + // Arrange: the metadata flows through the wrapper, but hosting cannot replace the inner + // workflow agent without discarding that wrapper. + var check = BuildCheckFor(new PassThroughAgent(BuildWorkflowAgent(executionEnvironment: null))); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.Equal(["WorkflowAgent"], Assert.IsType>(result.Data["incompatibleAgents"])); + } + [Fact] public async Task CheckHealthAsync_AgentThatDoesNotRunAWorkflow_IsPassedOverAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs index 2655b49c73..2f3422c282 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostingExtensionsTests.cs @@ -7,7 +7,7 @@ namespace Microsoft.Agents.AI.Workflows.UnitTests; /// /// Covers , which lets a host redirect -/// where a already-built workflow agent writes its checkpoints. +/// where an already-built workflow agent writes its checkpoints. /// public class WorkflowHostingExtensionsTests {