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.

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

@@ -48,7 +48,7 @@ internal static class AIAgentsAbstractionsExtensions
     /// 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

@@ -6,6 +6,7 @@ using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
@@ -31,140 +32,6 @@ internal sealed class HandoffAgentExecutorOptions
     public HandoffToolCallFilteringBehavior ToolCallFilteringBehavior { get; set; } = HandoffToolCallFilteringBehavior.HandoffOnly;
 }
 
-[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(List<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; }
-    }
-}
-
 internal struct AgentInvocationResult(AgentResponse agentResponse, string? handoffTargetId)
 {
     public AgentResponse Response => agentResponse;
@@ -175,19 +42,32 @@ internal struct AgentInvocationResult(AgentResponse agentResponse, string? hando
     public bool IsHandoffRequested => this.HandoffTargetId != null;
 }
 
-internal record HandoffAgentHostState(HandoffState? CurrentTurnState, List<ChatMessage> FilteredIncomingMessages, List<ChatMessage> TurnMessages)
+internal record HandoffAgentHostState(
+    HandoffState? IncomingState,
+    int ConversationBookmark,
+    AgentSession? AgentSession)
 {
-    public HandoffState PrepareHandoff(AgentInvocationResult invocationResult, string currentAgentId)
-    {
-        if (this.CurrentTurnState == null)
-        {
-            throw new InvalidOperationException("Cannot create a handoff request: Out of turn.");
-        }
+    [MemberNotNullWhen(true, nameof(IncomingState))]
+    [JsonIgnore]
+    public bool IsTakingTurn => this.IncomingState != null;
+}
 
-        IEnumerable<ChatMessage> allMessages = [.. this.CurrentTurnState.Messages, .. this.TurnMessages, .. invocationResult.Response.Messages];
+internal sealed record StateRef<TState>(string Key, string? ScopeName)
+{
+    public ValueTask InvokeWithStateAsync(Func<TState?, IWorkflowContext, CancellationToken, ValueTask<TState?>> invocation,
+                                          IWorkflowContext context,
+                                          CancellationToken cancellationToken)
+        => context.InvokeWithStateAsync(invocation, this.Key, this.ScopeName, cancellationToken);
 
-        return new(this.CurrentTurnState.TurnToken, invocationResult.HandoffTargetId, allMessages.ToList(), currentAgentId);
-    }
+    public ValueTask InvokeWithStateAsync(Func<TState?, IWorkflowContext, CancellationToken, ValueTask> invocation,
+                                                 IWorkflowContext context,
+                                                 CancellationToken cancellationToken)
+        => context.InvokeWithStateAsync<TState>(
+              async (state, ctx, ct) =>
+              {
+                  await invocation(state, ctx, ct).ConfigureAwait(false);
+                  return state;
+              }, this.Key, this.ScopeName, cancellationToken);
 }
 
 /// <summary>Executor used to represent an agent in a handoffs workflow, responding to <see cref="HandoffState"/> events.</summary>
@@ -208,7 +88,10 @@ internal sealed class HandoffAgentExecutor :
     private readonly HashSet<string> _handoffFunctionNames = [];
     private readonly Dictionary<string, string> _handoffFunctionToAgentId = [];
 
-    private static HandoffAgentHostState InitialStateFactory() => new(null, [], []);
+    private readonly StateRef<HandoffSharedState> _sharedStateRef = new(HandoffConstants.HandoffSharedStateKey,
+                                                                        HandoffConstants.HandoffSharedStateScope);
+
+    private static HandoffAgentHostState InitialStateFactory() => new(null, 0, null);
 
     public HandoffAgentExecutor(AIAgent agent, HashSet<HandoffTarget> handoffs, HandoffAgentExecutorOptions options)
         : base(IdFor(agent), InitialStateFactory)
@@ -291,13 +174,18 @@ internal sealed class HandoffAgentExecutor :
         // resumes can be processed in one invocation.
         return this.InvokeWithStateAsync((state, ctx, ct) =>
         {
-            state.TurnMessages.Add(new ChatMessage(ChatRole.User, [response])
+            if (!state.IsTakingTurn)
+            {
+                throw new InvalidOperationException("Cannot process user responses when not taking a turn in Handoff Orchestration.");
+            }
+
+            ChatMessage userMessage = new(ChatRole.User, [response])
             {
                 CreatedAt = DateTimeOffset.UtcNow,
                 MessageId = Guid.NewGuid().ToString("N"),
-            });
+            };
 
-            return this.ContinueTurnAsync(state, ctx, ct);
+            return this.ContinueTurnAsync(state, [userMessage], ctx, ct);
         }, context, skipCache: false, cancellationToken);
     }
 
