Adding support for events & shared state in durable workflows.

2026-06-16 21:04:09 +08:00 · 2026-02-17 14:41:11 -08:00
parent b62b1f2191
commit 8ffe7e6092
32 changed files with 2128 additions and 96 deletions
@@ -11,11 +11,20 @@ namespace Microsoft.Agents.AI.DurableTask.Workflows;
 /// A workflow context for durable activity execution.
 /// </summary>
 /// <remarks>
-/// Some of the methods are returning default for this version. Those method will be updated with real implementations in follow up PRs.
+/// State is passed in from the orchestration and updates are collected for return.
+/// Events emitted during execution are collected and returned to the orchestration
+/// as part of the activity output for streaming to callers.
 /// </remarks>
 [DebuggerDisplay("Executor = {_executor.Id}, StateEntries = {_initialState.Count}")]
 internal sealed class DurableActivityContext : IWorkflowContext
 {
+    /// <summary>
+    /// The default scope name used when no explicit scope is specified.
+    /// Scopes partition shared state into logical namespaces so that different
+    /// parts of a workflow can manage their state keys independently.
+    /// </summary>
+    private const string DefaultScopeName = "__default__";
+
    private readonly Dictionary<string, string> _initialState;
    private readonly Executor _executor;

@@ -33,12 +42,40 @@ internal sealed class DurableActivityContext : IWorkflowContext
    /// <summary>
    /// Gets the messages sent during activity execution via <see cref="SendMessageAsync"/>.
    /// </summary>
-    internal List<SentMessageInfo> SentMessages { get; } = [];
+    internal List<TypedPayload> SentMessages { get; } = [];
+
+    /// <summary>
+    /// Gets the events that were added during activity execution.
+    /// </summary>
+    internal List<WorkflowEvent> Events { get; } = [];
+
+    /// <summary>
+    /// Gets the state updates made during activity execution.
+    /// </summary>
+    internal Dictionary<string, string?> StateUpdates { get; } = [];
+
+    /// <summary>
+    /// Gets the scopes that were cleared during activity execution.
+    /// </summary>
+    internal HashSet<string> ClearedScopes { get; } = [];
+
+    /// <summary>
+    /// Gets a value indicating whether the executor requested a workflow halt.
+    /// </summary>
+    internal bool HaltRequested { get; private set; }

    /// <inheritdoc/>
    public ValueTask AddEventAsync(
        WorkflowEvent workflowEvent,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        if (workflowEvent is not null)
+        {
+            this.Events.Add(workflowEvent);
+        }
+
+        return default;
+    }

    /// <inheritdoc/>
    [UnconditionalSuppressMessage("AOT", "IL3050", Justification = "Serializing workflow message types registered at startup.")]
@@ -51,10 +88,10 @@ internal sealed class DurableActivityContext : IWorkflowContext
        if (message is not null)
        {
            Type messageType = message.GetType();
-            this.SentMessages.Add(new SentMessageInfo
+            this.SentMessages.Add(new TypedPayload
            {
-                Message = JsonSerializer.Serialize(message, messageType),
-                TypeName = messageType.FullName ?? messageType.Name
+                Data = JsonSerializer.Serialize(message, messageType, DurableSerialization.Options),
+                TypeName = messageType.AssemblyQualifiedName
            });
        }

@@ -64,44 +101,193 @@ internal sealed class DurableActivityContext : IWorkflowContext
    /// <inheritdoc/>
    public ValueTask YieldOutputAsync(
        object output,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        if (output is not null)
+        {
+            Type outputType = output.GetType();
+            if (!this._executor.CanOutput(outputType))
+            {
+                throw new InvalidOperationException(
+                    $"Cannot output object of type {outputType.Name}. " +
+                    $"Expecting one of [{string.Join(", ", this._executor.OutputTypes)}].");
+            }
+
+            this.Events.Add(new WorkflowOutputEvent(output, this._executor.Id));
+        }
+
+        return default;
+    }

    /// <inheritdoc/>
-    public ValueTask RequestHaltAsync() => default;
+    public ValueTask RequestHaltAsync()
+    {
+        this.HaltRequested = true;
+        this.Events.Add(new DurableHaltRequestedEvent(this._executor.Id));
+        return default;
+    }

    /// <inheritdoc/>
    public ValueTask<T?> ReadStateAsync<T>(
        string key,
        string? scopeName = null,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        string scopeKey = GetScopeKey(scopeName, key);
+        string normalizedScope = scopeName ?? DefaultScopeName;
+        bool scopeCleared = this.ClearedScopes.Contains(normalizedScope);
+
+        // Local updates take priority over initial state.
+        if (this.StateUpdates.TryGetValue(scopeKey, out string? updated))
+        {
+            return DeserializeStateAsync<T>(updated);
+        }
+
+        // If scope was cleared, ignore initial state
+        if (scopeCleared)
+        {
+            return ValueTask.FromResult<T?>(default);
+        }
+
+        // Fall back to initial state passed from orchestration
+        if (this._initialState.TryGetValue(scopeKey, out string? initial))
+        {
+            return DeserializeStateAsync<T>(initial);
+        }
+
+        return ValueTask.FromResult<T?>(default);
+    }

    /// <inheritdoc/>
-    public ValueTask<T> ReadOrInitStateAsync<T>(
+    public async ValueTask<T> ReadOrInitStateAsync<T>(
        string key,
        Func<T> initialStateFactory,
        string? scopeName = null,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        T? value = await this.ReadStateAsync<T>(key, scopeName, cancellationToken).ConfigureAwait(false);
+
+        if (value is not null)
+        {
+            return value;
+        }
+
+        T initialValue = initialStateFactory();
+        await this.QueueStateUpdateAsync(key, initialValue, scopeName, cancellationToken).ConfigureAwait(false);
+        return initialValue;
+    }

    /// <inheritdoc/>
    public ValueTask<HashSet<string>> ReadStateKeysAsync(
        string? scopeName = null,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        string scopePrefix = GetScopePrefix(scopeName);
+        int scopePrefixLength = scopePrefix.Length;
+        HashSet<string> keys = new(StringComparer.Ordinal);
+
+        bool scopeCleared = scopeName is null
+            ? this.ClearedScopes.Contains(DefaultScopeName)
+            : this.ClearedScopes.Contains(scopeName);
+
+        // Start with keys from initial state (skip if scope was cleared)
+        if (!scopeCleared)
+        {
+            foreach (string stateKey in this._initialState.Keys)
+            {
+                if (stateKey.StartsWith(scopePrefix, StringComparison.Ordinal))
+                {
+                    keys.Add(stateKey[scopePrefixLength..]);
+                }
+            }
+        }
+
+        // Merge local updates: add if non-null, remove if null (deleted)
+        foreach (KeyValuePair<string, string?> update in this.StateUpdates)
+        {
+            if (!update.Key.StartsWith(scopePrefix, StringComparison.Ordinal))
+            {
+                continue;
+            }
+
+            string key = update.Key[scopePrefixLength..];
+            if (update.Value is not null)
+            {
+                keys.Add(key);
+            }
+            else
+            {
+                keys.Remove(key);
+            }
+        }
+
+        return ValueTask.FromResult(keys);
+    }

    /// <inheritdoc/>
    public ValueTask QueueStateUpdateAsync<T>(
        string key,
        T? value,
        string? scopeName = null,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        string scopeKey = GetScopeKey(scopeName, key);
+        this.StateUpdates[scopeKey] = value is null ? null : SerializeState(value);
+        return default;
+    }

    /// <inheritdoc/>
    public ValueTask QueueClearScopeAsync(
        string? scopeName = null,
-        CancellationToken cancellationToken = default) => default;
+        CancellationToken cancellationToken = default)
+    {
+        this.ClearedScopes.Add(scopeName ?? DefaultScopeName);
+
+        // Remove any pending updates in this scope (snapshot keys to allow removal during iteration)
+        string scopePrefix = GetScopePrefix(scopeName);
+        foreach (string key in this.StateUpdates.Keys.ToArray())
+        {
+            if (key.StartsWith(scopePrefix, StringComparison.Ordinal))
+            {
+                this.StateUpdates.Remove(key);
+            }
+        }
+
+        return default;
+    }

    /// <inheritdoc/>
    public IReadOnlyDictionary<string, string>? TraceContext => null;

    /// <inheritdoc/>
    public bool ConcurrentRunsEnabled => false;
+
+    private static string GetScopeKey(string? scopeName, string key)
+        => $"{GetScopePrefix(scopeName)}{key}";
+
+    /// <summary>
+    /// Returns the key prefix for the given scope. Scopes partition shared state
+    /// into logical namespaces, allowing different workflow executors to manage
+    /// their state keys independently. When no scope is specified, the
+    /// <see cref="DefaultScopeName"/> is used.
+    /// </summary>
+    private static string GetScopePrefix(string? scopeName)
+        => scopeName is null ? $"{DefaultScopeName}:" : $"{scopeName}:";
+
+    [UnconditionalSuppressMessage("AOT", "IL3050", Justification = "Serializing workflow state types.")]
+    [UnconditionalSuppressMessage("Trimming", "IL2026", Justification = "Serializing workflow state types.")]
+    private static string SerializeState<T>(T value)
+        => JsonSerializer.Serialize(value, DurableSerialization.Options);
+
+    [UnconditionalSuppressMessage("AOT", "IL3050", Justification = "Deserializing workflow state types.")]
+    [UnconditionalSuppressMessage("Trimming", "IL2026", Justification = "Deserializing workflow state types.")]
+    private static ValueTask<T?> DeserializeStateAsync<T>(string? json)
+    {
+        if (json is null)
+        {
+            return ValueTask.FromResult<T?>(default);
+        }
+
+        return ValueTask.FromResult(JsonSerializer.Deserialize<T>(json, DurableSerialization.Options));
+    }
 }