.NET: Bug fixes for AGUI hosting and workflows

dotnet/samples/02-agents/AGUI/Step01_GettingStarted/Server/Program.cs

@@ -10,6 +10,11 @@
 builder.Services.AddHttpClient().AddLogging();
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"]

dotnet/samples/02-agents/AGUI/Step01_GettingStarted/Server/Server.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/02-agents/AGUI/Step02_BackendTools/Server/Program.cs

@@ -16,6 +16,11 @@
     options.SerializerOptions.TypeInfoResolverChain.Add(SampleJsonSerializerContext.Default));
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"]

dotnet/samples/02-agents/AGUI/Step02_BackendTools/Server/Server.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/02-agents/AGUI/Step03_FrontendTools/Server/Program.cs

@@ -10,6 +10,11 @@
 builder.Services.AddHttpClient().AddLogging();
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"]

dotnet/samples/02-agents/AGUI/Step03_FrontendTools/Server/Server.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/02-agents/AGUI/Step04_HumanInLoop/Server/Program.cs

@@ -27,6 +27,11 @@
     options.SerializerOptions.TypeInfoResolverChain.Add(ApprovalJsonContext.Default));
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 app.UseHttpLogging();

dotnet/samples/02-agents/AGUI/Step04_HumanInLoop/Server/Server.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/02-agents/AGUI/Step05_StateManagement/Server/Program.cs

@@ -17,6 +17,11 @@
 // Configure to listen on port 8888
 builder.WebHost.UseUrls("http://localhost:8888");
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"]

dotnet/samples/02-agents/AGUI/Step05_StateManagement/Server/Server.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/05-end-to-end/AGUIClientServer/AGUIDojoServer/AGUIDojoServer.csproj

@@ -15,6 +15,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>
 

dotnet/samples/05-end-to-end/AGUIClientServer/AGUIDojoServer/Program.cs

@@ -19,6 +19,11 @@
 builder.Services.ConfigureHttpJsonOptions(options => options.SerializerOptions.TypeInfoResolverChain.Add(AGUIDojoServerSerializerContext.Default));
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 app.UseHttpLogging();

dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs

@@ -49,8 +49,9 @@
                 AGUIServerSerializerContext.Default.Options)
         ]);
 
-// When running in production, make sure to use an SessionIsolationKeyProvider, e.g. ClaimsIdentity-based
-// if using Claims-based Identity for Authentication/Authorization
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
 // builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
 
 // Register the agent with the host and configure it to use an in-memory session store

dotnet/samples/05-end-to-end/AGUIWebChat/Server/AGUIWebChatServer.csproj

@@ -14,6 +14,7 @@
 
   <ItemGroup>
     <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore\Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.csproj" />
+    <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.Hosting.AspNetCore\Microsoft.Agents.AI.Hosting.AspNetCore.csproj" />
     <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.AGUI\Microsoft.Agents.AI.AGUI.csproj" />
     <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
   </ItemGroup>

dotnet/samples/05-end-to-end/AGUIWebChat/Server/Program.cs

@@ -12,6 +12,11 @@
 builder.Services.AddHttpClient().AddLogging();
 builder.Services.AddAGUI();
 
+// WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production,
+// make sure to also register a SessionIsolationKeyProvider to scope sessions by principal in multi-user
+// deployments, e.g.:
+// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });
+
 WebApplication app = builder.Build();
 
 string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"] ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");

dotnet/src/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore/AGUIEndpointRouteBuilderExtensions.cs

