.NET: Harden durable mailbox delivery and workflow output provenance

dotnet/src/Microsoft.Agents.AI.DurableTask/AgentEntityResultExpirySchedule.cs

@@ -0,0 +1,133 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using Microsoft.Agents.AI.DurableTask.State;
+
+namespace Microsoft.Agents.AI.DurableTask;
+
+/// <summary>An optional runtime profile, not part of the shared mailbox schema or history binding.</summary>
+internal sealed class AgentEntityResultExpirySchedule
+{
+    internal const string ExtensionName = "Microsoft.Agents.AI.DurableTask.resultExpiry";
+    private readonly Dictionary<string, JsonElement> _properties;
+
+    private AgentEntityResultExpirySchedule(
+        Dictionary<string, JsonElement> properties,
+        AgentEntityResultExpirationCheck? pending)
+    {
+        this._properties = properties;
+        this.Pending = pending;
+    }
+
+    public AgentEntityResultExpirationCheck? Pending { get; }
+
+    public static AgentEntityResultExpirySchedule? Read(DurableAgentState state, string entityId)
+    {
+        if (state.ExtensionData?.TryGetValue(ExtensionName, out JsonElement profile) != true)
+        {
+            return null;
+        }
+
+        if (profile.ValueKind != JsonValueKind.Object)
+        {
+            throw InvalidProfile();
+        }
+
+        Dictionary<string, JsonElement> properties = new(StringComparer.Ordinal);
+        foreach (JsonProperty property in profile.EnumerateObject())
+        {
+            if (!properties.TryAdd(property.Name, property.Value.Clone()))
+            {
+                throw InvalidProfile();
+            }
+        }
+
+        if (!properties.TryGetValue("version", out JsonElement version) ||
+            version.ValueKind != JsonValueKind.Number || !version.TryGetInt32(out int number) || number != 1 ||
+            !properties.TryGetValue("entityId", out JsonElement identity) ||
+            identity.ValueKind != JsonValueKind.String || identity.GetString() != entityId ||
+            !properties.TryGetValue("scheduledResultExpiryUtc", out JsonElement deadline) ||
+            !properties.TryGetValue("token", out JsonElement token))
+        {
+            throw InvalidProfile();
+        }
+
+        AgentEntityResultExpirationCheck? pending = null;
+        if (deadline.ValueKind != JsonValueKind.Null || token.ValueKind != JsonValueKind.Null)
+        {
+            if (deadline.ValueKind != JsonValueKind.String || !deadline.TryGetDateTimeOffset(out DateTimeOffset scheduledTime) ||
+                scheduledTime.Offset != TimeSpan.Zero ||
+                token.ValueKind != JsonValueKind.String || !Guid.TryParseExact(token.GetString(), "N", out Guid generation) ||
+                generation == Guid.Empty)
+            {
+                throw InvalidProfile();
+            }
+
+            pending = new AgentEntityResultExpirationCheck(scheduledTime, token.GetString(), entityId);
+        }
+
+        return new AgentEntityResultExpirySchedule(properties, pending);
+    }
+
+    public static DurableAgentState Write(
+        DurableAgentState state,
+        string entityId,
+        AgentEntityResultExpirySchedule? previous,
+        AgentEntityResultExpirationCheck? pending)
+    {
+        if (previous?.Pending == pending)
+        {
+            return state;
+        }
+
+        using MemoryStream buffer = new();
+        using (Utf8JsonWriter writer = new(buffer))
+        {
+            writer.WriteStartObject();
+            writer.WriteNumber("version", 1);
+            writer.WriteString("entityId", entityId);
+            if (pending is null)
+            {
+                writer.WriteNull("scheduledResultExpiryUtc");
+                writer.WriteNull("token");
+            }
+            else
+            {
+                writer.WriteString("scheduledResultExpiryUtc", pending.ScheduledTime);
+                writer.WriteString("token", pending.Token);
+            }
+
+            if (previous is not null)
+            {
+                foreach ((string key, JsonElement value) in previous._properties)
+                {
+                    if (key is not ("version" or "entityId" or "scheduledResultExpiryUtc" or "token"))
+                    {
+                        writer.WritePropertyName(key);
+                        value.WriteTo(writer);
+                    }
+                }
+            }
+
+            writer.WriteEndObject();
+        }
+
+        using JsonDocument profile = JsonDocument.Parse(buffer.ToArray());
+        Dictionary<string, JsonElement> extensions = state.ExtensionData is null
+            ? new(StringComparer.Ordinal)
+            : new(state.ExtensionData, StringComparer.Ordinal);
+        extensions[ExtensionName] = profile.RootElement.Clone();
+        return new DurableAgentState
+        {
+            SchemaVersion = state.SchemaVersion,
+            MailboxWritesAuthorized = state.MailboxWritesAuthorized,
+            Data = state.Data,
+            ExtensionData = extensions,
+            UnknownProperties = state.UnknownProperties,
+        };
+    }
+
+    private static InvalidOperationException InvalidProfile() =>
+        new($"The '{ExtensionName}' runtime profile is malformed, unsupported, or belongs to another entity. " +
+            "Result-expiry scheduling cannot safely continue; preserve the profile and use a compatible writer.");
+}

