.NET: Allow persisting partial interaction when streaming is cancelled

dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalNotRequiredFunctionBypassingChatClient.cs

@@ -159,6 +159,42 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
         }
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// 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
+    /// <see cref="ChatClientAgentOptions.PersistPartialMessagesOnCancellation"/>), 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 <see cref="FunctionInvokingChatClient"/>, so a
+    /// leftover call would be injected as approved and executed a second time.
+    /// </remarks>
+    /// <param name="session">The session holding the stored auto-approved function calls.</param>
+    /// <param name="settledCallIds">The call ids whose results have been persisted.</param>
+    internal static void RemoveSettledAutoApprovals(AgentSession session, HashSet<string> settledCallIds)
+    {
+        if (!session.StateBag.TryGetValue(
+            StateBagKey,
+            out List<ToolApprovalRequestContent>? 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);
+        }
+    }
+
     /// <summary>
     /// Attempts to get the current <see cref="AgentSession"/> from the ambient run context. When no run
     /// context or session is available, logs a warning (once per instance) and returns <see langword="false"/>

dotnet/src/Microsoft.Agents.AI/ChatClient/ApprovalResponseBindingChatClient.cs

@@ -611,6 +611,29 @@ private void MergePendingApprovalRequests(List<ToolApprovalRequestContent> emitt
         }
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// 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
+    /// <see cref="ChatClientAgentOptions.PersistPartialMessagesOnCancellation"/>), 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.
+    /// </remarks>
+    /// <param name="session">The session holding the recorded pending approval requests.</param>
+    /// <param name="settledCallIds">The call ids whose results have been persisted.</param>
+    internal static void RemoveSettledPendingApprovalRequests(AgentSession session, HashSet<string> settledCallIds)
+    {
+        var pendingRequests = LoadPendingApprovalRequests(session);
+        if (pendingRequests.RemoveAll(request => settledCallIds.Contains(request.ToolCall.CallId)) > 0)
+        {
+            SavePendingApprovalRequests(pendingRequests, session);
+        }
+    }
+
     private static List<ToolApprovalRequestContent> LoadPendingApprovalRequests(AgentSession session)
         => session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(StateBagKey, out var pendingRequests, AgentJsonUtilities.DefaultOptions)
             && pendingRequests is not null

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;
 /// </remarks>
 public sealed partial class ChatClientAgent : AIAgent
 {
+    /// <summary>
+    /// The key of the <see cref="ChatMessage.AdditionalProperties"/> entry that marks a response message persisted
+    /// from a streaming run that was cancelled before completion.
+    /// </summary>
+    /// <remarks>
+    /// This marker is only added when <see cref="ChatClientAgentOptions.PersistPartialMessagesOnCancellation"/> or
+    /// <see cref="ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation"/> is enabled. Its value is <see langword="true"/>.
+    /// </remarks>
+    [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
+    public const string PartialResponseMarkerKey = "agent_framework.partial_response";
+
     private readonly ChatClientAgentOptions? _agentOptions;
     private readonly HashSet<string> _aiContextProviderStateKeys;
     private readonly AIAgentMetadata _agentMetadata;
@@ -369,6 +382,11 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> 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<AgentResponseUpdate> 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,174 @@ private Task NotifyProvidersOfFailureAtEndOfRunAsync(
         return this.NotifyProvidersOfFailureAsync(session, ex, requestMessages, chatOptions, cancellationToken);
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// This method is evaluated in an exception filter, so it must not have side effects.
+    /// It returns <see langword="false"/> when per-service-call persistence is active (the
+    /// <see cref="PerServiceCallChatHistoryPersistingChatClient"/> handles the notification), when chat history is
+    /// 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.
+    /// </remarks>
+    private bool ShouldNotifyProvidersOfPartialMessagesAtEndOfRun(
+        Exception ex,
+        ChatClientAgentSession session,
+        ChatOptions? chatOptions,
+        AgentRunOptions? runOptions,
+        ChatClientAgentContinuationToken? continuationToken,
+        List<ChatResponseUpdate> 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) || update.ContinuationToken is not null)
+        && this.ShouldNotifyProvidersOfPartialMessagesOnCancellation(ex, runOptions, cancellationToken);
+
+    /// <summary>
+    /// Determines whether providers should be notified of the partial messages received before a cancellation.
+    /// </summary>
+    /// <remarks>
+    /// Returns <see langword="true"/> only when <paramref name="ex"/> is an <see cref="OperationCanceledException"/> caused by
+    /// the caller's <paramref name="cancellationToken"/>, and partial persistence on cancellation is enabled either per run via
+    /// <see cref="ChatClientAgentRunOptions.PersistPartialMessagesOnCancellation"/> or at agent level via
+    /// <see cref="ChatClientAgentOptions.PersistPartialMessagesOnCancellation"/>.
+    /// This method is evaluated in exception filters, so it must not have side effects.
+    /// It is also called by <see cref="PerServiceCallChatHistoryPersistingChatClient"/>.
+    /// </remarks>
+    internal bool ShouldNotifyProvidersOfPartialMessagesOnCancellation(Exception ex, AgentRunOptions? runOptions, CancellationToken cancellationToken) =>
+        ex is OperationCanceledException
+        && cancellationToken.IsCancellationRequested
+        && ((runOptions as ChatClientAgentRunOptions)?.PersistPartialMessagesOnCancellation
+            ?? this._agentOptions?.PersistPartialMessagesOnCancellation
+            ?? false);
+
+    /// <summary>
+    /// Notifies the <see cref="ChatHistoryProvider"/> and all <see cref="AIContextProviders"/> of the request messages and the
+    /// sanitized partial response of a cancelled streaming run, through the same channel used for successfully completed messages.
+    /// </summary>
+    /// <remarks>
+    /// The notification is best effort: it uses <see cref="CancellationToken.None"/> because the caller's token is already
+    /// cancelled, and any failure is logged and swallowed so that the original <see cref="OperationCanceledException"/>
+    /// is not masked. This method is also called by <see cref="PerServiceCallChatHistoryPersistingChatClient"/>.
+    /// 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.
+    /// </remarks>
+    /// <param name="session">The session of the cancelled run.</param>
+    /// <param name="requestMessages">The request messages to persist.</param>
+    /// <param name="responseUpdates">The response updates received before the cancellation.</param>
+    /// <param name="chatOptions">The chat options of the cancelled run.</param>
+    /// <param name="historyMessages">
+    /// The messages known to be stored in the chat history, including <paramref name="requestMessages"/>, used to find
+    /// the tool calls that are settled. When <see langword="null"/>, only <paramref name="requestMessages"/> are used.
+    /// </param>
+    internal async Task NotifyProvidersOfPartialMessagesOnCancellationAsync(
+        ChatClientAgentSession session,
+        IEnumerable<ChatMessage> requestMessages,
+        IEnumerable<ChatResponseUpdate> responseUpdates,
+        ChatOptions? chatOptions,
+        IEnumerable<ChatMessage>? historyMessages = null)
+    {
+        try
+        {
+            var requestMessagesList = requestMessages as IReadOnlyCollection<ChatMessage> ?? requestMessages.ToList();
+            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)
+        {
+            this._logger.LogAgentChatClientPartialMessagesNotificationFailed(this.Id, this.GetLoggingAgentName(), notifyException);
+        }
+    }
+
+    /// <summary>
+    /// 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.
+    /// </summary>
+    /// <remarks>
+    /// <see cref="ApprovalResponseBindingChatClient"/> and <see cref="ApprovalNotRequiredFunctionBypassingChatClient"/>
+    /// 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.
+    /// </remarks>
+    private static void RemoveSettledToolApprovalState(ChatClientAgentSession session, IEnumerable<ChatMessage> persistedMessages)
+    {
+        HashSet<string> settledCallIds = new(
+            persistedMessages.SelectMany(message => message.Contents).OfType<ToolResultContent>().Select(result => result.CallId),
+            StringComparer.Ordinal);
+
+        if (settledCallIds.Count == 0)
+        {
+            return;
+        }
+
+        ApprovalResponseBindingChatClient.RemoveSettledPendingApprovalRequests(session, settledCallIds);
+        ApprovalNotRequiredFunctionBypassingChatClient.RemoveSettledAutoApprovals(session, settledCallIds);
+    }
+
+    /// <summary>
+    /// Builds the response messages to persist for a cancelled streaming run from the updates received so far.
+    /// </summary>
+    /// <remarks>
+    /// 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, <see cref="ToolApprovalRequestContent"/>), and empty text.
+    /// Messages left without content are dropped, and every remaining message is marked with
+    /// <see cref="PartialResponseMarkerKey"/>.
+    /// </remarks>
+    internal static List<ChatMessage> CreatePartialResponseMessages(IEnumerable<ChatMessage> requestMessages, IEnumerable<ChatResponseUpdate> responseUpdates)
+    {
+        IList<ChatMessage> responseMessages = responseUpdates.ToChatResponse().Messages;
+
+        HashSet<string> toolCallIds = new(requestMessages.SelectMany(m => m.Contents).OfType<ToolCallContent>().Select(x => x.CallId), StringComparer.Ordinal);
+        HashSet<string> toolResultIds = new(StringComparer.Ordinal);
+
+        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<ChatMessage> partialMessages = [];
+        foreach (var message in responseMessages)
+        {
+            List<AIContent> 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;
+    }
+
     /// <summary>
     /// Gets a value indicating whether the agent is configured to simulate service-stored chat history.
     /// When <see langword="true"/>, end-of-run persistence and history loading are skipped because a

dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentLogMessages.cs

@@ -97,4 +97,16 @@ public static partial void LogAgentChatClientBackgroundResponseFallback(
         this ILogger logger,
         string agentId,
         string agentName);
+
+    /// <summary>
+    /// Logs a warning when notifying providers of the partial messages of a cancelled streaming run fails.
+    /// </summary>
+    [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);
 }