@@ -315,24 +203,44 @@ internal sealed class HandoffAgentExecutor :
         // resumes can be processed in one invocation.
         return this.InvokeWithStateAsync((state, ctx, ct) =>
         {
-            state.TurnMessages.Add(
-                new ChatMessage(ChatRole.Tool, [result])
-                {
-                    AuthorName = this._agent.Name ?? this._agent.Id,
-                    CreatedAt = DateTimeOffset.UtcNow,
-                    MessageId = Guid.NewGuid().ToString("N"),
-                });
+            if (!state.IsTakingTurn)
+            {
+                throw new InvalidOperationException("Cannot process user responses in when not taking a turn in Handoff Orchestration.");
+            }
 
-            return this.ContinueTurnAsync(state, ctx, ct);
+            ChatMessage toolMessage = new(ChatRole.Tool, [result])
+            {
+                AuthorName = this._agent.Name ?? this._agent.Id,
+                CreatedAt = DateTimeOffset.UtcNow,
+                MessageId = Guid.NewGuid().ToString("N"),
+            };
+
+            return this.ContinueTurnAsync(state, [toolMessage], ctx, ct);
         }, context, skipCache: false, cancellationToken);
     }
 
-    private async ValueTask<HandoffAgentHostState?> ContinueTurnAsync(HandoffAgentHostState state, IWorkflowContext context, CancellationToken cancellationToken)
+    private async ValueTask<HandoffAgentHostState?> ContinueTurnAsync(HandoffAgentHostState state, List<ChatMessage> incomingMessages, IWorkflowContext context, CancellationToken cancellationToken, bool skipAddIncoming = false)
     {
-        List<ChatMessage>? roleChanges = state.FilteredIncomingMessages.ChangeAssistantToUserForOtherParticipants(this._agent.Name ?? this._agent.Id);
+        if (!state.IsTakingTurn)
+        {
+            throw new InvalidOperationException("Cannot process user responses in when not taking a turn in Handoff Orchestration.");
+        }
 
-        bool emitUpdateEvents = state.CurrentTurnState!.ShouldEmitStreamingEvents(this._options.EmitAgentResponseUpdateEvents);
-        AgentInvocationResult result = await this.InvokeAgentAsync([.. state.FilteredIncomingMessages, .. state.TurnMessages], context, emitUpdateEvents, cancellationToken)
+        // If a handoff was invoked by a previous agent, filter out the handoff function call and tool result messages
+        // before sending to the underlying agent. These are internal workflow mechanics that confuse the target model
+        // into ignoring the original user question.
+        //
+        // This will not filter out tool responses and approval responses that are part of this agent's turn, which is
+        // the expected behavior since those are part of the agent's reasoning process.
+        HandoffMessagesFilter handoffMessagesFilter = new(this._options.ToolCallFilteringBehavior);
+        IEnumerable<ChatMessage> messagesForAgent = state.IncomingState.RequestedHandoffTargetAgentId is not null
+                                                  ? handoffMessagesFilter.FilterMessages(incomingMessages)
+                                                  : incomingMessages;
+
+        List<ChatMessage>? roleChanges = messagesForAgent.ChangeAssistantToUserForOtherParticipants(this._agent.Name ?? this._agent.Id);
+
+        bool emitUpdateEvents = state.IncomingState!.ShouldEmitStreamingEvents(this._options.EmitAgentResponseUpdateEvents);
+        AgentInvocationResult result = await this.InvokeAgentAsync(messagesForAgent, context, emitUpdateEvents, cancellationToken)
                                                      .ConfigureAwait(false);
 
         if (this.HasOutstandingRequests && result.IsHandoffRequested)
@@ -342,20 +250,40 @@ internal sealed class HandoffAgentExecutor :
 
         roleChanges.ResetUserToAssistantForChangedRoles();
 
+        int newConversationBookmark = state.ConversationBookmark;
+        await this._sharedStateRef.InvokeWithStateAsync(
+            (sharedState, ctx, ct) =>
+            {
+                if (sharedState == null)
+                {
+                    throw new InvalidOperationException("Handoff Orchestration shared state was not properly initialized.");
+                }
+
+                if (!skipAddIncoming)
+                {
+                    sharedState.Conversation.AddMessages(incomingMessages);
+                }
+
+                newConversationBookmark = sharedState.Conversation.AddMessages(result.Response.Messages);
+
+                return new ValueTask();
+            },
+            context,
+            cancellationToken).ConfigureAwait(false);
+
         // We send on the HandoffState even if handoff is not requested because we might be terminating the processing, but this only
         // happens if we have no outstanding requests.
         if (!this.HasOutstandingRequests)
         {
-            HandoffState outgoingState = state.PrepareHandoff(result, this._agent.Id);
+            HandoffState outgoingState = new(state.IncomingState.TurnToken, result.HandoffTargetId, this._agent.Id);
 
             await context.SendMessageAsync(outgoingState, cancellationToken).ConfigureAwait(false);
 
-            // reset the state for the next handoff (return-to-current is modeled as a new handoff turn, as opposed to "HITL", which
-            // can be a bit confusing.)
-            return null;
+            // reset the state for the next handoff, making sure to keep track of the conversation bookmark, and avoid resetting the
+            // agent session. (return-to-current is modeled as a new handoff turn, as opposed to "HITL", which can be a bit confusing.)
+            return state with { IncomingState = null, ConversationBookmark = newConversationBookmark };
         }
 
-        state.TurnMessages.AddRange(result.Response.Messages);
         return state;
     }
 
