From 8db651ac2b431a4a333557acdeb83336d0069bf7 Mon Sep 17 00:00:00 2001 From: Peter Ibekwe Date: Mon, 9 Mar 2026 17:33:49 -0700 Subject: [PATCH] Update pending requests to use typed properties instead of relying on StateBag. replying to PR feedback. --- .../WorkflowSession.cs | 41 +++++++++++----- .../WorkflowsJsonUtilities.cs | 1 + .../JsonSerializationTests.cs | 49 +++++++++++++++++++ 3 files changed, 80 insertions(+), 11 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs index 64354c58c4..f08f0c85a3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs @@ -25,9 +25,24 @@ internal sealed class WorkflowSession : AgentSession private InMemoryCheckpointManager? _inMemoryCheckpointManager; - // Key prefix used in StateBag to track pending external requests by their content ID - // This enables converting incoming response content back to ExternalResponse when resuming. - private const string PendingRequestKeyPrefix = "workflow.pendingRequest:"; + /// + /// Tracks pending external requests by their content ID (e.g., + /// or ). This mapping enables converting incoming response + /// content back to when resuming a workflow from a checkpoint. + /// + /// + /// + /// Entries are added when a is received during workflow execution, + /// and removed when a matching response is delivered via . + /// + /// + /// The number of entries is bounded by the number of outstanding external requests in a single workflow run. + /// When a session is abandoned, all pending requests are released with the session object. + /// Request-level timeouts, if needed, should be implemented in the workflow definition itself + /// (e.g., using a timer racing against an external event). + /// + /// + private readonly Dictionary _pendingRequests = []; internal static bool VerifyCheckpointingConfiguration(IWorkflowExecutionEnvironment executionEnvironment, [NotNullWhen(true)] out InProcessExecutionEnvironment? inProcEnv) { @@ -94,6 +109,7 @@ internal sealed class WorkflowSession : AgentSession this.LastCheckpoint = sessionState.LastCheckpoint; this.StateBag = sessionState.StateBag; + this._pendingRequests = sessionState.PendingRequests ?? []; } public CheckpointInfo? LastCheckpoint { get; set; } @@ -105,7 +121,8 @@ internal sealed class WorkflowSession : AgentSession this.SessionId, this.LastCheckpoint, this._inMemoryCheckpointManager, - this.StateBag); + this.StateBag, + this._pendingRequests); return marshaller.Marshal(info); } @@ -270,22 +287,22 @@ internal sealed class WorkflowSession : AgentSession }; /// - /// Tries to get a pending request from the state bag by content ID. + /// Tries to get a pending request by content ID. /// private ExternalRequest? TryGetPendingRequest(string contentId) => - this.StateBag.GetValue(PendingRequestKeyPrefix + contentId); + this._pendingRequests.TryGetValue(contentId, out ExternalRequest? request) ? request : null; /// - /// Adds a pending request to the state bag. + /// Adds a pending request indexed by content ID. /// private void AddPendingRequest(string contentId, ExternalRequest request) => - this.StateBag.SetValue(PendingRequestKeyPrefix + contentId, request); + this._pendingRequests[contentId] = request; /// - /// Removes a pending request from the state bag. + /// Removes a pending request by content ID. /// private void RemovePendingRequest(string contentId) => - this.StateBag.TryRemoveValue(PendingRequestKeyPrefix + contentId); + this._pendingRequests.Remove(contentId); internal async IAsyncEnumerable InvokeStageAsync( @@ -430,11 +447,13 @@ internal sealed class WorkflowSession : AgentSession string sessionId, CheckpointInfo? lastCheckpoint, InMemoryCheckpointManager? checkpointManager = null, - AgentSessionStateBag? stateBag = null) + AgentSessionStateBag? stateBag = null, + Dictionary? pendingRequests = null) { public string SessionId { get; } = sessionId; public CheckpointInfo? LastCheckpoint { get; } = lastCheckpoint; public InMemoryCheckpointManager? CheckpointManager { get; } = checkpointManager; public AgentSessionStateBag StateBag { get; } = stateBag ?? new(); + public Dictionary? PendingRequests { get; } = pendingRequests; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs index 08d1dcbcbb..5e9a86fbec 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs @@ -90,6 +90,7 @@ internal static partial class WorkflowsJsonUtilities [JsonSerializable(typeof(ChatMessage))] [JsonSerializable(typeof(ExternalRequest))] [JsonSerializable(typeof(ExternalResponse))] + [JsonSerializable(typeof(Dictionary))] [JsonSerializable(typeof(TurnToken))] // Built-in Executor State Types diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs index c2a538b302..572f858f13 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/JsonSerializationTests.cs @@ -673,6 +673,55 @@ public class JsonSerializationTests ValidateCheckpoint(retrievedCheckpoint, prototype); } + [Fact] + public void Test_SessionState_JsonRoundtrip_WithPendingRequests() + { + // Arrange + Dictionary pendingRequests = new() + { + ["call-1"] = TestExternalRequest, + ["call-2"] = ExternalRequest.Create(TestPort, "Request2", "OtherData"), + }; + + WorkflowSession.SessionState prototype = new( + sessionId: "test-session-123", + lastCheckpoint: TestParentCheckpointInfo, + pendingRequests: pendingRequests); + + // Act + WorkflowSession.SessionState result = RunJsonRoundtrip(prototype); + + // Assert + result.SessionId.Should().Be(prototype.SessionId); + result.LastCheckpoint.Should().Be(prototype.LastCheckpoint); + result.StateBag.Should().NotBeNull(); + result.PendingRequests.Should().NotBeNull() + .And.HaveCount(pendingRequests.Count); + + foreach (string key in pendingRequests.Keys) + { + result.PendingRequests.Should().ContainKey(key); + ValidateExternalRequest(result.PendingRequests![key], pendingRequests[key]); + } + } + + [Fact] + public void Test_SessionState_JsonRoundtrip_WithoutPendingRequests() + { + // Arrange + WorkflowSession.SessionState prototype = new( + sessionId: "test-session-456", + lastCheckpoint: null); + + // Act + WorkflowSession.SessionState result = RunJsonRoundtrip(prototype); + + // Assert + result.SessionId.Should().Be(prototype.SessionId); + result.LastCheckpoint.Should().BeNull(); + result.PendingRequests.Should().BeNull(); + } + /// /// Verifies that the default behavior (without AllowOutOfOrderMetadataProperties) fails /// when $type metadata is not the first property, demonstrating the PostgreSQL jsonb issue.