diff --git a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs index bf93832232e..7956b7a0daf 100644 --- a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs @@ -94,6 +94,8 @@ public override async IAsyncEnumerable GetStreamingResponseA throw; } + await using var _ = enumerator.ConfigureAwait(false); + bool hasUpdates; try { diff --git a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs index 6453209edda..21b5ebbba83 100644 --- a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs @@ -79,6 +79,8 @@ protected override async IAsyncEnumerable RunCoreStreamingA throw; } + await using var _ = enumerator.ConfigureAwait(false); + bool hasUpdates; try { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs index 5e65c4a1a6c..2bb29c8e4f0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs @@ -229,6 +229,50 @@ public async Task GetStreamingResponseAsync_OnSuccess_InvokedAsyncCalledAsync() Assert.Null(provider.LastInvokedContext!.InvokeException); } + [Fact] + public async Task GetStreamingResponseAsync_WhenConsumerStopsEarly_DisposesInnerEnumeratorAsync() + { + // Arrange + bool disposed = false; + + async IAsyncEnumerable StreamAsync() + { + try + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "Part1"); + yield return new ChatResponseUpdate(ChatRole.Assistant, "Part2"); + } + finally + { + disposed = true; + } + } + + var innerClient = CreateMockStreamingChatClient( + onGetStreamingResponse: (_, _, _) => StreamAsync()); + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + var agent = new TestAIAgent + { + RunAsyncFunc = async (messages, session, options, ct) => + { + await foreach (var _ in chatClient.GetStreamingResponseAsync(messages, cancellationToken: ct)) + { + break; + } + + return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]); + } + }; + + // Act + await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession); + + // Assert + Assert.True(disposed); + } + [Fact] public async Task GetStreamingResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs index 3ba170a8d5a..a2a111b4af5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs @@ -281,6 +281,40 @@ public async Task RunStreamingAsync_OnSuccess_InvokedContextContainsAccumulatedR Assert.True(responseMessages.Count > 0); } + [Fact] + public async Task RunStreamingAsync_WhenConsumerStopsEarly_DisposesInnerEnumeratorAsync() + { + // Arrange + bool disposed = false; + + async IAsyncEnumerable StreamAsync() + { + try + { + yield return new AgentResponseUpdate(ChatRole.Assistant, "Part1"); + yield return new AgentResponseUpdate(ChatRole.Assistant, "Part2"); + } + finally + { + disposed = true; + } + } + + var provider = new TestProvider(); + var innerAgent = CreateTestAgent( + runStreamingFunc: (_, _, _, _) => StreamAsync()); + var agent = new MessageAIContextProviderAgent(innerAgent, [provider]); + + // Act + await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession)) + { + break; + } + + // Assert + Assert.True(disposed); + } + [Fact] public async Task RunStreamingAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync() {