@@ -363,28 +291,36 @@ internal sealed class HandoffAgentExecutor :
     {
         return this.InvokeWithStateAsync(InvokeContinueTurnAsync, context, skipCache: false, cancellationToken);
 
-        ValueTask<HandoffAgentHostState?> InvokeContinueTurnAsync(HandoffAgentHostState state, IWorkflowContext context, CancellationToken cancellationToken)
+        async ValueTask<HandoffAgentHostState?> InvokeContinueTurnAsync(HandoffAgentHostState state, IWorkflowContext context, CancellationToken cancellationToken)
         {
             // Check that we are not getting this message while in the middle of a turn
-            if (state.CurrentTurnState != null)
+            if (state.IsTakingTurn)
             {
                 throw new InvalidOperationException("Cannot have multiple simultaneous conversations in Handoff Orchestration.");
             }
 
-            // If a handoff was invoked by a previous agent, filter out the handoff function
-            // call and tool result messages before sending to the underlying agent. These
-            // are internal workflow mechanics that confuse the target model into ignoring the
-            // original user question.
-            HandoffMessagesFilter handoffMessagesFilter = new(this._options.ToolCallFilteringBehavior);
-            IEnumerable<ChatMessage> messagesForAgent = message.RequestedHandoffTargetAgentId is not null
-                ? handoffMessagesFilter.FilterMessages(message.Messages)
-                : message.Messages;
+            IEnumerable<ChatMessage> newConversationMessages = [];
+            int newConversationBookmark = 0;
 
-            // This works because the runtime guarantees that a given executor instance will process messages serially,
-            // though there is no global cross-executor ordering guarantee (and in turn, no canonical message delivery order)
-            state = new(message, messagesForAgent.ToList(), []);
+            await this._sharedStateRef.InvokeWithStateAsync(
+                (sharedState, ctx, ct) =>
+                {
+                    if (sharedState == null)
+                    {
+                        throw new InvalidOperationException("Handoff Orchestration shared state was not properly initialized.");
+                    }
 
-            return this.ContinueTurnAsync(state, context, cancellationToken);
+                    (newConversationMessages, newConversationBookmark) = sharedState.Conversation.CollectNewMessages(state.ConversationBookmark);
+
+                    return new ValueTask();
+                },
+                context,
+                cancellationToken).ConfigureAwait(false);
+
+            state = state with { IncomingState = message, ConversationBookmark = newConversationBookmark };
+
+            return await this.ContinueTurnAsync(state, newConversationMessages.ToList(), context, cancellationToken, skipAddIncoming: true)
+                             .ConfigureAwait(false);
         }
     }
 
@@ -417,31 +353,46 @@ internal sealed class HandoffAgentExecutor :
 
         AIAgentUnservicedRequestsCollector collector = new(this._userInputHandler, this._functionCallHandler);
 
-        IAsyncEnumerable<AgentResponseUpdate> agentStream = this._agent.RunStreamingAsync(
-                messages,
-                options: this._agentOptions,
-                cancellationToken: cancellationToken);
-
         string? requestedHandoff = null;
         List<AgentResponseUpdate> updates = [];
         List<FunctionCallContent> candidateRequests = [];