@@ -103,7 +103,16 @@ public static IEndpointConventionBuilder MapAGUI(
         ArgumentNullException.ThrowIfNull(aiAgent);
 
         var agentSessionStore = endpoints.ServiceProvider.GetKeyedService<AgentSessionStore>(aiAgent.Name);
-        var hostAgent = new AIHostAgent(aiAgent, agentSessionStore ?? new NoopAgentSessionStore());
+
+        // Ensure that we have an IsolationKeyScopedAgentSessionStore registered.
+        var isolationKeyProvider = endpoints.ServiceProvider.GetService<SessionIsolationKeyProvider>();
+        if (agentSessionStore?.GetService<IsolationKeyScopedAgentSessionStore>() is null)
+        {
+            agentSessionStore ??= new NoopAgentSessionStore();
+            agentSessionStore = new IsolationKeyScopedAgentSessionStore(agentSessionStore, isolationKeyProvider, new() { Strict = isolationKeyProvider != null });
+        }
+
+        var hostAgent = new AIHostAgent(aiAgent, agentSessionStore);
 
         return endpoints.MapPost(pattern, async ([FromBody] RunAgentInput? input, HttpContext context, CancellationToken cancellationToken) =>
         {

dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/ObjectModel/InvokeMcpToolExecutor.cs

@@ -27,6 +27,14 @@ internal sealed class InvokeMcpToolExecutor(
     WorkflowFormulaState state) :
     DeclarativeActionExecutor<InvokeMcpTool>(model, state)
 {
+    private const string ApprovalSnapshotStateKey = nameof(_approvalSnapshot);
+
+    /// <summary>
+    /// Snapshot of evaluated parameters at approval-request time.
+    /// Used to prevent TOCTOU attacks where state mutates during the approval window.
+    /// </summary>
+    private ApprovalSnapshot? _approvalSnapshot;
+
     /// <summary>
     /// Step identifiers for the MCP tool invocation workflow.
     /// </summary>
@@ -75,6 +83,10 @@ public static bool RequiresNothing(object? message) =>
 
         if (requireApproval)
         {
+            // Snapshot the evaluated parameters to prevent TOCTOU attacks.
+            // If state mutates during the approval window, the approved values are used on resume.
+            this._approvalSnapshot = new ApprovalSnapshot(serverUrl, serverLabel, toolName, arguments, connectionName);
+
             // Create tool call content for approval request.
             // Transport headers (e.g. Authorization) are intentionally excluded from the
             // approval event: they must not cross into the externally-surfaced approval request.
@@ -137,13 +149,14 @@ public async ValueTask CaptureResponseAsync(
             return;
         }
 
-        // Approved - now invoke the tool
-        string serverUrl = this.GetServerUrl();
-        string? serverLabel = this.GetServerLabel();
-        string toolName = this.GetToolName();
-        Dictionary<string, object?>? arguments = this.GetArguments();
+        // Approved - use the snapshot from approval-request time to prevent TOCTOU attacks.
+        // Headers are re-evaluated (they may contain auth secrets that should not be persisted).
+        string serverUrl = this._approvalSnapshot?.ServerUrl ?? this.GetServerUrl();
+        string? serverLabel = this._approvalSnapshot?.ServerLabel ?? this.GetServerLabel();
+        string toolName = this._approvalSnapshot?.ToolName ?? this.GetToolName();
+        Dictionary<string, object?>? arguments = this._approvalSnapshot?.Arguments ?? this.GetArguments();
         Dictionary<string, string>? headers = this.GetHeaders();
-        string? connectionName = this.GetConnectionName();
+        string? connectionName = this._approvalSnapshot?.ConnectionName ?? this.GetConnectionName();
 
         McpServerToolResultContent resultContent = await mcpToolHandler.InvokeToolAsync(
             serverUrl,
@@ -162,9 +175,33 @@ public async ValueTask CaptureResponseAsync(
     /// </summary>
     public async ValueTask CompleteAsync(IWorkflowContext context, ActionExecutorResult message, CancellationToken cancellationToken)
     {
+        // Clear the approval snapshot after successful completion.
+        this._approvalSnapshot = null;
+        await ClearSnapshotStateAsync(context, cancellationToken).ConfigureAwait(false);
+
         await context.RaiseCompletionEventAsync(this.Model, cancellationToken).ConfigureAwait(false);
     }
 
+    /// <inheritdoc/>
+    /// <remarks>
+    /// Persists the approval snapshot to workflow state so it survives checkpoint/restore cycles.
+    /// </remarks>
+    protected override async ValueTask OnCheckpointingAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
+    {
+        await context.QueueStateUpdateAsync(ApprovalSnapshotStateKey, this._approvalSnapshot, null, cancellationToken).ConfigureAwait(false);
+        await base.OnCheckpointingAsync(context, cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <inheritdoc/>
+    /// <remarks>
+    /// Restores the approval snapshot from workflow state after a checkpoint restore.
+    /// </remarks>
+    protected override async ValueTask OnCheckpointRestoredAsync(IWorkflowContext context, CancellationToken cancellationToken = default)
+    {
+        await base.OnCheckpointRestoredAsync(context, cancellationToken).ConfigureAwait(false);
+        this._approvalSnapshot = await context.ReadStateAsync<ApprovalSnapshot>(ApprovalSnapshotStateKey, null, cancellationToken).ConfigureAwait(false);
+    }
+
     private async ValueTask ProcessResultAsync(IWorkflowContext context, McpServerToolResultContent resultContent, CancellationToken cancellationToken)
     {
         bool autoSend = this.GetAutoSendValue();
@@ -365,4 +402,24 @@ private bool GetAutoSendValue()
 
         return result;
     }
+
+    /// <summary>
+    /// Clears the persisted approval snapshot state after a successful tool invocation.
+    /// </summary>
+    private static async ValueTask ClearSnapshotStateAsync(IWorkflowContext context, CancellationToken cancellationToken)
+    {
+        await context.QueueStateUpdateAsync<ApprovalSnapshot?>(ApprovalSnapshotStateKey, null, null, cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <summary>
+    /// Stores the evaluated parameters at approval-request time so that
+    /// <see cref="CaptureResponseAsync"/> uses the values the user reviewed,
+    /// even if <see cref="WorkflowFormulaState"/> mutates during the approval window.
+    /// </summary>
+    internal sealed record ApprovalSnapshot(
+        string ServerUrl,
+        string? ServerLabel,
+        string ToolName,
+        Dictionary<string, object?>? Arguments,
+        string? ConnectionName);
 }