Skip to content

Commit 7638874

Browse files
quifoxquifox
authored
.NET: Dispose AI context streaming enumerators (microsoft#8812)
Co-authored-by: quifox <289420841+quifox@users.noreply.github.com>
1 parent 7a98aa9 commit 7638874

4 files changed

Lines changed: 82 additions & 0 deletions

File tree

‎dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
9494
throw;
9595
}
9696

97+
await using var _ = enumerator.ConfigureAwait(false);
98+
9799
bool hasUpdates;
98100
try
99101
{

‎dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
7979
throw;
8080
}
8181

82+
await using var _ = enumerator.ConfigureAwait(false);
83+
8284
bool hasUpdates;
8385
try
8486
{

‎dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,50 @@ public async Task GetStreamingResponseAsync_OnSuccess_InvokedAsyncCalledAsync()
229229
Assert.Null(provider.LastInvokedContext!.InvokeException);
230230
}
231231

232+
[Fact]
233+
public async Task GetStreamingResponseAsync_WhenConsumerStopsEarly_DisposesInnerEnumeratorAsync()
234+
{
235+
// Arrange
236+
bool disposed = false;
237+
238+
async IAsyncEnumerable<ChatResponseUpdate> StreamAsync()
239+
{
240+
try
241+
{
242+
yield return new ChatResponseUpdate(ChatRole.Assistant, "Part1");
243+
yield return new ChatResponseUpdate(ChatRole.Assistant, "Part2");
244+
}
245+
finally
246+
{
247+
disposed = true;
248+
}
249+
}
250+
251+
var innerClient = CreateMockStreamingChatClient(
252+
onGetStreamingResponse: (_, _, _) => StreamAsync());
253+
var provider = new TestAIContextProvider("key1");
254+
var chatClient = new AIContextProviderChatClient(innerClient, [provider]);
255+
256+
var agent = new TestAIAgent
257+
{
258+
RunAsyncFunc = async (messages, session, options, ct) =>
259+
{
260+
await foreach (var _ in chatClient.GetStreamingResponseAsync(messages, cancellationToken: ct))
261+
{
262+
break;
263+
}
264+
265+
return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]);
266+
}
267+
};
268+
269+
// Act
270+
await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession);
271+
272+
// Assert
273+
Assert.True(disposed);
274+
}
275+
232276
[Fact]
233277
public async Task GetStreamingResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
234278
{

‎dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,40 @@ public async Task RunStreamingAsync_OnSuccess_InvokedContextContainsAccumulatedR
281281
Assert.True(responseMessages.Count > 0);
282282
}
283283

284+
[Fact]
285+
public async Task RunStreamingAsync_WhenConsumerStopsEarly_DisposesInnerEnumeratorAsync()
286+
{
287+
// Arrange
288+
bool disposed = false;
289+
290+
async IAsyncEnumerable<AgentResponseUpdate> StreamAsync()
291+
{
292+
try
293+
{
294+
yield return new AgentResponseUpdate(ChatRole.Assistant, "Part1");
295+
yield return new AgentResponseUpdate(ChatRole.Assistant, "Part2");
296+
}
297+
finally
298+
{
299+
disposed = true;
300+
}
301+
}
302+
303+
var provider = new TestProvider();
304+
var innerAgent = CreateTestAgent(
305+
runStreamingFunc: (_, _, _, _) => StreamAsync());
306+
var agent = new MessageAIContextProviderAgent(innerAgent, [provider]);
307+
308+
// Act
309+
await foreach (var _ in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession))
310+
{
311+
break;
312+
}
313+
314+
// Assert
315+
Assert.True(disposed);
316+
}
317+
284318
[Fact]
285319
public async Task RunStreamingAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync()
286320
{

0 commit comments

Comments
 (0)