diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs index 8db6666c37..1c735539a4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs @@ -105,7 +105,7 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider State state = this._sessionState.GetOrInitializeState(context.Session); // Add request and response messages to the provider - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + var allNewMessages = (context.RequestMessages ?? []).Concat(context.ResponseMessages ?? []); state.Messages.AddRange(allNewMessages); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 44a136da3e..1133e10a8a 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -329,40 +329,18 @@ public sealed partial class ChatClientAgent : AIAgent this._logger.LogAgentChatClientInvokedStreamingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType); - bool hasUpdates; + // Ensure the inner enumerator is always disposed, even if the consumer breaks out early + // (e.g. ToolApprovalAgent does `yield break` after emitting an approval request). Without + // this, downstream decorators like PerServiceCallChatHistoryPersistingChatClient would be + // left suspended at `yield return`, never running their finally blocks, and any in-flight + // FunctionResultContent / FunctionCallContent state would not be persisted before the next + // turn, leaving the next request to the model with dangling tool calls. try { - // Ensure we start the streaming request - hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); - } - catch (Exception ex) - { - await this.NotifyProvidersOfFailureAtEndOfRunAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); - throw; - } - - while (hasUpdates) - { - var update = responseUpdatesEnumerator.Current; - if (update is not null) - { - update.AuthorName ??= this.Name; - - responseUpdates.Add(update); - - yield return new(update) - { - AgentId = this.Id, - ContinuationToken = WrapContinuationToken(update.ContinuationToken, GetInputMessages(inputMessages, continuationToken), responseUpdates) - }; - } - + bool hasUpdates; try { - // Re-ensure the run context has the resolved session before each MoveNextAsync. - // The base class RunStreamingAsync restores the original context (potentially with - // null session) after each yield, so we must re-establish it for the decorator. - EnsureRunContextHasSession(safeSession); + // Ensure we start the streaming request hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) @@ -370,20 +348,55 @@ public sealed partial class ChatClientAgent : AIAgent await this.NotifyProvidersOfFailureAtEndOfRunAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); throw; } + + while (hasUpdates) + { + var update = responseUpdatesEnumerator.Current; + if (update is not null) + { + update.AuthorName ??= this.Name; + + responseUpdates.Add(update); + + yield return new(update) + { + AgentId = this.Id, + ContinuationToken = WrapContinuationToken(update.ContinuationToken, GetInputMessages(inputMessages, continuationToken), responseUpdates) + }; + } + + try + { + // Re-ensure the run context has the resolved session before each MoveNextAsync. + // The base class RunStreamingAsync restores the original context (potentially with + // null session) after each yield, so we must re-establish it for the decorator. + EnsureRunContextHasSession(safeSession); + hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + await this.NotifyProvidersOfFailureAtEndOfRunAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); + throw; + } + } + + var chatResponse = responseUpdates.ToChatResponse(); + + var forceEndOfRunPersistence = continuationToken is not null || chatOptions?.AllowBackgroundResponses is true; + + // We can derive the type of supported session from whether we have a conversation id, + // so let's update it and set the conversation id for the service session case. + this.UpdateSessionConversationIdAtEndOfRun(safeSession, chatResponse.ConversationId, cancellationToken, forceUpdate: forceEndOfRunPersistence); + + // Notify providers of all new messages unless persistence is handled per-service-call by the decorator. + // When resuming from a continuation token or using background responses, force notification + // to send the combined data (per-service-call persistence is unreliable for these scenarios). + await this.NotifyProvidersOfNewMessagesAtEndOfRunAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, chatOptions, cancellationToken, forceNotify: forceEndOfRunPersistence).ConfigureAwait(false); + } + finally + { + await responseUpdatesEnumerator.DisposeAsync().ConfigureAwait(false); } - - var chatResponse = responseUpdates.ToChatResponse(); - - var forceEndOfRunPersistence = continuationToken is not null || chatOptions?.AllowBackgroundResponses is true; - - // We can derive the type of supported session from whether we have a conversation id, - // so let's update it and set the conversation id for the service session case. - this.UpdateSessionConversationIdAtEndOfRun(safeSession, chatResponse.ConversationId, cancellationToken, forceUpdate: forceEndOfRunPersistence); - - // Notify providers of all new messages unless persistence is handled per-service-call by the decorator. - // When resuming from a continuation token or using background responses, force notification - // to send the combined data (per-service-call persistence is unreliable for these scenarios). - await this.NotifyProvidersOfNewMessagesAtEndOfRunAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, chatOptions, cancellationToken, forceNotify: forceEndOfRunPersistence).ConfigureAwait(false); } /// diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs index ab62a38281..c2087b2d82 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs @@ -152,7 +152,14 @@ internal sealed class PerServiceCallChatHistoryPersistingChatClient : Delegating || options?.AllowBackgroundResponses is true; bool skipSimulation = isServiceManaged || isContinuationOrBackground; - var newMessages = messages as IList ?? messages.ToList(); + // Snapshot the input messages into a private list. The caller (typically + // FunctionInvokingChatClient) reuses a single mutable buffer across iterations, + // and the streaming path can defer persistence until after the caller has already + // mutated that buffer for the next iteration (e.g. on the cooperative early-exit + // path NotifyProvidersOfEarlyExitInputAsync). Aliasing the caller's list would + // then cause us to persist the wrong messages — losing FunctionResultContent and + // corrupting history with dangling FunctionCallContent. + var newMessages = messages.ToList(); // When simulating, load history and prepend it. When the service manages // history (real ConversationId) or this is a continuation/background run, @@ -174,45 +181,83 @@ internal sealed class PerServiceCallChatHistoryPersistingChatClient : Delegating throw; } - bool hasUpdates; + bool loopExitedNormally = false; + bool serviceErrorOccurred = false; try { - hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); - } - catch (Exception ex) - { - await agent.NotifyProvidersOfFailureAsync(session, ex, newMessages, options, cancellationToken).ConfigureAwait(false); - throw; - } - - while (hasUpdates) - { - var update = enumerator.Current; - responseUpdates.Add(update.Clone()); - - // If the service returned a real ConversationId on any update, remember that. - // Otherwise stamp our sentinel so FICC treats this as service-managed — - // unless this is a continuation/background run where the agent handles everything. - if (!string.IsNullOrEmpty(update.ConversationId)) - { - isServiceManaged = true; - } - else if (!skipSimulation) - { - update.ConversationId = LocalHistoryConversationId; - } - - yield return update; - + bool hasUpdates; try { hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) { + serviceErrorOccurred = true; await agent.NotifyProvidersOfFailureAsync(session, ex, newMessages, options, cancellationToken).ConfigureAwait(false); throw; } + + while (hasUpdates) + { + var update = enumerator.Current; + responseUpdates.Add(update.Clone()); + + // If the service returned a real ConversationId on any update, remember that. + // Otherwise stamp our sentinel so FICC treats this as service-managed — + // unless this is a continuation/background run where the agent handles everything. + if (!string.IsNullOrEmpty(update.ConversationId)) + { + isServiceManaged = true; + } + else if (!skipSimulation) + { + update.ConversationId = LocalHistoryConversationId; + } + + yield return update; + + try + { + hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + serviceErrorOccurred = true; + await agent.NotifyProvidersOfFailureAsync(session, ex, newMessages, options, cancellationToken).ConfigureAwait(false); + throw; + } + } + loopExitedNormally = true; + } + finally + { + // If the iterator was disposed by the consumer before completing — e.g. + // ToolApprovalAgent does `yield break` after emitting an approval request — persist + // the input messages so that any in-flight FunctionResultContent paired with + // previously-persisted FunctionCallContent is not lost between turns. We only do + // this on the cooperative-pause path; service errors deliberately do NOT persist + // input messages (history of failed calls is the caller's responsibility, e.g. + // by retrying or starting from an earlier point). + if (!loopExitedNormally && !serviceErrorOccurred) + { + // Prefer the original cancellation token so cleanup remains responsive; fall + // back to None only if the caller's token has already been canceled (otherwise + // the notify call would observe the cancellation, throw, and mask the + // original early-exit reason). + var persistToken = cancellationToken.IsCancellationRequested ? CancellationToken.None : cancellationToken; + try + { + await NotifyProvidersOfEarlyExitInputAsync(agent, session, newMessages, options, persistToken).ConfigureAwait(false); + } + catch + { + // Best-effort persistence; swallow to avoid masking the original exit reason. + } + } + + // Always dispose the underlying enumerator on every exit path (normal completion, + // exception, or early consumer disposal) to release the underlying HTTP/stream. + await enumerator.DisposeAsync().ConfigureAwait(false); } var chatResponse = responseUpdates.ToChatResponse(); @@ -236,6 +281,30 @@ internal sealed class PerServiceCallChatHistoryPersistingChatClient : Delegating } } + /// + /// Notifies s of the input messages only (no response + /// messages) on the cooperative early-exit path — e.g. when ToolApprovalAgent + /// does yield break after emitting an approval request. This ensures any + /// in-flight paired with previously-persisted + /// is not orphaned in the persisted chat history. + /// The notification is routed through the same success channel used at the end of a + /// normal run; the providers themselves decide how (or whether) to persist. + /// + private static async Task NotifyProvidersOfEarlyExitInputAsync( + ChatClientAgent agent, + ChatClientAgentSession session, + List newMessages, + ChatOptions? options, + CancellationToken cancellationToken) + { + if (newMessages.Count == 0) + { + return; + } + + await agent.NotifyProvidersOfNewMessagesAsync(session, newMessages, [], options, cancellationToken).ConfigureAwait(false); + } + /// /// Sets the sentinel on the response and session /// so that treats the conversation as service-managed. diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs index 8613d37747..8fcf99e17c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/PerServiceCallChatHistoryPersistingChatClientTests.cs @@ -1311,4 +1311,64 @@ public class PerServiceCallChatHistoryPersistingChatClientTests // Assert — session should NOT have the sentinel Assert.NotEqual(PerServiceCallChatHistoryPersistingChatClient.LocalHistoryConversationId, session!.ConversationId); } + + /// + /// Verifies that when the consumer abandons enumeration early (the streaming enumerator is + /// disposed before completing — e.g. ToolApprovalAgent.RunStreamingAsync doing a + /// yield break), the decorator still persists the input messages via its finally + /// block. This regression-guards the dropped-FunctionResultContent → HTTP 400 bug. + /// + [Fact] + public async Task RunStreamingAsync_PersistsInputMessages_WhenConsumerAbandonsEnumerationAsync() + { + // Arrange — emit multiple updates so the consumer can break after the first. + Mock mockService = new(); + mockService.Setup( + s => s.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(CreateAsyncEnumerableAsync( + new ChatResponseUpdate(ChatRole.Assistant, "first "), + new ChatResponseUpdate(ChatRole.Assistant, "second "), + new ChatResponseUpdate(ChatRole.Assistant, "third"))); + + Mock mockChatHistoryProvider = new(null, null, null); + mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]); + mockChatHistoryProvider + .Protected() + .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask>(ctx.RequestMessages.ToList())); + mockChatHistoryProvider + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + ChatHistoryProvider = mockChatHistoryProvider.Object, + RequirePerServiceCallChatHistoryPersistence = true, + }); + + // Act — consumer breaks out after the first update, mirroring ToolApprovalAgent's + // yield-break-on-approval-required path. + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + await foreach (var _ in agent.RunStreamingAsync([new(ChatRole.User, "frc-input")], session)) + { + break; + } + + // Assert — even though the consumer abandoned the stream, the input messages + // must still have been persisted (so we don't lose function-call/function-result + // pairings). + mockChatHistoryProvider + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), + ItExpr.Is(x => + x.RequestMessages.Any(m => m.Text == "frc-input") && + (x.ResponseMessages == null || !x.ResponseMessages.Any()) && + x.InvokeException == null), + ItExpr.IsAny()); + } }