-        await foreach (AgentResponseUpdate update in agentStream.ConfigureAwait(false))
-        {
-            await AddUpdateAsync(update, cancellationToken).ConfigureAwait(false);
 
-            collector.ProcessAgentResponseUpdate(update, CollectHandoffRequestsFilter);
-
-            bool CollectHandoffRequestsFilter(FunctionCallContent candidateHandoffRequest)
+        await this.InvokeWithStateAsync(
+            async (state, ctx, ct) =>
             {
-                bool isHandoffRequest = this._handoffFunctionNames.Contains(candidateHandoffRequest.Name);
-                if (isHandoffRequest)
+                if (state.AgentSession == null)
                 {
-                    candidateRequests.Add(candidateHandoffRequest);
+                    state = state with { AgentSession = await this._agent.CreateSessionAsync(ct).ConfigureAwait(false) };
                 }
 
-                return !isHandoffRequest;
-            }
-        }
+                IAsyncEnumerable<AgentResponseUpdate> agentStream =
+                    this._agent.RunStreamingAsync(messages,
+                                                  state.AgentSession,
+                                                  options: this._agentOptions,
+                                                  cancellationToken: ct);
+
+                await foreach (AgentResponseUpdate update in agentStream.ConfigureAwait(false))
+                {
+                    await AddUpdateAsync(update, ct).ConfigureAwait(false);
+
+                    collector.ProcessAgentResponseUpdate(update, CollectHandoffRequestsFilter);
+
+                    bool CollectHandoffRequestsFilter(FunctionCallContent candidateHandoffRequest)
+                    {
+                        bool isHandoffRequest = this._handoffFunctionNames.Contains(candidateHandoffRequest.Name);
+                        if (isHandoffRequest)
+                        {
+                            candidateRequests.Add(candidateHandoffRequest);
+                        }
+
+                        return !isHandoffRequest;
+                    }
+                }
+
+                return state;
+            },
+            context,
+            cancellationToken: cancellationToken).ConfigureAwait(false);
 
         if (candidateRequests.Count > 1)
         {

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 this._sharedStateRef.InvokeWithStateAsync(
+            async (HandoffSharedState? sharedState, IWorkflowContext context, CancellationToken cancellationToken) =>
+            {
+                if (sharedState == null)
+                {
+                    throw new InvalidOperationException("Handoff Orchestration shared state was not properly initialized.");
+                }
 
-        await context.YieldOutputAsync(handoff.Messages, cancellationToken).ConfigureAwait(false);
+                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/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 @@ internal sealed class HandoffStartExecutor(bool returnToPrevious) : ChatProtocol
 
     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 context.InvokeWithStateAsync(
+            async (HandoffSharedState? sharedState, IWorkflowContext context, CancellationToken cancellationToken) =>
+            {
+                sharedState ??= new HandoffSharedState();
+                sharedState.Conversation.AddMessages(messages);
 
-                    return previousAgentId;
-                },
-                HandoffConstants.PreviousAgentTrackerKey,
-                HandoffConstants.PreviousAgentTrackerScope,
-                cancellationToken);
-        }
+                string? previousAgentId = sharedState.PreviousAgentId;
 
-        HandoffState handoff = new(new(emitEvents), null, messages);
-        return context.SendMessageAsync(handoff, cancellationToken);
+                // 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 List<ChatMessage> CloneAllMessages() => this._history.ToList();
+
+    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;
+        }
+    }
+}

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/HandoffAgentExecutorTests.cs

@@ -7,6 +7,10 @@ using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.Checkpointing;
+using Microsoft.Agents.AI.Workflows.Execution;
+using Microsoft.Agents.AI.Workflows.InProc;
+using Microsoft.Agents.AI.Workflows.Sample;
 using Microsoft.Agents.AI.Workflows.Specialized;
 using Microsoft.Extensions.AI;
 
