.NET: fix: race concdition in workflow formula state

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/DeclarativeWorkflowBuilder.cs

@@ -74,6 +74,7 @@ public static Workflow Build<TInput>(
 
         WorkflowFormulaState state = new(options.CreateRecalcEngine());
         state.Initialize(workflowElement.WrapWithBot(), options.Configuration);
+        state.CaptureInitialState();
         DeclarativeWorkflowExecutor<TInput> rootExecutor =
             new(rootId,
                 options,

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeActionExecutor.cs

@@ -64,6 +64,7 @@ protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBui
     /// <inheritdoc/>
     public virtual ValueTask ResetAsync()
     {
+        this._state.Reset();
         return default;
     }
 

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DeclarativeWorkflowExecutor.cs

@@ -42,9 +42,14 @@ internal sealed class DeclarativeWorkflowExecutor<TInput>(
     /// <inheritdoc/>
     public ValueTask ResetAsync()
     {
+        state.Reset();
         return default;
     }
 
+    /// <inheritdoc/>
+    protected override ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default) =>
+        state.RestoreAsync(context, cancellationToken);
+
     /// <inheritdoc/>
     [SendsMessage(typeof(ActionExecutorResult))]
     public override ValueTask HandleAsync(TInput message, IWorkflowContext context, CancellationToken cancellationToken = default)

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Interpreter/DelegateActionExecutor.cs

@@ -52,6 +52,7 @@ protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBui
     /// <inheritdoc/>
     public ValueTask ResetAsync()
     {
+        this._state.Reset();
         return default;
     }
 

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/PowerFx/WorkflowFormulaState.cs

@@ -29,6 +29,8 @@ internal sealed class WorkflowFormulaState
 
     private readonly Dictionary<string, WorkflowScope> _scopes;
 
+    private Dictionary<string, WorkflowScope> _initialScopes;
+
     private int _isInitialized;
 
     public RecalcEngine Engine { get; }
@@ -38,6 +40,7 @@ internal sealed class WorkflowFormulaState
     public WorkflowFormulaState(RecalcEngine engine)
     {
         this._scopes = VariableScopeNames.AllScopes.ToDictionary(scopeName => GetScopeName(scopeName), _ => new WorkflowScope());
+        this._initialScopes = this.CreateScopeSnapshot();
 
         this.Engine = engine;
         this.Evaluator = new WorkflowExpressionEngine(engine);
@@ -61,13 +64,27 @@ public void Set(string variableName, FormulaValue value, string? scopeName = nul
 
     public bool SetInitialized() => Interlocked.CompareExchange(ref this._isInitialized, 1, 0) == 0;
 
+    public void CaptureInitialState()
+    {
+        this._initialScopes = this.CreateScopeSnapshot();
+    }
+
+    public void Reset()
+    {
+        this.RestoreInitialState();
+        Interlocked.Exchange(ref this._isInitialized, 0);
+        this.Bind();
+    }
+
     public async ValueTask RestoreAsync(IWorkflowContext context, CancellationToken cancellationToken)
     {
         if (!this.SetInitialized())
         {
             return;
         }
 
+        this.RestoreInitialState();
+
         Stopwatch timer = Stopwatch.StartNew();
         Debug.WriteLine("RESTORE CHECKPOINT - BEGIN");
         await Task.WhenAll(RestorableScopes.Select(scopeName => ReadScopeAsync(scopeName))).ConfigureAwait(false);
@@ -93,6 +110,22 @@ async Task ReadScopeAsync(string scopeName)
         }
     }
 
+    private Dictionary<string, WorkflowScope> CreateScopeSnapshot() =>
+        this._scopes.ToDictionary(scope => scope.Key, scope => new WorkflowScope(scope.Value));
+
+    private void RestoreInitialState()
+    {
+        foreach (KeyValuePair<string, WorkflowScope> initialScopeEntry in this._initialScopes)
+        {
+            WorkflowScope scope = this._scopes[initialScopeEntry.Key];
+            scope.Clear();
+            foreach (KeyValuePair<string, FormulaValue> initialValueEntry in initialScopeEntry.Value)
+            {
+                scope[initialValueEntry.Key] = initialValueEntry.Value;
+            }
+        }
+    }
+
     public void Bind(string? scopeNameToBind = null)
     {
         if (scopeNameToBind is not null)
@@ -143,5 +176,15 @@ public static string GetScopeName(string? scopeName)
     /// <summary>
     /// The set of variables for a specific action scope.
     /// </summary>
-    private sealed class WorkflowScope : Dictionary<string, FormulaValue>;
+    private sealed class WorkflowScope : Dictionary<string, FormulaValue>
+    {
+        public WorkflowScope()
+        {
+        }
+
+        public WorkflowScope(IDictionary<string, FormulaValue> values)
+            : base(values)
+        {
+        }
+    }
 }

dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.UnitTests/DeclarativeWorkflowTest.cs

@@ -79,6 +79,52 @@ public async Task EndConversationActionAsync()
         this.AssertNotExecuted("sendActivity_1");
     }
 
+    [Fact]
+    public async Task HostedWorkflowAgentIsolatesDeclarativeStateBySessionAsync()
+    {
+        // Arrange
+        RecordingAgentProvider provider = new();
+        AIAgent agent = CreateStateEchoWorkflow(provider).AsAIAgent(id: "host", name: "host");
+        AgentSession aliceSession = await agent.CreateSessionAsync();
+        AgentSession mallorySession = await agent.CreateSessionAsync();
+
+        // Act
+        AgentResponse aliceSeedResponse = await agent.RunAsync("EMBER-QUARTZ-7319", aliceSession);
+        AgentResponse aliceResumeResponse = await agent.RunAsync("inspect-alice", aliceSession);
+        AgentResponse malloryInspectResponse = await agent.RunAsync("inspect-mallory", mallorySession);
+        _ = await agent.RunAsync("ONYX-CEDAR-4826", mallorySession);
+        AgentResponse aliceFinalResponse = await agent.RunAsync("inspect-alice-again", aliceSession);
+
+        // Assert
+        Assert.NotSame(aliceSession, mallorySession);
+        Assert.Contains("Marker: \"\"", aliceSeedResponse.Text, StringComparison.Ordinal);
+        Assert.Contains("EMBER-QUARTZ-7319", aliceResumeResponse.Text, StringComparison.Ordinal);
+        Assert.DoesNotContain("EMBER-QUARTZ-7319", malloryInspectResponse.Text, StringComparison.Ordinal);
+        Assert.DoesNotContain("ONYX-CEDAR-4826", aliceFinalResponse.Text, StringComparison.Ordinal);
+        Assert.Contains("inspect-alice", aliceFinalResponse.Text, StringComparison.Ordinal);
+
+        Assert.True(provider.MessageConversations.Count >= 3);
+        Assert.Equal(provider.MessageConversations[0], provider.MessageConversations[1]);
+        Assert.NotEqual(provider.MessageConversations[0], provider.MessageConversations[2]);
+    }
+
+    [Fact]
+    public async Task HostedWorkflowAgentIsolatesDeclarativeStateForImplicitSessionsAsync()
+    {
+        // Arrange
+        RecordingAgentProvider provider = new();
+        AIAgent agent = CreateStateEchoWorkflow(provider).AsAIAgent(id: "host", name: "host");
+
+        // Act
+        _ = await agent.RunAsync("EMBER-QUARTZ-7319");
+        AgentResponse secondImplicitResponse = await agent.RunAsync("inspect-implicit");
+
+        // Assert
+        Assert.DoesNotContain("EMBER-QUARTZ-7319", secondImplicitResponse.Text, StringComparison.Ordinal);
+        Assert.True(provider.MessageConversations.Count >= 2);
+        Assert.NotEqual(provider.MessageConversations[0], provider.MessageConversations[1]);
+    }
+
     [Fact]
     public async Task GotoActionAsync()
     {
@@ -386,6 +432,81 @@ private Workflow CreateWorkflow<TInput>(string workflowPath, TInput workflowInpu
         return DeclarativeWorkflowBuilder.Build<TInput>(yamlReader, workflowContext);
     }
 
+    private static Workflow CreateStateEchoWorkflow(ResponseAgentProvider provider)
+    {
+        using StringReader yamlReader = new(
+            """
+                kind: Workflow
+                trigger:
+
+                  kind: OnConversationStart
+                  id: state_echo_workflow
+                  actions:
+
+                    - kind: SendActivity
+                      id: show_marker
+                      activity: |-
+                        Marker: "{Local.Marker}"
+
+                    - kind: SetVariable
+                      id: set_marker
+                      variable: Local.Marker
+                      value: =System.LastMessageText
+                """);
+        DeclarativeWorkflowOptions options = new(provider);
+
+        return DeclarativeWorkflowBuilder.Build<string>(yamlReader, options);
+    }
+
+    private sealed class RecordingAgentProvider : ResponseAgentProvider
+    {
+        public List<string> MessageConversations { get; } = [];
+
+        private int _conversationCount;
+
+        public override Task<string> CreateConversationAsync(CancellationToken cancellationToken = default) =>
+            Task.FromResult($"conversation-{Interlocked.Increment(ref this._conversationCount):D2}");
+
+        public override Task<ChatMessage> CreateMessageAsync(
+            string conversationId,
+            ChatMessage conversationMessage,
+            CancellationToken cancellationToken = default)
+        {
+            this.MessageConversations.Add(conversationId);
+            return Task.FromResult(conversationMessage);
+        }
+
+        public override Task<ChatMessage> GetMessageAsync(
+            string conversationId,
+            string messageId,
+            CancellationToken cancellationToken = default) =>
+            Task.FromResult(new ChatMessage(ChatRole.Assistant, string.Empty) { MessageId = messageId });
+
+        public override async IAsyncEnumerable<AgentResponseUpdate> InvokeAgentAsync(
+            string agentId,
+            string? agentVersion,
+            string? conversationId,
+            IEnumerable<ChatMessage>? messages,
+            IDictionary<string, object?>? inputArguments,
+            [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            await Task.CompletedTask;
+            yield break;
+        }
+
+        public override async IAsyncEnumerable<ChatMessage> GetMessagesAsync(
+            string conversationId,
+            int? limit = null,
+            string? after = null,
+            string? before = null,
+            bool newestFirst = false,
+            [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            await Task.CompletedTask;
+            yield break;
+        }
+    }
+
     private static Mock<ResponseAgentProvider> CreateMockProvider(string input)
     {
         Mock<ResponseAgentProvider> mockAgentProvider = new(MockBehavior.Strict);