diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index 725fd28f03b..32248cf30b1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -536,6 +536,11 @@ bool CheckNotAllowedStoreUsage() => } evt = enumerator.Current; + if (evt is ResponseCompletedEvent) + { + consentCts.Token.ThrowIfCancellationRequested(); + } + shutdownDetected = context.IsShutdownRequested && !emittedTerminal; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs index 070a6ca0f7b..2b7ae5b0061 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs @@ -666,15 +666,26 @@ public async Task CreateAsync_CancellationDuringExecution_PropagatesOperationCan .ReturnsAsync(Array.Empty()); using var cts = new CancellationTokenSource(); - cts.Cancel(); + var events = new List(); - // Act & Assert - await Assert.ThrowsAsync(async () => + async Task ExecuteAsync() { - await foreach (var _ in handler.CreateAsync(request, mockContext.Object, cts.Token)) + await foreach (var evt in handler.CreateAsync(request, mockContext.Object, cts.Token)) { + events.Add(evt); } - }); + } + + // Act + var execution = ExecuteAsync(); + await agent.Started.Task.WaitAsync(TimeSpan.FromSeconds(5)); + cts.Cancel(); + + // Assert + await Assert.ThrowsAnyAsync( + async () => await execution.WaitAsync(TimeSpan.FromSeconds(5))); + await agent.CancellationObserved.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.DoesNotContain(events, evt => evt is ResponseCompletedEvent); } [Fact] @@ -1643,15 +1654,29 @@ protected override ValueTask DeserializeSessionCoreAsync( private sealed class CancellationCheckingAgent : AIAgent { + public TaskCompletionSource Started { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource CancellationObserved { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + protected override async IAsyncEnumerable RunCoreStreamingAsync( IEnumerable messages, AgentSession? session, AgentRunOptions? options, [EnumeratorCancellation] CancellationToken cancellationToken = default) { - cancellationToken.ThrowIfCancellationRequested(); - yield return new AgentResponseUpdate { Contents = [new MeaiTextContent("test")] }; - await Task.CompletedTask; + this.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + this.CancellationObserved.TrySetResult(); + } + + yield break; } protected override Task RunCoreAsync( diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs index f5194f5f6d1..e3d255d700f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs @@ -161,6 +161,40 @@ public async Task WorkflowAgent_RegisteredWithKey_ResolvesCorrectlyAsync() Assert.True(events.Count >= 3, $"Expected at least 3 events, got {events.Count}"); } + [Fact] + public async Task CreateAsync_WorkflowCancellationDuringExecution_CancelsInnerAgentWithoutCompletionAsync() + { + // Arrange + var innerAgent = new CancellationCheckingWorkflowAgent("blocking-agent"); + var workflow = AgentWorkflowBuilder.BuildSequential("cancellation-workflow", innerAgent); + var workflowAgent = workflow.AsAIAgent( + id: "workflow-agent", + name: "Cancellation Workflow", + executionEnvironment: InProcessExecution.OffThread); + var (handler, request, context) = CreateHandlerWithAgent(workflowAgent, "Hello"); + using var cts = new CancellationTokenSource(); + var events = new List(); + + async Task ExecuteAsync() + { + await foreach (var evt in handler.CreateAsync(request, context, cts.Token)) + { + events.Add(evt); + } + } + + // Act + var execution = ExecuteAsync(); + await innerAgent.Started.Task.WaitAsync(TimeSpan.FromSeconds(5)); + cts.Cancel(); + + // Assert + await Assert.ThrowsAnyAsync( + async () => await execution.WaitAsync(TimeSpan.FromSeconds(5))); + await innerAgent.CancellationObserved.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.DoesNotContain(events, evt => evt is ResponseCompletedEvent); + } + private static (AgentFrameworkResponseHandler handler, CreateResponse request, ResponseContext context) CreateHandlerWithAgent( AIAgent agent, diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs index 7e10abc9ebd..0ca010d34fe 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs @@ -94,3 +94,62 @@ protected override ValueTask DeserializeSessionCoreAsync( CancellationToken cancellationToken = default) => throw new NotImplementedException(); } + +/// +/// A test agent that blocks until its run is cancelled. +/// +internal sealed class CancellationCheckingWorkflowAgent(string id) : AIAgent +{ + public new string Id => id; + + public TaskCompletionSource Started { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource CancellationObserved { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + this.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + this.CancellationObserved.TrySetResult(); + throw; + } + + yield break; + } + + protected override Task RunCoreAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + CancellationToken cancellationToken = default) => + throw new NotImplementedException(); + + protected override ValueTask CreateSessionCoreAsync( + CancellationToken cancellationToken = default) => + new(new TestAgentSession()); + + protected override ValueTask SerializeSessionCoreAsync( + AgentSession session, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(JsonSerializer.SerializeToElement(new { })); + + protected override ValueTask DeserializeSessionCoreAsync( + JsonElement serializedState, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(new TestAgentSession()); + + private sealed class TestAgentSession : AgentSession; +}