@@ -14,6 +18,27 @@ namespace Microsoft.Agents.AI.Workflows.UnitTests;
 
 public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
 {
+    private static async ValueTask<TestRunContext> PrepareHandoffSharedStateAsync(TestRunContext? runContext = null, IEnumerable<ChatMessage>? messages = null)
+    {
+        runContext ??= new();
+
+        HandoffSharedState sharedState = new();
+
+        if (messages != null)
+        {
+            sharedState.Conversation.AddMessages(messages);
+        }
+
+        await runContext.BindWorkflowContext(nameof(HandoffStartExecutor))
+                        .QueueStateUpdateAsync(HandoffConstants.HandoffSharedStateKey,
+                                               sharedState,
+                                               HandoffConstants.HandoffSharedStateScope);
+
+        await runContext.StateManager.PublishUpdatesAsync(null);
+
+        return runContext;
+    }
+
     [Theory]
     [InlineData(null, null)]
     [InlineData(null, true)]
@@ -27,7 +52,7 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
     public async Task Test_HandoffAgentExecutor_EmitsStreamingUpdatesIFFConfiguredAsync(bool? executorSetting, bool? turnSetting)
     {
         // Arrange
-        TestRunContext testContext = new();
+        TestRunContext testContext = await PrepareHandoffSharedStateAsync();
         TestReplayAgent agent = new(TestMessages, TestAgentId, TestAgentName);
 
         HandoffAgentExecutorOptions options = new("",
@@ -39,7 +64,7 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
         testContext.ConfigureExecutor(executor);
 
         // Act
-        HandoffState message = new(new(turnSetting), null, []);
+        HandoffState message = new(new(turnSetting), null, null);
         await executor.HandleAsync(message, testContext.BindWorkflowContext(executor.Id));
 
         // Assert
@@ -55,7 +80,7 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
     public async Task Test_HandoffAgentExecutor_EmitsResponseIFFConfiguredAsync(bool executorSetting)
     {
         // Arrange
-        TestRunContext testContext = new();
+        TestRunContext testContext = await PrepareHandoffSharedStateAsync();
         TestReplayAgent agent = new(TestMessages, TestAgentId, TestAgentName);
 
         HandoffAgentExecutorOptions options = new("",
@@ -67,7 +92,7 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
         testContext.ConfigureExecutor(executor);
 
         // Act
-        HandoffState message = new(new(false), null, []);
+        HandoffState message = new(new(false), null, null);
         await executor.HandleAsync(message, testContext.BindWorkflowContext(executor.Id));
 
         // Assert
@@ -75,6 +100,82 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
         CheckResponseEventsAgainstTestMessages(updates, expectingResponse: executorSetting, agent.GetDescriptiveId());
     }
 
+    [Fact]
+    public async Task Test_HandoffAgentExecutor_ComposesWithHITLSubworkflowAsync()
+    {
+        // Arrange
+        TestRunContext testContext = await PrepareHandoffSharedStateAsync();
+
+        SendsRequestExecutor challengeSender = new();
+        Workflow subworkflow = new WorkflowBuilder(challengeSender)
+                                   .AddExternalRequest<Challenge, Response>(challengeSender, "SendChallengeToUser")
+                                   .WithOutputFrom(challengeSender)
+                                   .Build();
+
+        InProcessExecutionEnvironment environment = InProcessExecution.Lockstep.WithCheckpointing(CheckpointManager.CreateInMemory());
+        AIAgent subworkflowAgent = subworkflow.AsAIAgent(includeWorkflowOutputsInResponse: true, name: "Subworkflow", executionEnvironment: environment);
+        HandoffAgentExecutorOptions options = new("",
+                                                  emitAgentResponseEvents: true,
+                                                  emitAgentResponseUpdateEvents: true,
+                                                  HandoffToolCallFilteringBehavior.None);
+
+        HandoffAgentExecutor executor = new(subworkflowAgent, [], options);
+        Workflow fakeWorkflow = new(executor.Id) { ExecutorBindings = { { executor.Id, executor } } };
+        EdgeMap map = new(testContext, fakeWorkflow, null);
+
+        testContext.ConfigureExecutor(executor, map);
+
+        // Validate that our test assumptions hold
+        string functionCallPortId = $"{HandoffAgentExecutor.IdFor(subworkflowAgent)}_FunctionCall";
+        map.TryGetResponsePortExecutorId(functionCallPortId, out string? responsePortExecutorId).Should().BeTrue();
+        responsePortExecutorId.Should().Be(executor.Id);
+
+        // Act
+        HandoffState message = new(new(false), null, null);
+        await executor.HandleAsync(message, testContext.BindWorkflowContext(executor.Id));
+
+        await testContext.StateManager.PublishUpdatesAsync(null);
+
+        // Assert
+        testContext.ExternalRequests.Should().HaveCount(1)
+                                .And.ContainSingle(request => request.IsDataOfType<FunctionCallContent>());
+
+        FunctionCallContent functionCallContent = testContext.ExternalRequests.Single().Data.As<FunctionCallContent>()!;
+        object? requestData = functionCallContent.Arguments!["data"];
+
+        Challenge? challenge = null;
+        if (requestData is PortableValue pv)
+        {
+            challenge = pv.As<Challenge>();
+        }
+        else
+        {
+            challenge = requestData as Challenge;
+        }
+
+        if (challenge is null)
+        {
+            Assert.Fail($"Expected request data to be of type {typeof(Challenge).FullName}, but was {requestData?.GetType().FullName ?? "null"}");
+            return; // Unreachable, but analysis cannot infer that Debug.Fail will throw/exit, and UnreachableException is not available on net472
+        }
+
+        // Act 2
+        string challengeResponse = new(challenge.Value.Reverse().ToArray());
+        FunctionResultContent responseContent = new(functionCallContent.CallId, new Response(challengeResponse));
+
+        RequestPortInfo requestPortInfo = new(new(typeof(Challenge)), new(typeof(Response)), functionCallPortId);
+        string requestId = $"{functionCallPortId.Length}:{functionCallPortId}:{functionCallContent.CallId}";
+        DeliveryMapping? mapping = await map.PrepareDeliveryForResponseAsync(new(requestPortInfo, requestId, new(responseContent)));
+
+        mapping!.Deliveries.Should().HaveCount(1);
+
+        MessageDelivery delivery = mapping!.Deliveries.Single();
+
+        object? result = await executor.ExecuteCoreAsync(delivery.Envelope.Message,
+                                                         delivery.Envelope.MessageType,
+                                                         testContext.BindWorkflowContext(executor.Id));
+    }
+
     [Fact]
     public async Task Test_HandoffAgentExecutor_PreservesExistingInstructionsAndToolsAsync()
     {
@@ -92,80 +193,113 @@ public class HandoffAgentExecutorTests : AIAgentHostingExecutorTestsBase
         HandoffTarget handoff = new(targetAgent);
         HandoffAgentExecutor executor = new(handoffAgent, [handoff], options);
 
-        TestWorkflowContext testContext = new(executor.Id);
-        HandoffState state = new(new(false), null, [], null);
+        TestRunContext runContext = await PrepareHandoffSharedStateAsync();
+        IWorkflowContext testContext = runContext.BindWorkflowContext(executor.Id);
+        HandoffState state = new(new(false), null);
 
         // Act / Assert
         Func<Task> runStreamingAsync = async () => await executor.HandleAsync(state, testContext);
         await runStreamingAsync.Should().NotThrowAsync();
     }
+}
 
-    private sealed class OptionValidatingChatClient(string baseInstructions, string handoffInstructions, AITool baseTool) : IChatClient
+internal sealed record Challenge(string Value);
+internal sealed record Response(string Value);
+
+[SendsMessage(typeof(Challenge))]
+internal sealed partial class SendsRequestExecutor(string? id = null) : ChatProtocolExecutor(id ?? nameof(SendsRequestExecutor), s_chatOptions)
+{
+    internal const string ChallengeString = "{C7A762AE-7DAA-4D9C-A647-E64E6DBC35AE}";
+    private static string ResponseKey { get; } = new(ChallengeString.Reverse().ToArray());
+
+    private static readonly ChatProtocolExecutorOptions s_chatOptions = new()
     {
-        public void Dispose()
+        AutoSendTurnToken = false
+    };
+
+    protected override ValueTask TakeTurnAsync(List<ChatMessage> messages, IWorkflowContext context, bool? emitEvents, CancellationToken cancellationToken = default)
+        => context.SendMessageAsync(new Challenge(ChallengeString), cancellationToken);
+
+    [MessageHandler]
+    public async ValueTask HandleChallengeResponseAsync(Response response, IWorkflowContext context, CancellationToken cancellationToken = default)
+    {
+        if (response.Value != ResponseKey)
         {
+            throw new InvalidOperationException($"Incorrect response received. Expected '{ResponseKey}' but got '{response.Value}'");
         }
 
-        private void CheckOptions(ChatOptions? options)
-        {
-            options.Should().NotBeNull();
+        await context.SendMessageAsync(new ChatMessage(ChatRole.Assistant, "Correct response."), cancellationToken)
+                     .ConfigureAwait(false);
 
-            options.Instructions.Should().NotBeNullOrEmpty("Handoff orchestration should preserve and augment instructions.")
-                                     .And.Contain(baseInstructions, because: "Handoff orchestration should preserve existing instructions.")
-                                     .And.Contain(handoffInstructions, because: "Handoff orchestration should inject handoff instructions.");
+        await context.SendMessageAsync(new TurnToken(false), cancellationToken).ConfigureAwait(false);
+    }
+}
 
-            options.Tools.Should().NotBeNullOrEmpty("Handoff orchestration should preserve and augment tools.")
-                                  .And.Contain(tool => tool.Name == baseTool.Name, "Handoff orchestration should preserve existing tools.")
-                                  .And.Contain(tool => tool.Name.StartsWith(HandoffWorkflowBuilder.FunctionPrefix, StringComparison.Ordinal),
-                                               because: "Handoff orchestration should inject handoff tools.");
-        }
+internal sealed class OptionValidatingChatClient(string baseInstructions, string handoffInstructions, AITool baseTool) : IChatClient
+{
+    public void Dispose()
+    {
+    }
 
-        private List<ChatMessage> ResponseMessages =>
-            [
-                new ChatMessage(ChatRole.Assistant, "Ok")
+    private void CheckOptions(ChatOptions? options)
+    {
+        options.Should().NotBeNull();
+
+        options.Instructions.Should().NotBeNullOrEmpty("Handoff orchestration should preserve and augment instructions.")
+                                 .And.Contain(baseInstructions, because: "Handoff orchestration should preserve existing instructions.")
+                                 .And.Contain(handoffInstructions, because: "Handoff orchestration should inject handoff instructions.");
+
+        options.Tools.Should().NotBeNullOrEmpty("Handoff orchestration should preserve and augment tools.")
+                              .And.Contain(tool => tool.Name == baseTool.Name, "Handoff orchestration should preserve existing tools.")
+                              .And.Contain(tool => tool.Name.StartsWith(HandoffWorkflowBuilder.FunctionPrefix, StringComparison.Ordinal),
+                                           because: "Handoff orchestration should inject handoff tools.");
+    }
+
+    private List<ChatMessage> ResponseMessages =>
+        [
+            new ChatMessage(ChatRole.Assistant, "Ok")
                 {
                     MessageId = Guid.NewGuid().ToString(),
                     AuthorName = nameof(OptionValidatingChatClient)
                 }
-            ];
+        ];
 
-        public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+    public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        this.CheckOptions(options);
+
+        ChatResponse response = new(this.ResponseMessages)
         {
-            this.CheckOptions(options);
+            ResponseId = Guid.NewGuid().ToString("N"),
+            CreatedAt = DateTimeOffset.Now
+        };
 
-            ChatResponse response = new(this.ResponseMessages)
+        return Task.FromResult(response);
+    }
+
+    public object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        if (serviceType == typeof(OptionValidatingChatClient))
+        {
+            return this;
+        }
+
+        return null;
+    }
+
+    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        this.CheckOptions(options);
+
+        string responseId = Guid.NewGuid().ToString("N");
+        foreach (ChatMessage message in this.ResponseMessages)
+        {
+            yield return new(message.Role, message.Contents)
             {
-                ResponseId = Guid.NewGuid().ToString("N"),
+                ResponseId = responseId,
+                MessageId = message.MessageId,
                 CreatedAt = DateTimeOffset.Now
             };
-
-            return Task.FromResult(response);
-        }
-
-        public object? GetService(Type serviceType, object? serviceKey = null)
-        {
-            if (serviceType == typeof(OptionValidatingChatClient))
-            {
-                return this;
-            }
-
-            return null;
-        }
-
-        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
-        {
-            this.CheckOptions(options);
-
-            string responseId = Guid.NewGuid().ToString("N");
-            foreach (ChatMessage message in this.ResponseMessages)
-            {
-                yield return new(message.Role, message.Contents)
-                {
-                    ResponseId = responseId,
-                    MessageId = message.MessageId,
-                    CreatedAt = DateTimeOffset.Now
-                };
-            }
         }
     }
 }

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/12_HandOff_HostAsAgent.cs

