diff --git a/dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/ResilientWorkflowAgent.cs b/dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/ResilientWorkflowAgent.cs index 949b2ed8f6..fea345f179 100644 --- a/dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/ResilientWorkflowAgent.cs +++ b/dotnet/tests/Foundry.Hosting.IntegrationTests.TestContainer/ResilientWorkflowAgent.cs @@ -92,6 +92,9 @@ public override async ValueTask HandleAsync( { if (TryCreateCrashMarker(token, out string crashedProcessIncarnation)) { + await Task.Delay( + TimeSpan.FromSeconds(GetCrashDelaySeconds()), + cancellationToken).ConfigureAwait(false); Console.Out.Flush(); Console.Error.Flush(); Environment.Exit(70); @@ -115,6 +118,21 @@ private static int GetLongRunningDelaySeconds() string? value = Environment.GetEnvironmentVariable("IT_LONG_RUNNING_DELAY_SECONDS"); return int.TryParse(value, out int seconds) && seconds > 0 ? seconds : DefaultDelaySeconds; } + + private static int GetCrashDelaySeconds() + { + const int DefaultDelaySeconds = 5; + string? value = Environment.GetEnvironmentVariable( + "IT_CRASH_DELAY_SECONDS"); + return int.TryParse( + value, + NumberStyles.None, + CultureInfo.InvariantCulture, + out int seconds) + && seconds >= 0 + ? seconds + : DefaultDelaySeconds; + } } [SendsMessage(typeof(string))] diff --git a/dotnet/tests/Foundry.Hosting.IntegrationTests/Fixtures/ResilientWorkflowHostedAgentFixture.cs b/dotnet/tests/Foundry.Hosting.IntegrationTests/Fixtures/ResilientWorkflowHostedAgentFixture.cs index 6281f8f23b..bacd65ffc6 100644 --- a/dotnet/tests/Foundry.Hosting.IntegrationTests/Fixtures/ResilientWorkflowHostedAgentFixture.cs +++ b/dotnet/tests/Foundry.Hosting.IntegrationTests/Fixtures/ResilientWorkflowHostedAgentFixture.cs @@ -18,6 +18,7 @@ public sealed class ResilientWorkflowHostedAgentFixture : HostedAgentFixture protected override void ConfigureEnvironment(IDictionary environment) { environment["IT_LONG_RUNNING_DELAY_SECONDS"] = "20"; + environment["IT_CRASH_DELAY_SECONDS"] = "5"; environment["IT_COUNTDOWN_DELAY_MILLISECONDS"] = "250"; environment["IT_COUNTDOWN_CRASH_DELAY_SECONDS"] = "5"; } diff --git a/dotnet/tests/Foundry.Hosting.IntegrationTests/ResilientWorkflowHostedAgentTests.cs b/dotnet/tests/Foundry.Hosting.IntegrationTests/ResilientWorkflowHostedAgentTests.cs index 570167aea6..85a4abbeba 100644 --- a/dotnet/tests/Foundry.Hosting.IntegrationTests/ResilientWorkflowHostedAgentTests.cs +++ b/dotnet/tests/Foundry.Hosting.IntegrationTests/ResilientWorkflowHostedAgentTests.cs @@ -67,7 +67,7 @@ public async Task BackgroundResponse_ProcessCrash_RecoversAndCompletesAsync() ResponseWaitResult waitResult = await WaitForTerminalAsync(responses, accepted.Id, s_completionTimeout); // Assert: this token is emitted only after a new process observes the crash marker written - // immediately before Environment.Exit. + // before Environment.Exit. Assert.True(accepted.Status is ResponseStatus.Queued or ResponseStatus.InProgress); Assert.True( waitResult.SawSessionNotReady @@ -260,7 +260,7 @@ private static ResponseContinuationToken CreateReplayFromStartToken( } private static bool IsTransientRecoveryStatus(int status) => - status is 404 or 424 or 500 or 502 or 503; + status is 404 or 409 or 424 or 500 or 502 or 503; private static int CountCountdownUpdates(IEnumerable texts) => texts.Count(text => text != "Countdown complete."); @@ -296,6 +296,13 @@ private static async Task WaitForTerminalAsync( await Task.Delay(TimeSpan.FromSeconds(2)); continue; } + catch (ClientResultException ex) when (ex.Status == 409) + { + longestPollDuration = Max(longestPollDuration, pollStopwatch.Elapsed); + sawSessionNotReady = true; + await Task.Delay(TimeSpan.FromSeconds(2)); + continue; + } longestPollDuration = Max(longestPollDuration, pollStopwatch.Elapsed); if (response.Status is ResponseStatus.Completed) diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs index 8d72a2e4c3..6c1bf39e80 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs @@ -11,6 +11,7 @@ using System.Text.Json; using System.Threading; using System.Threading.Tasks; +using Azure.AI.AgentServer.Responses; using Microsoft.Agents.AI.Workflows; using Microsoft.Agents.AI.Workflows.Checkpointing; using Microsoft.AspNetCore.Builder; @@ -18,6 +19,9 @@ using Microsoft.AspNetCore.TestHost; using Microsoft.Extensions.AI; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.DependencyInjection.Extensions; +using CreateResponse = Azure.AI.AgentServer.Responses.Models.CreateResponse; +using ResponseStreamEvent = Azure.AI.AgentServer.Responses.Models.ResponseStreamEvent; namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; @@ -148,7 +152,8 @@ public async Task StoppedHost_RecoversWorkflowWithCompleteOrderedOutputAsync() { WebApplication firstHost = await StartServerAsync( BuildCountdownWorkflowAgent(coordinator, checkpointStore), - new FoundryAgentSessionStore(storeName: sessionStoreName)); + new FoundryAgentSessionStore(storeName: sessionStoreName), + coordinator); try { using HttpClient firstClient = GetClient(firstHost); @@ -158,12 +163,8 @@ public async Task StoppedHost_RecoversWorkflowWithCompleteOrderedOutputAsync() agentName: "countdown-workflow", input: "Count down from 6"); await coordinator.Blocked.Task.WaitAsync(TimeSpan.FromSeconds(15)); - await WaitForResponseProgressAsync( - firstClient, - responseId, - ["6", "5", "4"], - minimumOutputItems: 12, - timeout: TimeSpan.FromSeconds(15)); + await coordinator.ExpectedCheckpointProcessed.Task.WaitAsync( + TimeSpan.FromSeconds(15)); using CancellationTokenSource stopTimeout = new(TimeSpan.FromSeconds(15)); @@ -224,7 +225,8 @@ await WaitForResponseProgressAsync( private static async Task StartServerAsync( AIAgent agent, - AgentSessionStore sessionStore) + AgentSessionStore sessionStore, + CountdownRecoveryCoordinator? recoveryCoordinator = null) { WebApplicationBuilder builder = WebApplication.CreateBuilder(); builder.WebHost.UseTestServer(); @@ -235,6 +237,15 @@ private static async Task StartServerAsync( builder.Services.AddSingleton( new FakeHostedSessionIsolationKeyProvider()); builder.Services.AddLogging(); + if (recoveryCoordinator is not null) + { + builder.Services.RemoveAll(); + builder.Services.AddSingleton(serviceProvider => + new CheckpointObservingResponseHandler( + ActivatorUtilities.CreateInstance( + serviceProvider), + recoveryCoordinator)); + } WebApplication app = builder.Build(); app.MapFoundryResponses(); @@ -310,53 +321,23 @@ private static async Task WaitForTerminalAsync( $"Response '{responseId}' did not complete. Last response: {last}"); } - private static async Task WaitForResponseProgressAsync( - HttpClient client, - string responseId, - IReadOnlyList expected, - int minimumOutputItems, - TimeSpan timeout) - { - var deadline = DateTimeOffset.UtcNow + timeout; - List last = []; - while (DateTimeOffset.UtcNow < deadline) - { - using HttpResponseMessage response = await client.GetAsync( - new Uri($"/responses/{responseId}", UriKind.Relative)); - if (response.StatusCode == HttpStatusCode.OK) - { - using JsonDocument document = JsonDocument.Parse( - await response.Content.ReadAsStringAsync()); - JsonElement root = document.RootElement; - last = GetOutputTexts(root); - if (last.Count == expected.Count - && last.SequenceEqual(expected) - && root.GetProperty("output").GetArrayLength() >= minimumOutputItems) - { - return; - } - } - - await Task.Delay(TimeSpan.FromMilliseconds(25)); - } - - throw new TimeoutException( - $"Response '{responseId}' did not reach the expected checkpointed output. " + - $"Expected: {string.Join(", ", expected)}. Last: {string.Join(", ", last)}."); - } - private static JsonElement ReadPersistedResponse( string stateRoot, string responseId) { - string path = Path.Combine( + string path = GetPersistedResponsePath(stateRoot, responseId); + using JsonDocument document = JsonDocument.Parse(File.ReadAllBytes(path)); + return document.RootElement.GetProperty("envelope").Clone(); + } + + private static string GetPersistedResponsePath( + string stateRoot, + string responseId) => + Path.Combine( stateRoot, "responses", "envelopes", $"{responseId}.json"); - using JsonDocument document = JsonDocument.Parse(File.ReadAllBytes(path)); - return document.RootElement.GetProperty("envelope").Clone(); - } private static string GetOutputText(JsonElement response) { @@ -637,13 +618,53 @@ private sealed class RecoveryCoordinator private sealed class CountdownRecoveryCoordinator(int target, int blockAt) { private int _blocked; + private int _processedCheckpoints; public int Target { get; } = target; + public int BlockAt { get; } = blockAt; + public TaskCompletionSource Blocked { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + public TaskCompletionSource ExpectedCheckpointProcessed { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + public bool ShouldBlock(int value) => - value == blockAt && Interlocked.CompareExchange(ref this._blocked, 1, 0) == 0; + value == this.BlockAt && Interlocked.CompareExchange(ref this._blocked, 1, 0) == 0; + + public void OnCheckpointProcessed() + { + int expectedCheckpointCount = this.Target - this.BlockAt + 1; + if (Interlocked.Increment(ref this._processedCheckpoints) == expectedCheckpointCount) + { + this.ExpectedCheckpointProcessed.TrySetResult(); + } + } + } + + private sealed class CheckpointObservingResponseHandler( + ResponseHandler inner, + CountdownRecoveryCoordinator coordinator) : ResponseHandler + { + public override async IAsyncEnumerable CreateAsync( + CreateResponse request, + ResponseContext context, + [EnumeratorCancellation] CancellationToken cancellationToken) + { + await foreach (ResponseStreamEvent responseEvent in inner + .CreateAsync(request, context, cancellationToken) + .WithCancellation(cancellationToken) + .ConfigureAwait(false)) + { + bool isCheckpoint = + responseEvent.GetType().Name == "ResponseCheckpointEvent"; + yield return responseEvent; + if (isCheckpoint) + { + coordinator.OnCheckpointProcessed(); + } + } + } } }