.NET: fix: Add session support for Handoff-hosted Agents (#5280)

* fix: Add session support for Handoff-hosted Agents

In order to better support using `Workflows` hosted as `AIAgents` inside of Handoff workflows, we need to make proper use of AgentSession. This causes potential issues around checkpointing and making sure that we properly compute only the new incoming messages for each agent invocation.

* fix: AgentSession checkpointing using AIAgent's Serialize/Deserialize methods

We cannot rely on implicit serialization through `HandoffHostState` because we are missing type information.

* fix: Thread safety issue in `MultiPartyConversation.AllMessages`

* fix: Enable unwrapping of FunctionResultContent when ExternalRequest was wrapped into FunctionCallContent

Author: lokitoth

dotnet/src/Microsoft.Agents.AI.Workflows/AIAgentsAbstractionsExtensions.cs

@@ -48,7 +48,7 @@ private static ChatMessage ChatAssistantToUserIfNotFromNamed(this ChatMessage me
     /// any that have a different <see cref="ChatMessage.AuthorName"/> from <paramref name="targetAgentName"/> to
     /// <see cref="ChatRole.User"/>.
     /// </summary>
-    public static List<ChatMessage>? ChangeAssistantToUserForOtherParticipants(this List<ChatMessage> messages, string targetAgentName)
+    public static List<ChatMessage>? ChangeAssistantToUserForOtherParticipants(this IEnumerable<ChatMessage> messages, string targetAgentName)
     {
         List<ChatMessage>? roleChanged = null;
         foreach (var m in messages)

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffAgentExecutor.cs

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft. All rights reserved.
 
+using System;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -12,23 +13,33 @@ internal sealed class HandoffEndExecutor(bool returnToPrevious) : Executor(Execu
 {
     public const string ExecutorId = "HandoffEnd";
 
+    private readonly StateRef<HandoffSharedState> _sharedStateRef = new(HandoffConstants.HandoffSharedStateKey,
+                                                                        HandoffConstants.HandoffSharedStateScope);
+
     protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder) =>
-        protocolBuilder.ConfigureRoutes(routeBuilder => routeBuilder.AddHandler<HandoffState>((handoff, context, cancellationToken) =>
-                                            this.HandleAsync(handoff, context, cancellationToken)))
+        protocolBuilder.ConfigureRoutes(routeBuilder => routeBuilder.AddHandler<HandoffState>(
+                                            (handoff, context, cancellationToken) => this.HandleAsync(handoff, context, cancellationToken)))
                        .YieldsOutput<List<ChatMessage>>();
 
     private async ValueTask HandleAsync(HandoffState handoff, IWorkflowContext context, CancellationToken cancellationToken)
     {
-        if (returnToPrevious)
-        {
-            await context.QueueStateUpdateAsync<string?>(HandoffConstants.PreviousAgentTrackerKey,
-                                                         handoff.PreviousAgentId,
-                                                         HandoffConstants.PreviousAgentTrackerScope,
-                                                         cancellationToken)
-                         .ConfigureAwait(false);
-        }
-
-        await context.YieldOutputAsync(handoff.Messages, cancellationToken).ConfigureAwait(false);
+        await this._sharedStateRef.InvokeWithStateAsync(
+            async (HandoffSharedState? sharedState, IWorkflowContext context, CancellationToken cancellationToken) =>
+            {
+                if (sharedState == null)
+                {
+                    throw new InvalidOperationException("Handoff Orchestration shared state was not properly initialized.");
+                }
+
+                if (returnToPrevious)
+                {
+                    sharedState.PreviousAgentId = handoff.PreviousAgentId;
+                }
+
+                await context.YieldOutputAsync(sharedState.Conversation.CloneAllMessages(), cancellationToken).ConfigureAwait(false);
+
+                return sharedState;
+            }, context, cancellationToken).ConfigureAwait(false);
     }
 
     public ValueTask ResetAsync() => default;

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffEndExecutor.cs

@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI.Workflows.Specialized;
+
+[Experimental(DiagnosticConstants.ExperimentalFeatureDiagnostic)]
+internal sealed class HandoffMessagesFilter
+{
+    private readonly HandoffToolCallFilteringBehavior _filteringBehavior;
+
+    public HandoffMessagesFilter(HandoffToolCallFilteringBehavior filteringBehavior)
+    {
+        this._filteringBehavior = filteringBehavior;
+    }
+
+    [Experimental(DiagnosticConstants.ExperimentalFeatureDiagnostic)]
+    internal static bool IsHandoffFunctionName(string name)
+    {
+        return name.StartsWith(HandoffWorkflowBuilder.FunctionPrefix, StringComparison.Ordinal);
+    }
+
+    public IEnumerable<ChatMessage> FilterMessages(IEnumerable<ChatMessage> messages)
+    {
+        if (this._filteringBehavior == HandoffToolCallFilteringBehavior.None)
+        {
+            return messages;
+        }
+
+        Dictionary<string, FilterCandidateState> filteringCandidates = new();
+        List<ChatMessage> filteredMessages = [];
+        HashSet<int> messagesToRemove = [];
+
+        bool filterHandoffOnly = this._filteringBehavior == HandoffToolCallFilteringBehavior.HandoffOnly;
+        foreach (ChatMessage unfilteredMessage in messages)
+        {
+            ChatMessage filteredMessage = unfilteredMessage.Clone();
+
+            // .Clone() is shallow, so we cannot modify the contents of the cloned message in place.
+            List<AIContent> contents = [];
+            contents.Capacity = unfilteredMessage.Contents?.Count ?? 0;
+            filteredMessage.Contents = contents;
+
+            // Because this runs after the role changes from assistant to user for the target agent, we cannot rely on tool calls
+            // originating only from messages with the Assistant role. Instead, we need to inspect the contents of all non-Tool (result)
+            // FunctionCallContent.
+            if (unfilteredMessage.Role != ChatRole.Tool)
+            {
+                for (int i = 0; i < unfilteredMessage.Contents!.Count; i++)
+                {
+                    AIContent content = unfilteredMessage.Contents[i];
+                    if (content is not FunctionCallContent fcc || (filterHandoffOnly && !IsHandoffFunctionName(fcc.Name)))
+                    {
+                        filteredMessage.Contents.Add(content);
+
+                        // Track non-handoff function calls so their tool results are preserved in HandoffOnly mode
+                        if (filterHandoffOnly && content is FunctionCallContent nonHandoffFcc)
+                        {
+                            filteringCandidates[nonHandoffFcc.CallId] = new FilterCandidateState(nonHandoffFcc.CallId)
+                            {
+                                IsHandoffFunction = false,
+                            };
+                        }
+                    }
+                    else if (filterHandoffOnly)
+                    {
+                        if (!filteringCandidates.TryGetValue(fcc.CallId, out FilterCandidateState? candidateState))
+                        {
+                            filteringCandidates[fcc.CallId] = new FilterCandidateState(fcc.CallId)
+                            {
+                                IsHandoffFunction = true,
+                            };
+                        }
+                        else
+                        {
+                            candidateState.IsHandoffFunction = true;
+                            (int messageIndex, int contentIndex) = candidateState.FunctionCallResultLocation!.Value;
+                            ChatMessage messageToFilter = filteredMessages[messageIndex];
+                            messageToFilter.Contents.RemoveAt(contentIndex);
+                            if (messageToFilter.Contents.Count == 0)
+                            {
+                                messagesToRemove.Add(messageIndex);
+                            }
+                        }
+                    }
+                    else
+                    {
+                        // All mode: strip all FunctionCallContent
+                    }
+                }
+            }
+            else
+            {
+                if (!filterHandoffOnly)
+                {
+                    continue;
+                }
+
+                for (int i = 0; i < unfilteredMessage.Contents!.Count; i++)
+                {
+                    AIContent content = unfilteredMessage.Contents[i];
+                    if (content is not FunctionResultContent frc
+                        || (filteringCandidates.TryGetValue(frc.CallId, out FilterCandidateState? candidateState)
+                            && candidateState.IsHandoffFunction is false))
+                    {
+                        // Either this is not a function result content, so we should let it through, or it is a FRC that
+                        // we know is not related to a handoff call. In either case, we should include it.
+                        filteredMessage.Contents.Add(content);
+                    }
+                    else if (candidateState is null)
+                    {
+                        // We haven't seen the corresponding function call yet, so add it as a candidate to be filtered later
+                        filteringCandidates[frc.CallId] = new FilterCandidateState(frc.CallId)
+                        {
+                            FunctionCallResultLocation = (filteredMessages.Count, filteredMessage.Contents.Count),
+                        };
+                    }
+                    // else we have seen the corresponding function call and it is a handoff, so we should filter it out.
+                }
+            }
+
+            if (filteredMessage.Contents.Count > 0)
+            {
+                filteredMessages.Add(filteredMessage);
+            }
+        }
+
+        return filteredMessages.Where((_, index) => !messagesToRemove.Contains(index));
+    }
+
+    private class FilterCandidateState(string callId)
+    {
+        public (int MessageIndex, int ContentIndex)? FunctionCallResultLocation { get; set; }
+
+        public string CallId => callId;
+
+        public bool? IsHandoffFunction { get; set; }
+    }
+}

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffMessagesFilter.cs

@@ -9,8 +9,23 @@ namespace Microsoft.Agents.AI.Workflows.Specialized;
 
 internal static class HandoffConstants
 {
+    internal const string HandoffOrchestrationSharedScope = "HandoffOrchestration";
+
     internal const string PreviousAgentTrackerKey = "LastAgentId";
-    internal const string PreviousAgentTrackerScope = "HandoffOrchestration";
+    internal const string PreviousAgentTrackerScope = HandoffOrchestrationSharedScope;
+
+    internal const string MultiPartyConversationKey = "MultiPartyConversation";
+    internal const string MultiPartyConversationScope = HandoffOrchestrationSharedScope;
+
+    internal const string HandoffSharedStateKey = "SharedState";
+    internal const string HandoffSharedStateScope = HandoffOrchestrationSharedScope;
+}
+
+internal sealed class HandoffSharedState
+{
+    public MultiPartyConversation Conversation { get; } = new();
+
+    public string? PreviousAgentId { get; set; }
 }
 
 /// <summary>Executor used at the start of a handoffs workflow to accumulate messages and emit them as HandoffState upon receiving a turn token.</summary>
@@ -29,23 +44,25 @@ protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBui
 
     protected override ValueTask TakeTurnAsync(List<ChatMessage> messages, IWorkflowContext context, bool? emitEvents, CancellationToken cancellationToken = default)
     {
-        if (returnToPrevious)
-        {
-            return context.InvokeWithStateAsync(
-                async (string? previousAgentId, IWorkflowContext context, CancellationToken cancellationToken) =>
-                {
-                    HandoffState handoffState = new(new(emitEvents), null, messages, previousAgentId);
-                    await context.SendMessageAsync(handoffState, cancellationToken).ConfigureAwait(false);
-
-                    return previousAgentId;
-                },
-                HandoffConstants.PreviousAgentTrackerKey,
-                HandoffConstants.PreviousAgentTrackerScope,
-                cancellationToken);
-        }
-
-        HandoffState handoff = new(new(emitEvents), null, messages);
-        return context.SendMessageAsync(handoff, cancellationToken);
+        return context.InvokeWithStateAsync(
+            async (HandoffSharedState? sharedState, IWorkflowContext context, CancellationToken cancellationToken) =>
+            {
+                sharedState ??= new HandoffSharedState();
+                sharedState.Conversation.AddMessages(messages);
+
+                string? previousAgentId = sharedState.PreviousAgentId;
+
+                // If we are configured to return to the previous agent, include the previous agent id in the handoff state.
+                // If there was no previousAgent, it will still be null.
+                HandoffState turnState = new(new(emitEvents), null, returnToPrevious ? previousAgentId : null);
+
+                await context.SendMessageAsync(turnState, cancellationToken).ConfigureAwait(false);
+
+                return sharedState;
+            },
+            HandoffConstants.HandoffSharedStateKey,
+            HandoffConstants.HandoffSharedStateScope,
+            cancellationToken);
     }
 
     public new ValueTask ResetAsync() => base.ResetAsync();

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffStartExecutor.cs

@@ -1,12 +1,8 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-using System.Collections.Generic;
-using Microsoft.Extensions.AI;
-
 namespace Microsoft.Agents.AI.Workflows.Specialized;
 
 internal sealed record class HandoffState(
     TurnToken TurnToken,
     string? RequestedHandoffTargetAgentId,
-    List<ChatMessage> Messages,
     string? PreviousAgentId = null);

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffState.cs

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI.Workflows.Specialized;
+
+internal sealed class MultiPartyConversation
+{
+    private readonly List<ChatMessage> _history = [];
+    private readonly object _mutex = new();
+
+    public List<ChatMessage> CloneAllMessages()
+    {
+        lock (this._mutex)
+        {
+            return this._history.ToList();
+        }
+    }
+
+    public (ChatMessage[], int) CollectNewMessages(int bookmark)
+    {
+        lock (this._mutex)
+        {
+            int count = this._history.Count - bookmark;
+            if (count < 0)
+            {
+                throw new InvalidOperationException($"Bookmark value too large: {bookmark} vs count={count}");
+            }
+
+            return (this._history.Skip(bookmark).ToArray(), this.CurrentBookmark);
+        }
+    }
+
+    private int CurrentBookmark => this._history.Count;
+
+    public int AddMessages(IEnumerable<ChatMessage> messages)
+    {
+        lock (this._mutex)
+        {
+            this._history.AddRange(messages);
+            return this.CurrentBookmark;
+        }
+    }
+
+    public int AddMessage(ChatMessage message)
+    {
+        lock (this._mutex)
+        {
+            this._history.Add(message);
+            return this.CurrentBookmark;
+        }
+    }
+}

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/MultiPartyConversation.cs

@@ -41,7 +41,7 @@ internal static FunctionCallContent ToFunctionCall(this ExternalRequest request)
     {
         Dictionary<string, object?> parameters = new()
         {
-            { "data", request.Data}
+            { "data", request.Data }
         };
 
         return new FunctionCallContent(request.RequestId, request.PortInfo.PortId, parameters);