From cc98ce14a26b77ff25bef2a7f7ac451ac72668f1 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 28 Oct 2026 09:29:55 +0000 Subject: [PATCH 1/5] .NET: Allow persisting partial interaction when streaming is cancelled Adds an opt-in PersistPartialMessagesOnCancellation setting on ChatClientAgentOptions (agent-level) and ChatClientAgentRunOptions (per-run override). When a streaming run is cancelled by the caller's token, the request messages and the sanitized partial response are persisted through the success channel, marked with ChatClientAgent.PartialResponseMarkerKey. Works with both end-of-run and per-service-call persistence without duplicate writes; the original OperationCanceledException is always rethrown. --- .../ChatClient/ChatClientAgent.cs | 159 +++++++ .../ChatClient/ChatClientAgentLogMessages.cs | 12 + .../ChatClient/ChatClientAgentOptions.cs | 36 ++ .../ChatClient/ChatClientAgentRunOptions.cs | 12 + ...viceCallChatHistoryPersistingChatClient.cs | 13 + .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net472/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 5 + .../netstandard2.0/PublicAPI.Unshipped.txt | 5 + .../ChatClient/ChatClientAgentOptionsTests.cs | 2 + .../ChatClientAgentRunOptionsTests.cs | 20 + ...t_PartialPersistenceOnCancellationTests.cs | 415 ++++++++++++++++++ 13 files changed, 694 insertions(+) create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 5988737629e..2fe11dcc74b 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Runtime.CompilerServices; using System.Text.Json; @@ -10,6 +11,7 @@ using Microsoft.Extensions.AI; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -56,6 +58,17 @@ namespace Microsoft.Agents.AI; /// public sealed partial class ChatClientAgent : AIAgent { + /// + /// The key of the entry that marks a response message persisted + /// from a streaming run that was cancelled before completion. + /// + /// + /// This marker is only added when or + /// is enabled. Its value is . + /// + [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] + public const string PartialResponseMarkerKey = "agent_framework.partial_response"; + private readonly ChatClientAgentOptions? _agentOptions; private readonly HashSet _aiContextProviderStateKeys; private readonly AIAgentMetadata _agentMetadata; @@ -369,6 +382,11 @@ protected override async IAsyncEnumerable RunCoreStreamingA // Ensure we start the streaming request hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); } + catch (Exception ex) when (this.ShouldNotifyProvidersOfPartialMessagesAtEndOfRun(ex, safeSession, chatOptions, options, continuationToken, responseUpdates, cancellationToken)) + { + await this.NotifyProvidersOfPartialMessagesOnCancellationAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), responseUpdates, chatOptions).ConfigureAwait(false); + throw; + } catch (Exception ex) { await this.NotifyProvidersOfFailureAtEndOfRunAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); @@ -399,6 +417,11 @@ protected override async IAsyncEnumerable RunCoreStreamingA EnsureRunContextHasSession(safeSession); hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); } + catch (Exception ex) when (this.ShouldNotifyProvidersOfPartialMessagesAtEndOfRun(ex, safeSession, chatOptions, options, continuationToken, responseUpdates, cancellationToken)) + { + await this.NotifyProvidersOfPartialMessagesOnCancellationAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), responseUpdates, chatOptions).ConfigureAwait(false); + throw; + } catch (Exception ex) { await this.NotifyProvidersOfFailureAtEndOfRunAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); @@ -942,6 +965,142 @@ private Task NotifyProvidersOfFailureAtEndOfRunAsync( return this.NotifyProvidersOfFailureAsync(session, ex, requestMessages, chatOptions, cancellationToken); } + /// + /// Determines whether providers should be notified of the partial messages of a cancelled streaming run at the end of the run, + /// instead of being notified of a failure. + /// + /// + /// This method is evaluated in an exception filter, so it must not have side effects. + /// It returns when per-service-call persistence is active (the + /// handles the notification), when chat history is + /// managed by the service, or when the run involves continuation tokens or background responses. + /// + private bool ShouldNotifyProvidersOfPartialMessagesAtEndOfRun( + Exception ex, + ChatClientAgentSession session, + ChatOptions? chatOptions, + AgentRunOptions? runOptions, + ChatClientAgentContinuationToken? continuationToken, + List responseUpdates, + CancellationToken cancellationToken) => + !this.RequiresPerServiceCallChatHistoryPersistence + && continuationToken is null + && chatOptions?.AllowBackgroundResponses is not true + && string.IsNullOrWhiteSpace(chatOptions?.ConversationId) + && string.IsNullOrWhiteSpace(session.ConversationId) + && !responseUpdates.Any(update => !string.IsNullOrWhiteSpace(update.ConversationId)) + && this.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken); + + /// + /// Determines whether providers should be notified of the partial messages received before a cancellation. + /// + /// + /// Returns only when is an caused by + /// the caller's , and partial persistence on cancellation is enabled either per run via + /// or at agent level via + /// . + /// This method is evaluated in exception filters, so it must not have side effects. + /// It is also called by . + /// + internal bool ShouldNotifyProvidersOfPartialMessagesOnCancellation(Exception ex, AgentRunOptions? runOptions, CancellationToken cancellationToken) => + ex is OperationCanceledException + && cancellationToken.IsCancellationRequested + && ((runOptions as ChatClientAgentRunOptions)?.PersistPartialMessagesOnCancellation + ?? this._agentOptions?.PersistPartialMessagesOnCancellation + ?? false); + + /// + /// Notifies the and all of the request messages and the + /// sanitized partial response of a cancelled streaming run, through the same channel used for successfully completed messages. + /// + /// + /// The notification is best effort: it uses because the caller's token is already + /// cancelled, and any failure is logged and swallowed so that the original + /// is not masked. This method is also called by . + /// + internal async Task NotifyProvidersOfPartialMessagesOnCancellationAsync( + ChatClientAgentSession session, + IEnumerable requestMessages, + IEnumerable responseUpdates, + ChatOptions? chatOptions) + { + try + { + var requestMessagesList = requestMessages as IReadOnlyCollection ?? requestMessages.ToList(); + var partialResponseMessages = CreatePartialResponseMessages(requestMessagesList, responseUpdates); + + await this.NotifyProvidersOfNewMessagesAsync(session, requestMessagesList, partialResponseMessages, chatOptions, CancellationToken.None).ConfigureAwait(false); + } + catch (Exception notifyException) + { + this._logger.LogAgentChatClientPartialMessagesNotificationFailed(this.Id, this.GetLoggingAgentName(), notifyException); + } + } + + /// + /// Builds the response messages to persist for a cancelled streaming run from the updates received so far. + /// + /// + /// Incomplete content is removed so that the stored history remains valid for the next service call: + /// tool calls without a matching tool result in the partial response, tool results without a matching tool call, + /// pending input requests (for example, ), and empty text. + /// Messages left without content are dropped, and every remaining message is marked with + /// . + /// + internal static List CreatePartialResponseMessages(IEnumerable requestMessages, IEnumerable responseUpdates) + { + IList responseMessages = responseUpdates.ToChatResponse().Messages; + + HashSet toolCallIds = new(StringComparer.Ordinal); + HashSet toolResultIds = new(StringComparer.Ordinal); + + foreach (var content in requestMessages.SelectMany(m => m.Contents)) + { + if (content is ToolCallContent toolCall) + { + toolCallIds.Add(toolCall.CallId); + } + } + + foreach (var content in responseMessages.SelectMany(m => m.Contents)) + { + switch (content) + { + case ToolCallContent toolCall: + toolCallIds.Add(toolCall.CallId); + break; + case ToolResultContent toolResult: + toolResultIds.Add(toolResult.CallId); + break; + } + } + + List partialMessages = []; + foreach (var message in responseMessages) + { + List contents = message.Contents.Where(content => content switch + { + ToolCallContent toolCall => toolResultIds.Contains(toolCall.CallId), + ToolResultContent toolResult => toolCallIds.Contains(toolResult.CallId), + InputRequestContent => false, + TextContent text => !string.IsNullOrEmpty(text.Text), + _ => true, + }).ToList(); + + if (contents.Count == 0) + { + continue; + } + + message.Contents = contents; + message.AdditionalProperties ??= new(); + message.AdditionalProperties[PartialResponseMarkerKey] = true; + partialMessages.Add(message); + } + + return partialMessages; + } + /// /// Gets a value indicating whether the agent is configured to simulate service-stored chat history. /// When , end-of-run persistence and history loading are skipped because a diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentLogMessages.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentLogMessages.cs index dde1d97ba4b..f012f50a0b9 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentLogMessages.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentLogMessages.cs @@ -97,4 +97,16 @@ public static partial void LogAgentChatClientBackgroundResponseFallback( this ILogger logger, string agentId, string agentName); + + /// + /// Logs a warning when notifying providers of the partial messages of a cancelled streaming run fails. + /// + [LoggerMessage( + Level = LogLevel.Warning, + Message = "Agent {AgentId}/{AgentName}: Failed to notify providers of the partial messages of a cancelled streaming run.")] + public static partial void LogAgentChatClientPartialMessagesNotificationFailed( + this ILogger logger, + string agentId, + string agentName, + Exception exception); } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs index 75127b41944..97f8acdf402 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs @@ -311,6 +311,41 @@ public sealed class ChatClientAgentOptions [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] public bool EnableInvocableFunctionBypassing { get; set; } + /// + /// Gets or sets a value indicating whether the should persist the request + /// messages and the partial response received so far when a streaming run is cancelled via its + /// . + /// + /// + /// + /// By default, a cancelled run is treated as a failed run: providers are notified with + /// set, and nothing is stored. + /// + /// + /// When set to and a streaming run is cancelled by the caller's cancellation token, the agent + /// builds a response from the updates received so far, removes incomplete content (for example, a + /// without a matching , or a pending + /// ), marks each persisted response message with + /// in , + /// and notifies the and the + /// as for a successful invocation. The original is always rethrown, + /// and failures while persisting the partial interaction are logged and swallowed. + /// + /// + /// This setting only applies to framework-managed chat history. It has no effect on non-streaming runs, + /// on service-managed chat history, on runs that use continuation tokens or background responses, or when + /// the stream is disposed early by the consumer without cancelling the token. + /// + /// + /// This value can be overridden per run via . + /// + /// + /// + /// Default is . + /// + [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] + public bool PersistPartialMessagesOnCancellation { get; set; } + /// /// Creates a new instance of with the same values as this instance. /// @@ -333,5 +368,6 @@ public ChatClientAgentOptions Clone() DisableApprovalNotRequiredFunctionBypassing = this.DisableApprovalNotRequiredFunctionBypassing, DisableApprovalResponseBinding = this.DisableApprovalResponseBinding, EnableInvocableFunctionBypassing = this.EnableInvocableFunctionBypassing, + PersistPartialMessagesOnCancellation = this.PersistPartialMessagesOnCancellation, }; } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs index cf35aa80a16..08e23040adc 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs @@ -1,7 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Diagnostics.CodeAnalysis; using Microsoft.Extensions.AI; +using Microsoft.Shared.DiagnosticIds; namespace Microsoft.Agents.AI; @@ -35,6 +37,7 @@ private ChatClientAgentRunOptions(ChatClientAgentRunOptions options) { this.ChatOptions = options.ChatOptions?.Clone(); this.ChatClientFactory = options.ChatClientFactory; + this.PersistPartialMessagesOnCancellation = options.PersistPartialMessagesOnCancellation; } /// @@ -62,6 +65,15 @@ private ChatClientAgentRunOptions(ChatClientAgentRunOptions options) /// public Func? ChatClientFactory { get; set; } + /// + /// Gets or sets a per-run override for . + /// + /// + /// (default) to use the agent-level setting; otherwise the value to use for this run only. + /// + [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] + public bool? PersistPartialMessagesOnCancellation { get; set; } + /// public override AgentRunOptions Clone() => new ChatClientAgentRunOptions(this); } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs index c2087b2d829..894b98ed1ca 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs @@ -145,6 +145,7 @@ public override async IAsyncEnumerable GetStreamingResponseA [EnumeratorCancellation] CancellationToken cancellationToken = default) { var (agent, session) = GetRequiredAgentAndSession(); + var runOptions = AIAgent.CurrentRunContext?.RunOptions; options = StripLocalHistoryConversationId(options); bool isServiceManaged = !string.IsNullOrEmpty(options?.ConversationId); @@ -190,6 +191,12 @@ public override async IAsyncEnumerable GetStreamingResponseA { hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); } + catch (Exception ex) when (!skipSimulation && !isServiceManaged && agent.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken)) + { + serviceErrorOccurred = true; + await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); + throw; + } catch (Exception ex) { serviceErrorOccurred = true; @@ -220,6 +227,12 @@ public override async IAsyncEnumerable GetStreamingResponseA { hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); } + catch (Exception ex) when (!skipSimulation && !isServiceManaged && agent.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken)) + { + serviceErrorOccurred = true; + await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); + throw; + } catch (Exception ex) { serviceErrorOccurred = true; diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 1bd39868305..e53770fb248 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -25,3 +25,8 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.get -> bool +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.get -> bool? +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]const Microsoft.Agents.AI.ChatClientAgent.PartialResponseMarkerKey = "agent_framework.partial_response" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt index 1bd39868305..e53770fb248 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -25,3 +25,8 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.get -> bool +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.get -> bool? +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]const Microsoft.Agents.AI.ChatClientAgent.PartialResponseMarkerKey = "agent_framework.partial_response" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt index 1bd39868305..e53770fb248 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -25,3 +25,8 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.get -> bool +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.get -> bool? +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]const Microsoft.Agents.AI.ChatClientAgent.PartialResponseMarkerKey = "agent_framework.partial_response" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt index 1bd39868305..e53770fb248 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -25,3 +25,8 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.get -> bool +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.get -> bool? +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]const Microsoft.Agents.AI.ChatClientAgent.PartialResponseMarkerKey = "agent_framework.partial_response" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index 1bd39868305..e53770fb248 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -25,3 +25,8 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.get -> bool +[MAAI001]Microsoft.Agents.AI.ChatClientAgentOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.get -> bool? +[MAAI001]Microsoft.Agents.AI.ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation.set -> void +[MAAI001]const Microsoft.Agents.AI.ChatClientAgent.PartialResponseMarkerKey = "agent_framework.partial_response" -> string! diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs index 27d9370cc9b..09a238a17dd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs @@ -138,6 +138,7 @@ public void Clone_CreatesDeepCopyWithSameValues() ThrowOnChatHistoryProviderConflict = false, DisableApprovalNotRequiredFunctionBypassing = true, EnableInvocableFunctionBypassing = true, + PersistPartialMessagesOnCancellation = true, }; // Act @@ -157,6 +158,7 @@ public void Clone_CreatesDeepCopyWithSameValues() Assert.Equal(original.ThrowOnChatHistoryProviderConflict, clone.ThrowOnChatHistoryProviderConflict); Assert.Equal(original.DisableApprovalNotRequiredFunctionBypassing, clone.DisableApprovalNotRequiredFunctionBypassing); Assert.Equal(original.EnableInvocableFunctionBypassing, clone.EnableInvocableFunctionBypassing); + Assert.Equal(original.PersistPartialMessagesOnCancellation, clone.PersistPartialMessagesOnCancellation); // ChatOptions should be cloned, not the same reference Assert.NotSame(original.ChatOptions, clone.ChatOptions); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs index 7a00a9f796a..dabfcb4141d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs @@ -45,6 +45,26 @@ public void ChatOptionsPropertyIsReadOnly() Assert.Equal(200, retrievedOptions.MaxOutputTokens); // Ensure the change is reflected } + /// + /// Verify that Clone copies PersistPartialMessagesOnCancellation. + /// + [Theory] + [InlineData(null)] + [InlineData(true)] + [InlineData(false)] + public void Clone_CopiesPersistPartialMessagesOnCancellation(bool? value) + { + // Arrange + var runOptions = new ChatClientAgentRunOptions { PersistPartialMessagesOnCancellation = value }; + + // Act + var clone = Assert.IsType(runOptions.Clone()); + + // Assert + Assert.NotSame(runOptions, clone); + Assert.Equal(value, clone.PersistPartialMessagesOnCancellation); + } + #region ChatClientFactory Tests /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs new file mode 100644 index 00000000000..f0c7fcb0a2e --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs @@ -0,0 +1,415 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; +using Moq.Protected; + +namespace Microsoft.Agents.AI.UnitTests; + +/// +/// Contains unit tests that verify the behavior of +/// and when a streaming run is cancelled. +/// +public class ChatClientAgent_PartialPersistenceOnCancellationTests +{ + [Fact] + public async Task RunStreamingAsync_FlagDisabled_DoesNotPersistOnCancellationAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new()); + var session = await agent.CreateSessionAsync(); + + // Act + var received = await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + Assert.Equal("partial", string.Concat(received.Select(u => u.Text))); + Assert.Empty(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PersistsRequestAndPartialResponseAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("par", "m1"), TextUpdate("tial", "m1")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + var received = await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + Assert.Equal("partial", string.Concat(received.Select(u => u.Text))); + var history = GetHistory(agent, session); + Assert.Equal(2, history.Count); + Assert.Equal(ChatRole.User, history[0].Role); + Assert.Equal("Hello", history[0].Text); + Assert.False(IsMarkedAsPartial(history[0])); + Assert.Equal(ChatRole.Assistant, history[1].Role); + Assert.Equal("partial", history[1].Text); + Assert.True(IsMarkedAsPartial(history[1])); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_CancelledBeforeFirstUpdate_PersistsOnlyRequestAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + var received = await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + Assert.Empty(received); + var history = GetHistory(agent, session); + var message = Assert.Single(history); + Assert.Equal(ChatRole.User, message.Role); + Assert.Equal("Hello", message.Text); + } + + [Fact] + public async Task RunStreamingAsync_RunOptionsFalse_OverridesAgentLevelTrueAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, new ChatClientAgentRunOptions { PersistPartialMessagesOnCancellation = false }, cts.Token); + + // Assert + Assert.Empty(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_RunOptionsTrue_OverridesAgentLevelFalseAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new()); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, new ChatClientAgentRunOptions { PersistPartialMessagesOnCancellation = true }, cts.Token); + + // Assert + var history = GetHistory(agent, session); + Assert.Equal(2, history.Count); + Assert.Equal("partial", history[1].Text); + Assert.True(IsMarkedAsPartial(history[1])); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_RemovesUnmatchedFunctionCallAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var tool = AIFunctionFactory.Create(() => "Sunny", "GetWeather"); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall( + [ + TextUpdate("Let me check", "m1"), + new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }, + ], + CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { ChatOptions = new() { Tools = [tool] }, PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + var history = GetHistory(agent, session); + Assert.Equal(2, history.Count); + Assert.Equal(ChatRole.User, history[0].Role); + Assert.Equal(ChatRole.Assistant, history[1].Role); + Assert.Equal("Let me check", history[1].Text); + Assert.DoesNotContain(history.SelectMany(m => m.Contents), c => c is FunctionCallContent); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_KeepsMatchedFunctionCallAndResultAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var tool = AIFunctionFactory.Create(() => "Sunny", "GetWeather"); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([TextUpdate("It is sun", "m2")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { ChatOptions = new() { Tools = [tool] }, PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + AssertToolLoopPartialHistory(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_NonCallerCancellation_DoesNotPersistAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.NonCallerCancellation)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + Assert.False(cts.IsCancellationRequested); + Assert.Empty(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PersistenceFailure_RethrowsOriginalCancellationAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels)); + + 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()) + .Throws(new InvalidOperationException("Storage failure")); + + ChatClientAgent agent = new(chatClient, options: new() + { + ChatHistoryProvider = mockChatHistoryProvider.Object, + PersistPartialMessagesOnCancellation = true, + }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + mockChatHistoryProvider + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), + ItExpr.Is(x => x.InvokeException == null && x.ResponseMessages!.Count() == 1), + ItExpr.IsAny()); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_PersistsEachMessageOnceAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var tool = AIFunctionFactory.Create(() => "Sunny", "GetWeather"); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([TextUpdate("It is sun", "m2")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() + { + ChatOptions = new() { Tools = [tool] }, + RequirePerServiceCallChatHistoryPersistence = true, + PersistPartialMessagesOnCancellation = true, + }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + AssertToolLoopPartialHistory(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagDisabled_PerServiceCallPersistence_DoesNotPersistCancelledCallAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var tool = AIFunctionFactory.Create(() => "Sunny", "GetWeather"); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([TextUpdate("It is sun", "m2")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() + { + ChatOptions = new() { Tools = [tool] }, + RequirePerServiceCallChatHistoryPersistence = true, + }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + var history = GetHistory(agent, session); + Assert.Equal(2, history.Count); + Assert.Equal(ChatRole.User, history[0].Role); + Assert.Contains(history[1].Contents, c => c is FunctionCallContent); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_BackgroundResponses_DoesNotPersistAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, new ChatClientAgentRunOptions { AllowBackgroundResponses = true }, cts.Token); + + // Assert + Assert.Empty(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_ServiceManagedHistory_DoesNotPersistAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient(cts, new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, "partial") { MessageId = "m1", ConversationId = "conv-1" }], CallEnd.CallerCancels)); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = (ChatClientAgentSession)await agent.CreateSessionAsync(); + + // Act + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Assert + Assert.Null(session.ConversationId); + Assert.Empty(GetHistory(agent, session)); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_NextRunIncludesPersistedPartialInteractionAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([TextUpdate("partial", "m1")], CallEnd.CallerCancels), + new StreamingCall([TextUpdate("complete", "m2")])); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + + // Act + await foreach (var _ in agent.RunStreamingAsync([new(ChatRole.User, "Continue")], session)) + { + } + + // Assert + var secondRequest = chatClient.Requests[1]; + Assert.Equal(["Hello", "partial", "Continue"], secondRequest.Select(m => m.Text)); + Assert.Equal(4, GetHistory(agent, session).Count); + } + + private static void AssertToolLoopPartialHistory(List history) + { + Assert.Equal(4, history.Count); + Assert.Equal(ChatRole.User, history[0].Role); + Assert.Equal(ChatRole.Assistant, history[1].Role); + Assert.Equal("call1", Assert.IsType(Assert.Single(history[1].Contents)).CallId); + Assert.Equal(ChatRole.Tool, history[2].Role); + Assert.Equal("call1", Assert.IsType(Assert.Single(history[2].Contents)).CallId); + Assert.Equal(ChatRole.Assistant, history[3].Role); + Assert.Equal("It is sun", history[3].Text); + Assert.True(IsMarkedAsPartial(history[3])); + } + + private static async Task> ConsumeUntilCancelledAsync(ChatClientAgent agent, AgentSession session, AgentRunOptions? options, CancellationToken cancellationToken) + { + List received = []; + await Assert.ThrowsAnyAsync(async () => + { + await foreach (var update in agent.RunStreamingAsync([new(ChatRole.User, "Hello")], session, options, cancellationToken)) + { + received.Add(update); + } + }); + + return received; + } + + private static List GetHistory(ChatClientAgent agent, AgentSession session) => + Assert.IsType(agent.ChatHistoryProvider).GetMessages(session).ToList(); + + private static bool IsMarkedAsPartial(ChatMessage message) => + message.AdditionalProperties?.TryGetValue(ChatClientAgent.PartialResponseMarkerKey, out var value) is true && value is true; + + private static ChatResponseUpdate TextUpdate(string text, string messageId) => + new(ChatRole.Assistant, text) { MessageId = messageId }; + + private enum CallEnd + { + Complete, + CallerCancels, + NonCallerCancellation, + } + + private sealed record StreamingCall(ChatResponseUpdate[] Updates, CallEnd End = CallEnd.Complete); + + /// + /// A chat client that streams scripted updates for each service call and optionally ends the stream + /// with a cancellation, either triggered on the caller's token or independent of it. + /// + private sealed class ScriptedStreamingChatClient(CancellationTokenSource callerCancellation, params StreamingCall[] calls) : IChatClient + { + private int _callIndex; + + public List> Requests { get; } = []; + + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + this.Requests.Add(messages.ToList()); + var call = calls[this._callIndex++]; + + foreach (var update in call.Updates) + { + await Task.Yield(); + yield return update; + } + + switch (call.End) + { + case CallEnd.CallerCancels: + callerCancellation.Cancel(); + cancellationToken.ThrowIfCancellationRequested(); + break; + case CallEnd.NonCallerCancellation: + throw new OperationCanceledException("Cancelled independently of the caller's token."); + } + } + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() + { + } + } +} From fb95ef47fcf7d0bcc97262e711fd3168d4c0a2a5 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 28 Oct 2026 10:47:25 +0000 Subject: [PATCH 2/5] .NET: Skip partial persistence when a continuation token was received --- .../ChatClient/ChatClientAgent.cs | 7 ++-- ...t_PartialPersistenceOnCancellationTests.cs | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 2fe11dcc74b..d5780c33b94 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -973,7 +973,10 @@ private Task NotifyProvidersOfFailureAtEndOfRunAsync( /// This method is evaluated in an exception filter, so it must not have side effects. /// It returns when per-service-call persistence is active (the /// handles the notification), when chat history is - /// managed by the service, or when the run involves continuation tokens or background responses. + /// managed by the service, or when the run involves background responses or continuation tokens, either supplied by + /// the caller or received in the response updates. In the latter case, the caller can resume the run with the received + /// token, and the resumed run persists the request and the updates received so far, so persisting them here would + /// duplicate them in the chat history. /// private bool ShouldNotifyProvidersOfPartialMessagesAtEndOfRun( Exception ex, @@ -988,7 +991,7 @@ private bool ShouldNotifyProvidersOfPartialMessagesAtEndOfRun( && chatOptions?.AllowBackgroundResponses is not true && string.IsNullOrWhiteSpace(chatOptions?.ConversationId) && string.IsNullOrWhiteSpace(session.ConversationId) - && !responseUpdates.Any(update => !string.IsNullOrWhiteSpace(update.ConversationId)) + && !responseUpdates.Any(update => !string.IsNullOrWhiteSpace(update.ConversationId) || update.ContinuationToken is not null) && this.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken); /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs index f0c7fcb0a2e..4b7bd0fbfba 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs @@ -299,6 +299,39 @@ public async Task RunStreamingAsync_FlagEnabled_ServiceManagedHistory_DoesNotPer Assert.Empty(GetHistory(agent, session)); } + [Fact] + public async Task RunStreamingAsync_FlagEnabled_ContinuationTokenReceived_DoesNotDuplicateHistoryOnResumeAsync() + { + // Arrange + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, "partial") { MessageId = "m1", ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) }], CallEnd.CallerCancels), + new StreamingCall([TextUpdate(" rest", "m1")])); + ChatClientAgent agent = new(chatClient, options: new() { PersistPartialMessagesOnCancellation = true }); + var session = await agent.CreateSessionAsync(); + + // Act + var received = await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + var historyAfterCancellation = GetHistory(agent, session); + + var continuationToken = Assert.Single(received).ContinuationToken; + Assert.NotNull(continuationToken); + await foreach (var _ in agent.RunStreamingAsync(session, options: new AgentRunOptions { ContinuationToken = continuationToken })) + { + } + + // Assert + Assert.Empty(historyAfterCancellation); + var history = GetHistory(agent, session); + Assert.Equal(2, history.Count); + Assert.Equal(ChatRole.User, history[0].Role); + Assert.Equal("Hello", history[0].Text); + Assert.Equal(ChatRole.Assistant, history[1].Role); + Assert.Equal("partial rest", history[1].Text); + Assert.False(IsMarkedAsPartial(history[1])); + } + [Fact] public async Task RunStreamingAsync_FlagEnabled_NextRunIncludesPersistedPartialInteractionAsync() { From 14460af38e4d185c1d0f57ad2dba1ef7ee50f222 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 28 Oct 2026 11:21:11 +0000 Subject: [PATCH 3/5] .NET: Skip per-service-call partial persistence when a continuation token was received --- ...viceCallChatHistoryPersistingChatClient.cs | 31 ++++++++++++++- ...t_PartialPersistenceOnCancellationTests.cs | 39 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs index 894b98ed1ca..50f775dd965 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs @@ -184,6 +184,7 @@ public override async IAsyncEnumerable GetStreamingResponseA bool loopExitedNormally = false; bool serviceErrorOccurred = false; + bool continuationTokenReceived = false; try { bool hasUpdates; @@ -191,7 +192,7 @@ public override async IAsyncEnumerable GetStreamingResponseA { hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); } - catch (Exception ex) when (!skipSimulation && !isServiceManaged && agent.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken)) + catch (Exception ex) when (ShouldNotifyProvidersOfPartialMessages(agent, ex, runOptions, skipSimulation || isServiceManaged || continuationTokenReceived, cancellationToken)) { serviceErrorOccurred = true; await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); @@ -209,6 +210,12 @@ public override async IAsyncEnumerable GetStreamingResponseA var update = enumerator.Current; responseUpdates.Add(update.Clone()); + // Track continuation tokens on the original update, because ChatResponseUpdate.Clone() does not copy them. + if (update.ContinuationToken is not null) + { + continuationTokenReceived = true; + } + // 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. @@ -227,7 +234,7 @@ public override async IAsyncEnumerable GetStreamingResponseA { hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); } - catch (Exception ex) when (!skipSimulation && !isServiceManaged && agent.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken)) + catch (Exception ex) when (ShouldNotifyProvidersOfPartialMessages(agent, ex, runOptions, skipSimulation || isServiceManaged || continuationTokenReceived, cancellationToken)) { serviceErrorOccurred = true; await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); @@ -294,6 +301,26 @@ public override async IAsyncEnumerable GetStreamingResponseA } } + /// + /// Determines whether providers should be notified of the partial messages of the current service call + /// after a cancellation, instead of being notified of a failure. + /// + /// + /// This method is evaluated in exception filters, so it must not have side effects. + /// Callers pass as when chat history is not simulated by this + /// decorator (service-managed history, continuation or background runs) or when any received update carried a continuation + /// token. In the latter case, the caller can resume the run with the received token, and the resumed run persists the + /// request and the updates received so far, so persisting them here would duplicate them in the chat history. + /// + private static bool ShouldNotifyProvidersOfPartialMessages( + ChatClientAgent agent, + Exception ex, + AgentRunOptions? runOptions, + bool skipPartialMessages, + CancellationToken cancellationToken) => + !skipPartialMessages + && agent.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken); + /// /// Notifies s of the input messages only (no response /// messages) on the cooperative early-exit path — e.g. when ToolApprovalAgent diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs index 4b7bd0fbfba..65097b48751 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs @@ -332,6 +332,19 @@ public async Task RunStreamingAsync_FlagEnabled_ContinuationTokenReceived_DoesNo Assert.False(IsMarkedAsPartial(history[1])); } + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_ContinuationTokenReceived_DoesNotDuplicateHistoryOnResumeAsync() + { + // Arrange & Act + var (historyAfterCancellation, history) = await CancelAndResumePerServiceCallRunAsync(persistPartialMessagesOnCancellation: true); + var (_, baselineHistory) = await CancelAndResumePerServiceCallRunAsync(persistPartialMessagesOnCancellation: false); + + // Assert + Assert.Empty(historyAfterCancellation); + Assert.DoesNotContain(history, IsMarkedAsPartial); + Assert.Equal(baselineHistory.Select(m => (m.Role, m.Text)), history.Select(m => (m.Role, m.Text))); + } + [Fact] public async Task RunStreamingAsync_FlagEnabled_NextRunIncludesPersistedPartialInteractionAsync() { @@ -369,6 +382,32 @@ private static void AssertToolLoopPartialHistory(List history) Assert.True(IsMarkedAsPartial(history[3])); } + private static async Task<(List HistoryAfterCancellation, List History)> CancelAndResumePerServiceCallRunAsync(bool persistPartialMessagesOnCancellation) + { + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, "partial") { MessageId = "m1", ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }) }], CallEnd.CallerCancels), + new StreamingCall([TextUpdate(" rest", "m1")])); + ChatClientAgent agent = new(chatClient, options: new() + { + RequirePerServiceCallChatHistoryPersistence = true, + PersistPartialMessagesOnCancellation = persistPartialMessagesOnCancellation, + }); + var session = await agent.CreateSessionAsync(); + + var received = await ConsumeUntilCancelledAsync(agent, session, null, cts.Token); + var historyAfterCancellation = GetHistory(agent, session); + + var continuationToken = Assert.Single(received).ContinuationToken; + Assert.NotNull(continuationToken); + await foreach (var _ in agent.RunStreamingAsync(session, options: new AgentRunOptions { ContinuationToken = continuationToken })) + { + } + + return (historyAfterCancellation, GetHistory(agent, session)); + } + private static async Task> ConsumeUntilCancelledAsync(ChatClientAgent agent, AgentSession session, AgentRunOptions? options, CancellationToken cancellationToken) { List received = []; From 95315cb18f4ae3e3957e6f1883196a0d8b0fec7d Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 28 Oct 2026 15:21:19 +0000 Subject: [PATCH 4/5] .NET: Consume settled tool approval state when partial messages are persisted on cancellation When a cancelled streaming run persists its partial interaction, remove from the session the pending approval requests and stored auto-approved calls whose results are now in the chat history, keeping unfinished calls for retry. With per-service-call persistence the next run handles approvals before the history is loaded, so leftover state would settle or execute the same calls again. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...lNotRequiredFunctionBypassingChatClient.cs | 36 +++++ .../ApprovalResponseBindingChatClient.cs | 23 +++ .../ChatClient/ChatClientAgent.cs | 39 ++++- ...viceCallChatHistoryPersistingChatClient.cs | 4 +- ...equiredFunctionBypassingChatClientTests.cs | 46 ++++++ .../ApprovalResponseBindingChatClientTests.cs | 36 +++++ ...t_PartialPersistenceOnCancellationTests.cs | 137 ++++++++++++++++++ 7 files changed, 318 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalNotRequiredFunctionBypassingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalNotRequiredFunctionBypassingChatClient.cs index 6bbfd9815f5..c6bc4f1bf81 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalNotRequiredFunctionBypassingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalNotRequiredFunctionBypassingChatClient.cs @@ -159,6 +159,42 @@ public override async IAsyncEnumerable GetStreamingResponseA } } + /// + /// Removes from the stored auto-approved function calls those whose call already has a result in the chat history, + /// keeping the calls that are still unfinished. + /// + /// + /// The stored calls are normally cleared only when a run completes, so that a failed run injects them again on the + /// next run. When a cancelled streaming run still persists its partial interaction (see + /// ), the results of the injected calls + /// are stored even though the run did not complete. Those calls must then be removed: with per-service-call + /// persistence the next run does not see the stored history before , so a + /// leftover call would be injected as approved and executed a second time. + /// + /// The session holding the stored auto-approved function calls. + /// The call ids whose results have been persisted. + internal static void RemoveSettledAutoApprovals(AgentSession session, HashSet settledCallIds) + { + if (!session.StateBag.TryGetValue( + StateBagKey, + out List? pendingRequests, + AgentJsonUtilities.DefaultOptions) + || pendingRequests is not { Count: > 0 } + || pendingRequests.RemoveAll(request => settledCallIds.Contains(request.ToolCall.CallId)) == 0) + { + return; + } + + if (pendingRequests.Count > 0) + { + session.StateBag.SetValue(StateBagKey, pendingRequests, AgentJsonUtilities.DefaultOptions); + } + else + { + session.StateBag.TryRemoveValue(StateBagKey); + } + } + /// /// Attempts to get the current from the ambient run context. When no run /// context or session is available, logs a warning (once per instance) and returns diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalResponseBindingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalResponseBindingChatClient.cs index ea5664c9f48..4dc67d32e55 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalResponseBindingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalResponseBindingChatClient.cs @@ -611,6 +611,29 @@ private void MergePendingApprovalRequests(List emitt } } + /// + /// Removes from the recorded pending approval requests those whose tool call already has a result in the chat + /// history, keeping the requests whose calls are still unfinished. + /// + /// + /// The records are normally consumed only when a run completes, so that a failed run can be retried with the same + /// approval response. When a cancelled streaming run still persists its partial interaction (see + /// ), the results of the approved or + /// rejected calls are stored even though the run did not complete. Their records must then be consumed: with + /// per-service-call persistence the next run does not see the stored history before this decorator, so a leftover + /// record would be rejected again and produce a second result for a call that is already settled. + /// + /// The session holding the recorded pending approval requests. + /// The call ids whose results have been persisted. + internal static void RemoveSettledPendingApprovalRequests(AgentSession session, HashSet settledCallIds) + { + var pendingRequests = LoadPendingApprovalRequests(session); + if (pendingRequests.RemoveAll(request => settledCallIds.Contains(request.ToolCall.CallId)) > 0) + { + SavePendingApprovalRequests(pendingRequests, session); + } + } + private static List LoadPendingApprovalRequests(AgentSession session) => session.StateBag.TryGetValue>(StateBagKey, out var pendingRequests, AgentJsonUtilities.DefaultOptions) && pendingRequests is not null diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index d5780c33b94..32a3edf17f5 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -1020,12 +1020,23 @@ ex is OperationCanceledException /// The notification is best effort: it uses because the caller's token is already /// cancelled, and any failure is logged and swallowed so that the original /// is not masked. This method is also called by . + /// When the notification succeeds, the tool approval state kept in the session for calls whose results are now in the + /// chat history is consumed, as it would be at the end of a completed run, while the state of unfinished calls is kept. /// + /// The session of the cancelled run. + /// The request messages to persist. + /// The response updates received before the cancellation. + /// The chat options of the cancelled run. + /// + /// The messages known to be stored in the chat history, including , used to find + /// the tool calls that are settled. When , only are used. + /// internal async Task NotifyProvidersOfPartialMessagesOnCancellationAsync( ChatClientAgentSession session, IEnumerable requestMessages, IEnumerable responseUpdates, - ChatOptions? chatOptions) + ChatOptions? chatOptions, + IEnumerable? historyMessages = null) { try { @@ -1033,6 +1044,8 @@ internal async Task NotifyProvidersOfPartialMessagesOnCancellationAsync( var partialResponseMessages = CreatePartialResponseMessages(requestMessagesList, responseUpdates); await this.NotifyProvidersOfNewMessagesAsync(session, requestMessagesList, partialResponseMessages, chatOptions, CancellationToken.None).ConfigureAwait(false); + + RemoveSettledToolApprovalState(session, (historyMessages ?? requestMessagesList).Concat(partialResponseMessages)); } catch (Exception notifyException) { @@ -1040,6 +1053,30 @@ internal async Task NotifyProvidersOfPartialMessagesOnCancellationAsync( } } + /// + /// Removes the tool approval state kept in the session for the tool calls whose results are contained in the + /// specified messages, which must all be stored in the chat history. + /// + /// + /// and + /// keep their state when a streaming run does not complete, so that it can be retried. Once the results of those calls + /// are persisted, the state must be consumed, otherwise the next run would settle or execute the same calls again. + /// + private static void RemoveSettledToolApprovalState(ChatClientAgentSession session, IEnumerable persistedMessages) + { + HashSet settledCallIds = new( + persistedMessages.SelectMany(message => message.Contents).OfType().Select(result => result.CallId), + StringComparer.Ordinal); + + if (settledCallIds.Count == 0) + { + return; + } + + ApprovalResponseBindingChatClient.RemoveSettledPendingApprovalRequests(session, settledCallIds); + ApprovalNotRequiredFunctionBypassingChatClient.RemoveSettledAutoApprovals(session, settledCallIds); + } + /// /// Builds the response messages to persist for a cancelled streaming run from the updates received so far. /// diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs index 50f775dd965..a4c614c0084 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/PerServiceCallChatHistoryPersistingChatClient.cs @@ -195,7 +195,7 @@ public override async IAsyncEnumerable GetStreamingResponseA catch (Exception ex) when (ShouldNotifyProvidersOfPartialMessages(agent, ex, runOptions, skipSimulation || isServiceManaged || continuationTokenReceived, cancellationToken)) { serviceErrorOccurred = true; - await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); + await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options, messagesForService).ConfigureAwait(false); throw; } catch (Exception ex) @@ -237,7 +237,7 @@ public override async IAsyncEnumerable GetStreamingResponseA catch (Exception ex) when (ShouldNotifyProvidersOfPartialMessages(agent, ex, runOptions, skipSimulation || isServiceManaged || continuationTokenReceived, cancellationToken)) { serviceErrorOccurred = true; - await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options).ConfigureAwait(false); + await agent.NotifyProvidersOfPartialMessagesOnCancellationAsync(session, newMessages, responseUpdates, options, messagesForService).ConfigureAwait(false); throw; } catch (Exception ex) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs index ed37bf6dc44..37ca2d21590 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalNotRequiredFunctionBypassingChatClientTests.cs @@ -814,6 +814,52 @@ public void WithDefaultAgentMiddleware_DisableApprovalNotRequiredFunctionBypassi #endregion + #region Settled Auto-Approval Tests + + [Fact] + public void RemoveSettledAutoApprovals_RemovesOnlySettledCalls() + { + // Arrange + var session = new ChatClientAgentSession(); + session.StateBag.SetValue( + ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, + new List + { + new("req1", new FunctionCallContent("call1", "toolA")), + new("req2", new FunctionCallContent("call2", "toolB")), + }, + AgentJsonUtilities.DefaultOptions); + + // Act + ApprovalNotRequiredFunctionBypassingChatClient.RemoveSettledAutoApprovals(session, new HashSet(StringComparer.Ordinal) { "call1" }); + + // Assert + Assert.True(session.StateBag.TryGetValue>( + ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions)); + var remaining = Assert.Single(stored!); + Assert.Equal("call2", remaining.ToolCall.CallId); + } + + [Fact] + public void RemoveSettledAutoApprovals_AllSettled_RemovesState() + { + // Arrange + var session = new ChatClientAgentSession(); + session.StateBag.SetValue( + ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, + new List { new("req1", new FunctionCallContent("call1", "toolA")) }, + AgentJsonUtilities.DefaultOptions); + + // Act + ApprovalNotRequiredFunctionBypassingChatClient.RemoveSettledAutoApprovals(session, new HashSet(StringComparer.Ordinal) { "call1" }); + + // Assert + Assert.False(session.StateBag.TryGetValue>( + ApprovalNotRequiredFunctionBypassingChatClient.StateBagKey, out _, AgentJsonUtilities.DefaultOptions)); + } + + #endregion + #region Usage Pass-Through Tests /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalResponseBindingChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalResponseBindingChatClientTests.cs index e30cfc0d58e..427fe66717c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalResponseBindingChatClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ApprovalResponseBindingChatClientTests.cs @@ -1060,6 +1060,42 @@ private static async Task RecordRequestAsync(ChatClientAgentSession session, Too await RunAsync(decorator, session, [new ChatMessage(ChatRole.User, "Hi")]); } + [Fact] + public async Task RemoveSettledPendingApprovalRequests_RemovesOnlySettledRequestsAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + await RecordRequestsAsync(session, + [ + new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "toolA")), + new ToolApprovalRequestContent("req2", new FunctionCallContent("call2", "toolB")), + ]); + + // Act + ApprovalResponseBindingChatClient.RemoveSettledPendingApprovalRequests(session, new HashSet(StringComparer.Ordinal) { "call1" }); + + // Assert + Assert.True(session.StateBag.TryGetValue>( + ApprovalResponseBindingChatClient.StateBagKey, out var pending)); + var remaining = Assert.Single(pending!); + Assert.Equal("req2", remaining.RequestId); + } + + [Fact] + public async Task RemoveSettledPendingApprovalRequests_AllSettled_RemovesStateAsync() + { + // Arrange + var session = new ChatClientAgentSession(); + await RecordRequestsAsync(session, [new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "toolA"))]); + + // Act + ApprovalResponseBindingChatClient.RemoveSettledPendingApprovalRequests(session, new HashSet(StringComparer.Ordinal) { "call1" }); + + // Assert + Assert.False(session.StateBag.TryGetValue>( + ApprovalResponseBindingChatClient.StateBagKey, out _)); + } + private static async Task RecordRequestsAsync(ChatClientAgentSession session, IList requests) { var inner = CreateMockChatClient((_, _, _) => diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs index 65097b48751..972117b7a27 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_PartialPersistenceOnCancellationTests.cs @@ -345,6 +345,108 @@ public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_Contin Assert.Equal(baselineHistory.Select(m => (m.Role, m.Text)), history.Select(m => (m.Role, m.Text))); } + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_ApprovedToolThenCancellation_NextTurnDoesNotSettleCallAgainAsync() + { + // Arrange + int executions = 0; + var tool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => { executions++; return "Sunny"; }, "GetWeather")); + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([TextUpdate("It is sun", "m2")], CallEnd.CallerCancels), + new StreamingCall([TextUpdate("You're welcome", "m3")])); + ChatClientAgent agent = new(chatClient, options: new() + { + ChatOptions = new() { Tools = [tool] }, + RequirePerServiceCallChatHistoryPersistence = true, + PersistPartialMessagesOnCancellation = true, + }); + var session = await agent.CreateSessionAsync(); + + var approvalRequest = Assert.Single(await RunAndCollectApprovalRequestsAsync(agent, session, new ChatMessage(ChatRole.User, "What's the weather?"))); + + // Act + await RunUntilCancelledAsync(agent, session, new ChatMessage(ChatRole.User, [approvalRequest.CreateResponse(approved: true)]), cts.Token); + await RunToCompletionAsync(agent, session, new ChatMessage(ChatRole.User, "Thanks")); + + // Assert + Assert.Equal(1, executions); + AssertSingleResultPerCall(GetHistory(agent, session), "call1"); + AssertSingleResultPerCall(chatClient.Requests[2], "call1"); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_ApprovedToolThenCancellationInLaterServiceCall_NextTurnDoesNotSettleCallAgainAsync() + { + // Arrange + int executions = 0; + var weatherTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => { executions++; return "Sunny"; }, "GetWeather")); + var timeTool = AIFunctionFactory.Create(() => "12:00", "GetTime"); + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call2", "GetTime")]) { MessageId = "m2" }]), + new StreamingCall([TextUpdate("It is sun", "m3")], CallEnd.CallerCancels), + new StreamingCall([TextUpdate("You're welcome", "m4")])); + ChatClientAgent agent = new(chatClient, options: new() + { + ChatOptions = new() { Tools = [weatherTool, timeTool] }, + RequirePerServiceCallChatHistoryPersistence = true, + PersistPartialMessagesOnCancellation = true, + }); + var session = await agent.CreateSessionAsync(); + + var approvalRequest = Assert.Single(await RunAndCollectApprovalRequestsAsync(agent, session, new ChatMessage(ChatRole.User, "What's the weather and the time?"))); + + // Act + await RunUntilCancelledAsync(agent, session, new ChatMessage(ChatRole.User, [approvalRequest.CreateResponse(approved: true)]), cts.Token); + await RunToCompletionAsync(agent, session, new ChatMessage(ChatRole.User, "Thanks")); + + // Assert + Assert.Equal(1, executions); + AssertSingleResultPerCall(GetHistory(agent, session), "call1", "call2"); + AssertSingleResultPerCall(chatClient.Requests[3], "call1", "call2"); + } + + [Fact] + public async Task RunStreamingAsync_FlagEnabled_PerServiceCallPersistence_AutoApprovedToolThenCancellation_NextTurnDoesNotExecuteAgainAsync() + { + // Arrange + int bookExecutions = 0; + int weatherExecutions = 0; + var bookTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => { bookExecutions++; return "Booked"; }, "BookTrip")); + var weatherTool = AIFunctionFactory.Create(() => { weatherExecutions++; return "Sunny"; }, "GetWeather"); + using CancellationTokenSource cts = new(); + var chatClient = new ScriptedStreamingChatClient( + cts, + new StreamingCall([new ChatResponseUpdate(ChatRole.Assistant, [new FunctionCallContent("call1", "BookTrip"), new FunctionCallContent("call2", "GetWeather")]) { MessageId = "m1" }]), + new StreamingCall([TextUpdate("Your trip is bo", "m2")], CallEnd.CallerCancels), + new StreamingCall([TextUpdate("You're welcome", "m3")])); + ChatClientAgent agent = new(chatClient, options: new() + { + ChatOptions = new() { Tools = [bookTool, weatherTool] }, + RequirePerServiceCallChatHistoryPersistence = true, + PersistPartialMessagesOnCancellation = true, + }); + var session = await agent.CreateSessionAsync(); + + var approvalRequest = Assert.Single(await RunAndCollectApprovalRequestsAsync(agent, session, new ChatMessage(ChatRole.User, "Book a trip and check the weather"))); + Assert.Equal("call1", approvalRequest.ToolCall.CallId); + + // Act + await RunUntilCancelledAsync(agent, session, new ChatMessage(ChatRole.User, [approvalRequest.CreateResponse(approved: true)]), cts.Token); + await RunToCompletionAsync(agent, session, new ChatMessage(ChatRole.User, "Thanks")); + + // Assert + Assert.Equal(1, bookExecutions); + Assert.Equal(1, weatherExecutions); + AssertSingleResultPerCall(GetHistory(agent, session), "call1", "call2"); + AssertSingleResultPerCall(chatClient.Requests[2], "call1", "call2"); + } + [Fact] public async Task RunStreamingAsync_FlagEnabled_NextRunIncludesPersistedPartialInteractionAsync() { @@ -408,6 +510,41 @@ private static void AssertToolLoopPartialHistory(List history) return (historyAfterCancellation, GetHistory(agent, session)); } + private static async Task> RunAndCollectApprovalRequestsAsync(ChatClientAgent agent, AgentSession session, ChatMessage message) + { + List requests = []; + await foreach (var update in agent.RunStreamingAsync([message], session)) + { + requests.AddRange(update.Contents.OfType()); + } + + return requests; + } + + private static async Task RunUntilCancelledAsync(ChatClientAgent agent, AgentSession session, ChatMessage message, CancellationToken cancellationToken) => + await Assert.ThrowsAnyAsync(async () => + { + await foreach (var _ in agent.RunStreamingAsync([message], session, cancellationToken: cancellationToken)) + { + } + }); + + private static async Task RunToCompletionAsync(ChatClientAgent agent, AgentSession session, ChatMessage message) + { + await foreach (var _ in agent.RunStreamingAsync([message], session)) + { + } + } + + private static void AssertSingleResultPerCall(IEnumerable messages, params string[] callIds) + { + var resultCallIds = messages.SelectMany(m => m.Contents).OfType().Select(r => r.CallId).ToList(); + foreach (var callId in callIds) + { + Assert.Single(resultCallIds, id => string.Equals(id, callId, StringComparison.Ordinal)); + } + } + private static async Task> ConsumeUntilCancelledAsync(ChatClientAgent agent, AgentSession session, AgentRunOptions? options, CancellationToken cancellationToken) { List received = []; From 31e70d7a53af3836a9f54225e45ae0284e53e9e4 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Wed, 28 Oct 2026 19:09:53 +0100 Subject: [PATCH 5/5] Update dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs Co-authored-by: Vincent Biret --- .../Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 32a3edf17f5..57a24c03941 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -1091,17 +1091,9 @@ internal static List CreatePartialResponseMessages(IEnumerable responseMessages = responseUpdates.ToChatResponse().Messages; - HashSet toolCallIds = new(StringComparer.Ordinal); + HashSet toolCallIds = new(requestMessages.SelectMany(m => m.Contents).OfType().Select(x => x.CallId), StringComparer.Ordinal); HashSet toolResultIds = new(StringComparer.Ordinal); - foreach (var content in requestMessages.SelectMany(m => m.Contents)) - { - if (content is ToolCallContent toolCall) - { - toolCallIds.Add(toolCall.CallId); - } - } - foreach (var content in responseMessages.SelectMany(m => m.Contents)) { switch (content)