dotnet/src/Microsoft.Agents.AI.DurableTask/AgentRunHandle.cs

@@ -14,15 +14,20 @@ internal sealed class AgentRunHandle
 {
     private readonly DurableTaskClient _client;
     private readonly ILogger _logger;
+    private readonly TimeProvider _timeProvider;
 
     internal AgentRunHandle(
         DurableTaskClient client,
         ILogger logger,
         AgentSessionId sessionId,
-        string correlationId)
+        string correlationId,
+        TimeProvider? timeProvider = null)
     {
+        ArgumentException.ThrowIfNullOrWhiteSpace(correlationId);
+
         this._client = client;
         this._logger = logger;
+        this._timeProvider = timeProvider ?? TimeProvider.System;
         this.SessionId = sessionId;
         this.CorrelationId = correlationId;
     }
@@ -39,12 +44,19 @@ internal AgentRunHandle(
 
     /// <summary>
     /// Reads the agent response for this request by polling the entity state until the response is found.
-    /// Uses an exponential backoff polling strategy with a maximum interval of 1 second.
+    /// Uses an exponential backoff polling strategy with a maximum interval of 3 seconds.
     /// </summary>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The agent response corresponding to this request.</returns>
     /// <exception cref="InvalidOperationException">Thrown when the response is not found after polling.</exception>
     public async Task<AgentResponse> ReadAgentResponseAsync(CancellationToken cancellationToken = default)
+    {
+        DurableAgentRunOutcome outcome = await this.ReadAgentOutcomeAsync(cancellationToken);
+        return outcome.GetResponse(this.CorrelationId);
+    }
+
+    internal async Task<DurableAgentRunOutcome> ReadAgentOutcomeAsync(
+        CancellationToken cancellationToken = default)
     {
         TimeSpan pollInterval = TimeSpan.FromMilliseconds(50); // Start with 50ms
         TimeSpan maxPollInterval = TimeSpan.FromSeconds(3); // Maximum 3 seconds
@@ -59,17 +71,29 @@ public async Task<AgentResponse> ReadAgentResponseAsync(CancellationToken cancel
                 cancellation: cancellationToken);
             DurableAgentState? state = entityResponse?.State;
 
-            if (state?.Data.ConversationHistory is not null)
+            if (state is not null)
             {
-                // Look for an agent response with matching CorrelationId
-                DurableAgentStateResponse? response = state.Data.ConversationHistory
-                    .OfType<DurableAgentStateResponse>()
-                    .FirstOrDefault(r => r.CorrelationId == this.CorrelationId);
+                DurableAgentRunOutcome outcome;
+                try
+                {
+                    outcome = DurableAgentStateOutcomeResolver.Resolve(
+                        state,
+                        this.CorrelationId,
+                        this._timeProvider.GetUtcNow());
+                }
+                catch (DurableAgentStateCorruptionException exception)
+                {
+                    this._logger.LogDurableOutcomeStateCorruption(
+                        exception,
+                        this.SessionId,
+                        this.CorrelationId);
+                    throw;
+                }
 
-                if (response is not null)
+                if (outcome.Kind != DurableAgentRunOutcomeKind.Pending)
                 {
                     this._logger.LogDonePollingForResponse(this.SessionId, this.CorrelationId);
-                    return response.ToResponse();
+                    return outcome;
                 }
             }
 

dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [Unreleased]
 
+- Hardened durable-agent mailbox delivery, duplicate correlation handling, working-state rollback, token-bounded durable result expiry (including zero-write stale cleanup on legacy state), and stale-safe TTL deletion scheduling; added gated real-backend state/outbox atomicity coverage with locally tested SDK failure handling while keeping schema-2 production activation inaccessible; preserved historical message boundaries, opaque schema-2 history/state profiles, committed failure metadata, and legacy response metadata through mailbox promotion; isolated untrusted response text and malformed activity/child payloads from workflow controls with all-or-nothing typed-message validation, preserved opaque child fallback string routing, and rejected unknown target type hints instead of selecting an unrelated handler ([#94](https://github.com/microsoft/agent-framework-durable-extension/pull/94))
 - Fail durable workflows with a `MaxSuperstepsExceededException` when they reach the configurable `MaxSupersteps` limit with work still queued, instead of returning a successful partial result ([#84](https://github.com/microsoft/agent-framework-durable-extension/pull/84))
 - Added passive .NET DTO, converter, validation, and source-generation support for the proposed durable agent state 2.0 contract ([tamirdresher/agent-framework-durable-extension#1](https://github.com/tamirdresher/agent-framework-durable-extension/pull/1))
 - Fixed `ConfigureDurableAgents` and `ConfigureDurableWorkflows` ignoring the `workerBuilder` or `clientBuilder` supplied to a later call when no earlier call supplied one, so the Durable Task worker and client are now registered whichever configuration call provides them. The first non-null delegate wins; later ones are still ignored so a builder passed to several calls is only applied once. Registering an agent that a workflow already referenced now promotes it to an explicitly registered agent instead of throwing, so agents and workflows can be configured in either order ([#67](https://github.com/microsoft/agent-framework-durable-extension/pull/67))

dotnet/src/Microsoft.Agents.AI.DurableTask/DefaultDurableAgentClient.cs

@@ -6,17 +6,28 @@
 
 namespace Microsoft.Agents.AI.DurableTask;
 
-internal class DefaultDurableAgentClient(DurableTaskClient client, ILoggerFactory loggerFactory) : IDurableAgentClient
+internal class DefaultDurableAgentClient(
+    DurableTaskClient client,
+    ILoggerFactory loggerFactory,
+    TimeProvider? timeProvider = null) : IDurableAgentClient
 {
     private readonly DurableTaskClient _client = client ?? throw new ArgumentNullException(nameof(client));
     private readonly ILogger _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<DefaultDurableAgentClient>();
+    private readonly TimeProvider _timeProvider = timeProvider ?? TimeProvider.System;
 
     public async Task<AgentRunHandle> RunAgentAsync(
         AgentSessionId sessionId,
         RunRequest request,
         CancellationToken cancellationToken = default)
     {
         ArgumentNullException.ThrowIfNull(request);
+        string correlationId = request.CorrelationId;
+        if (string.IsNullOrWhiteSpace(correlationId))
+        {
+            throw new ArgumentException(
+                "A non-empty correlation ID is required to run a durable agent request.",
+                nameof(request));
+        }
 
         this._logger.LogSignallingAgent(sessionId);
 
@@ -26,6 +37,6 @@ await this._client.Entities.SignalEntityAsync(
             request,
             cancellation: cancellationToken);
 
-        return new AgentRunHandle(this._client, this._logger, sessionId, request.CorrelationId);
+        return new AgentRunHandle(this._client, this._logger, sessionId, correlationId, this._timeProvider);
     }
 }

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs

@@ -142,6 +142,10 @@ protected override async Task<AgentResponse> RunCoreAsync(
         {
             throw new AgentNotRegisteredException(this._agentName, e);
         }
+        catch (Exception e) when (DurableAgentFailure.TryRestore(e, out Exception? failure))
+        {
+            throw failure;
+        }
     }
 
     /// <summary>
@@ -291,6 +295,8 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
         // the orchestration.
         AgentResponse response = await this.RunAsync(messages, session, options, cancellationToken);
 
-        return new AgentResponse<T>(response, serializerOptions) { IsWrappedInObject = isWrappedInObject };
+        AgentResponse<T> typedResponse = new(response, serializerOptions) { IsWrappedInObject = isWrappedInObject };
+        DurableAgentJsonUtilities.CopyRetainedResult(response, typedResponse);
+        return typedResponse;
     }
 }

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentFailure.cs

@@ -0,0 +1,103 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Diagnostics.CodeAnalysis;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using Microsoft.Agents.AI.DurableTask.State;
+using Microsoft.DurableTask;
+using Microsoft.DurableTask.Entities;
+
+namespace Microsoft.Agents.AI.DurableTask;
+
+// Durable Task serializes exception type, message, and inner failure, but not custom CLR
+// properties. This framework-only inner exception carries a versioned metadata snapshot.
+internal static class DurableAgentFailure
+{
+    internal static Exception CreateMetadataException(DurableAgentFailureData data, Exception? innerException = null) =>
+        new DurableAgentFailureMetadataException(
+            JsonSerializer.Serialize(data, DurableAgentJsonUtilities.JsonContext.Default.DurableAgentFailureData), innerException);
+
+    internal static bool TryRestore(Exception exception, [NotNullWhen(true)] out Exception? restored)
+    {
+        restored = null;
+        TaskFailureDetails? failure = exception switch
+        {
+            EntityOperationFailedException entityFailure => entityFailure.FailureDetails,
+            TaskFailedException taskFailure => taskFailure.FailureDetails,
+            _ => null,
+        };
+        bool terminal = failure?.ErrorType == typeof(DurableAgentTerminalException).FullName;
+        bool unavailable = failure?.ErrorType == typeof(DurableAgentResultUnavailableException).FullName;
+        if (!terminal && !unavailable)
+        {
+            return false;
+        }
+
+        TaskFailureDetails? metadata = failure!.InnerFailure;
+        if (metadata?.ErrorType != typeof(DurableAgentFailureMetadataException).FullName)
+        {
+            // Older or message-only failures still fail with a typed exception. Never try
+            // to extract a contract from their human/model-authored error message.
+            restored = terminal
+                ? new DurableAgentTerminalException(failure.ErrorMessage, exception)
+                : new DurableAgentResultUnavailableException(failure.ErrorMessage, exception);
+            return true;
+        }
+
+        try
+        {
+            DurableAgentFailureData? data = JsonSerializer.Deserialize(
+                metadata!.ErrorMessage, DurableAgentJsonUtilities.JsonContext.Default.DurableAgentFailureData);
+            if (data is null || data.Version != 1 || string.IsNullOrWhiteSpace(data.CorrelationId))
+            {
+                return false;
+            }
+
+            if (terminal && !string.IsNullOrEmpty(data.Code) && data.SerializedResponse is not null)
+            {
+                AgentResponse? response = new DurableDataConverter().Deserialize(data.SerializedResponse, typeof(AgentResponse)) as AgentResponse;
+                if (response is not null)
+                {
+                    restored = new DurableAgentTerminalException(
+                        data.CorrelationId, data.Code, failure.ErrorMessage, data.Details, response, exception);
+                }
+            }
+            else if (unavailable && data.CompletedAt is DateTimeOffset completedAt &&
+                data.Outcome is DurableAgentStateCompletionReceipt.SucceededOutcome or DurableAgentStateCompletionReceipt.FailedOutcome)
+            {
+                restored = new DurableAgentResultUnavailableException(
+                    data.CorrelationId, completedAt, data.ResultExpiresAt, data.Outcome, exception);
+            }
+        }
+        catch (JsonException)
+        {
+            // Unsupported/malformed metadata must leave the original SDK failure intact.
+        }
+        catch (InvalidOperationException)
+        {
+            // Includes invalid canonical response metadata; never degrade it to success.
+        }
+
+        return restored is not null;
+    }
+}
+
+internal sealed class DurableAgentFailureData
+{
+    public required int Version { get; init; }
+
+    public required string CorrelationId { get; init; }
+
+    public string? Code { get; init; }
+
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
+    public JsonElement Details { get; init; }
+
+    public string? SerializedResponse { get; init; }
+
+    public DateTimeOffset? CompletedAt { get; init; }
+
+    public DateTimeOffset? ResultExpiresAt { get; init; }
+
+    public string? Outcome { get; init; }
+}