@@ -73,6 +73,7 @@ internal static class Step12EntryPoint
         foreach (string input in inputs)
         {
             AgentResponse response;
+
             ResponseContinuationToken? continuationToken = null;
             do
             {

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/SampleSmokeTest.cs

@@ -314,6 +314,38 @@ public class SampleSmokeTest
         Action<string> CreateValidator(string expected) => actual => actual.Should().Be(expected);
     }
 
+    public class Step12ExpectedOutputCalculator(int agentCount)
+    {
+        private readonly int[] _bookmarks = new int[agentCount];
+        private readonly List<string> _history = new();
+        private readonly HashSet<int> _skipIndices = new();
+
+        public IEnumerable<string> ExpectedOutputs =>
+            this._history.Where((element, index) => !this._skipIndices.Contains(index));
+
+        public void ProcessInput(string newInput)
+        {
+            this._skipIndices.Add(this._history.Count);
+            this._history.Add(newInput);
+
+            for (int i = 0; i < agentCount; i++)
+            {
+                int agentId = i + 1;
+                int agentBookmark = this._bookmarks[i];
+                int count = this._history.Count - agentBookmark;
+
+                count.Should().BeGreaterThanOrEqualTo(0);
+
+                foreach (string input in this._history.Skip(agentBookmark).ToList())
+                {
+                    this._history.Add($"{agentId}:{input}");
+                }
+
+                this._bookmarks[i] = this._history.Count;
+            }
+        }
+    }
+
     [Theory]
     [InlineData(ExecutionEnvironment.InProcess_Lockstep)]
     [InlineData(ExecutionEnvironment.InProcess_OffThread)]
@@ -322,14 +354,10 @@ public class SampleSmokeTest
     {
         List<string> inputs = ["1", "2", "3"];
 
-        using StringWriter writer = new();
-        await Step12EntryPoint.RunAsync(writer, environment.ToWorkflowExecutionEnvironment(), inputs);
-
-        string[] lines = writer.ToString().Split(['\r', '\n'], StringSplitOptions.RemoveEmptyEntries);
-
         // The expectation is that each agent will echo each input along with every echo from previous agents
         // E.g.:
         // (user): 1
+        // ----- outputs below
         // (a1): 1:1
         // (a2): 2:1
         // (a2): 2:1:1
@@ -340,7 +368,35 @@ public class SampleSmokeTest
         // (a3): 3:2:1
         // (a3): 3:2:1:1
 
-        string[] expected = inputs.SelectMany(input => EchoesForInput(input)).ToArray();
+        // If there are multiple inputs (there are), then each successive input adds to the depth of the previous
+        // ones, so, for example, once we do input = "1", "2":
+
+        // (user): 1
+        // (a1): 1:1     <- a1 "last seen"
+        // (a2): 2:1
+        // (a2): 2:1:1   <- a2 "last seen"
+        // (user): 2
+        // ----- outputs below
+        // (a1): 1:2:1
+        // (a1): 1:2:1:1
+        // (a1): 1:2     <- from user input, a1 "last seen"
+        // (a2): 2:2     <- from user input (note that a2 seems like it is seeing these in a different "order" than a1 - but it is not)
+        // (a2): 2:1:2:1
+        // (a2): 2:1:2:1:1
+        // (a2): 2:1:2   <- from a1's first echo, a2 "last seen"
+
+        Step12ExpectedOutputCalculator outputGenerator = new(Step12EntryPoint.AgentCount);
+        foreach (string input in inputs)
+        {
+            outputGenerator.ProcessInput(input);
+        }
+
+        string[] expected = outputGenerator.ExpectedOutputs.ToArray();
+
+        using StringWriter writer = new();
+        await Step12EntryPoint.RunAsync(writer, environment.ToWorkflowExecutionEnvironment(), inputs);
+
+        string[] lines = writer.ToString().Split(['\r', '\n'], StringSplitOptions.RemoveEmptyEntries);
 
         Console.Error.WriteLine("Expected lines: ");
         foreach (string expectedLine in expected)
@@ -357,19 +413,6 @@ public class SampleSmokeTest
         Assert.Collection(lines,
                           expected.Select(CreateValidator).ToArray());
 
-        IEnumerable<string> EchoesForInput(string input)
-        {
-            List<string> echoes = [$"{Step12EntryPoint.EchoPrefixForAgent(1)}{input}"];
-            for (int i = 2; i <= Step12EntryPoint.AgentCount; i++)
-            {
-                string agentPrefix = Step12EntryPoint.EchoPrefixForAgent(i);
-                List<string> newEchoes = [$"{agentPrefix}{input}", .. echoes.Select(echo => $"{agentPrefix}{echo}")];
-                echoes.AddRange(newEchoes);
-            }
-
-            return echoes;
-        }
-
         Action<string> CreateValidator(string expected) => actual => actual.Should().Be(expected);
     }
 

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRunContext.cs

@@ -27,6 +27,9 @@ public class TestRunContext : IRunnerContext
 
     internal TestRunContext ConfigureExecutor(Executor executor, EdgeMap? map = null)
     {
+        // Ensure that we have run the ProtocolBuilder
+        _ = executor.Protocol.Describe();
+
         executor.AttachRequestContext(new TestExternalRequestContext(this, executor.Id, map));
         this.Executors.Add(executor.Id, executor);
         return this;
@@ -42,6 +45,7 @@ public class TestRunContext : IRunnerContext
         return this;
     }
 
+    internal StateManager StateManager { get; } = new();
     private sealed class BoundContext(
         string executorId,
         TestRunContext runnerContext,
@@ -70,16 +74,16 @@ public class TestRunContext : IRunnerContext
             => this.AddEventAsync(new RequestHaltEvent());
 
         public ValueTask QueueClearScopeAsync(string? scopeName = null, CancellationToken cancellationToken = default)
-            => default;
+            => runnerContext.StateManager.ClearStateAsync(executorId, scopeName);
 
         public ValueTask QueueStateUpdateAsync<T>(string key, T? value, string? scopeName = null, CancellationToken cancellationToken = default)
-            => default;
+            => runnerContext.StateManager.WriteStateAsync(new ScopeId(executorId, scopeName), key, value);
 
         public ValueTask<T?> ReadStateAsync<T>(string key, string? scopeName = null, CancellationToken cancellationToken = default)
-            => new(default(T?));
+            => runnerContext.StateManager.ReadStateAsync<T>(new ScopeId(executorId, scopeName), key);
 
         public ValueTask<HashSet<string>> ReadStateKeysAsync(string? scopeName = null, CancellationToken cancellationToken = default)
-            => new([]);
+            => runnerContext.StateManager.ReadKeysAsync(new ScopeId(executorId, scopeName));
 
         public ValueTask SendMessageAsync(object message, string? targetId = null, CancellationToken cancellationToken = default)
             => runnerContext.SendMessageAsync(executorId, message, targetId, cancellationToken);