.NET: fix: Add session support for Handoff-hosted Agents

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/HandoffEndExecutor.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.AllMessages, cancellationToken).ConfigureAwait(false);
+
+                return sharedState;
+            }, context, cancellationToken).ConfigureAwait(false);
     }
 
     public ValueTask ResetAsync() => default;

dotnet/src/Microsoft.Agents.AI.Workflows/Specialized/HandoffMessagesFilter.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/HandoffStartExecutor.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/HandoffState.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/MultiPartyConversation.cs

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+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 IReadOnlyList<ChatMessage> AllMessages => this._history;
+
+    public (ChatMessage[], int) CollectNewMessages(int bookmark)
+    {
+        lock (this._mutex)
+        {
+            int count = this._history.Count - bookmark;
+            if (count > 0)
+            {
+                return (this._history.Skip(bookmark).ToArray(), this.CurrentBookmark);
+            }
+
+            return ([], 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;
+        }
+    }
+}