Add InitializeStateAsync support for async entity state initialization

Co-authored-by: YunchuWang <[email protected]>

Author: Copilot

src/Abstractions/Entities/TaskEntity.cs

@@ -119,25 +119,40 @@ public abstract class TaskEntity<TState> : ITaskEntity
     protected TaskEntityContext Context { get; private set; } = null!;
 
     /// <inheritdoc/>
-    public ValueTask<object?> RunAsync(TaskEntityOperation operation)
+    public async ValueTask<object?> RunAsync(TaskEntityOperation operation)
     {
         Check.NotNull(operation);
         this.Context = operation.Context;
         object? state = operation.State.GetState(typeof(TState));
-        this.State = state is null ? this.InitializeState(operation) : (TState)state;
+        this.State = state is null ? await this.InitializeStateAsync(operation) : (TState)state;
         if (!operation.TryDispatch(this, out object? result, out Type returnType)
             && !this.TryDispatchState(operation, out result, out returnType))
         {
             if (TryDispatchImplicit(operation, out ValueTask<object?> task))
             {
                 // We do not go into UnwrapAsync for implicit tasks
-                return task;
+                return await task;
             }
 
             throw new NotSupportedException($"No suitable method found for entity operation '{operation}'.");
         }
 
-        return TaskEntityHelpers.UnwrapAsync(operation.State, () => this.State, result, returnType);
+        return await TaskEntityHelpers.UnwrapAsync(operation.State, () => this.State, result, returnType);
+    }
+
+    /// <summary>
+    /// Initializes the entity state asynchronously. This is only called when there is no current state for this entity.
+    /// </summary>
+    /// <param name="entityOperation">The entity operation to be executed.</param>
+    /// <returns>A task that resolves to the entity state.</returns>
+    /// <remarks>
+    /// The default implementation calls <see cref="InitializeState"/> for backward compatibility.
+    /// Override this method to perform async initialization operations, such as loading state from a database or
+    /// external storage.
+    /// </remarks>
+    protected virtual ValueTask<TState> InitializeStateAsync(TaskEntityOperation entityOperation)
+    {
+        return new ValueTask<TState>(this.InitializeState(entityOperation));
     }
 
     /// <summary>

test/Abstractions.Tests/Entities/StateTaskEntityTests.cs

@@ -172,13 +172,103 @@ public async Task ExplicitDelete_Overridden(string op)
         operation.State.GetState(typeof(TestState)).Should().BeOfType<TestState>().Which.Value.Should().Be(0);
     }
 
+    [Fact]
+    public async Task AsyncInitializeState_Called()
+    {
+        TestEntityOperation operation = new("get0", new TestEntityState(null), default);
+        AsyncInitEntity entity = new();
+
+        object? result = await entity.RunAsync(operation);
+
+        result.Should().BeOfType<int>().Which.Should().Be(42);
+        entity.InitializeAsyncCalled.Should().BeTrue();
+    }
+
+    [Fact]
+    public async Task AsyncInitializeState_WithYield_Succeeds()
+    {
+        TestEntityOperation operation = new("get0", new TestEntityState(null), default);
+        AsyncInitWithYieldEntity entity = new();
+
+        object? result = await entity.RunAsync(operation);
+
+        result.Should().BeOfType<int>().Which.Should().Be(99);
+        entity.InitializeAsyncCalled.Should().BeTrue();
+    }
+
+    [Fact]
+    public async Task AsyncInitializeState_ValueTask_Succeeds()
+    {
+        TestEntityOperation operation = new("get0", new TestEntityState(null), default);
+        AsyncInitValueTaskEntity entity = new();
+
+        object? result = await entity.RunAsync(operation);
+
+        result.Should().BeOfType<int>().Which.Should().Be(77);
+    }
+
+    [Fact]
+    public async Task SyncInitializeState_StillWorks()
+    {
+        TestEntityOperation operation = new("get0", new TestEntityState(null), default);
+        SyncInitEntity entity = new();
+
+        object? result = await entity.RunAsync(operation);
+
+        result.Should().BeOfType<int>().Which.Should().Be(55);
+        entity.SyncInitCalled.Should().BeTrue();
+    }
+
     static TestState State(int value) => new() { Value = value };
 
     class NullStateEntity : TestEntity
     {
         protected override TestState InitializeState(TaskEntityOperation entityOperation) => null!;
     }
 
+    class AsyncInitEntity : TestEntity
+    {
+        public bool InitializeAsyncCalled { get; private set; }
+
+        protected override async ValueTask<TestState> InitializeStateAsync(TaskEntityOperation entityOperation)
+        {
+            await Task.CompletedTask;
+            this.InitializeAsyncCalled = true;
+            return new TestState { Value = 42 };
+        }
+    }
+
+    class AsyncInitWithYieldEntity : TestEntity
+    {
+        public bool InitializeAsyncCalled { get; private set; }
+
+        protected override async ValueTask<TestState> InitializeStateAsync(TaskEntityOperation entityOperation)
+        {
+            await Task.Yield();
+            this.InitializeAsyncCalled = true;
+            return new TestState { Value = 99 };
+        }
+    }
+
+    class AsyncInitValueTaskEntity : TestEntity
+    {
+        protected override ValueTask<TestState> InitializeStateAsync(TaskEntityOperation entityOperation)
+        {
+            return new ValueTask<TestState>(new TestState { Value = 77 });
+        }
+    }
+
+    class SyncInitEntity : TestEntity
+    {
+        public bool SyncInitCalled { get; private set; }
+
+        protected override TestState InitializeState(TaskEntityOperation entityOperation)
+        {
+            this.SyncInitCalled = true;
+            return new TestState { Value = 55 };
+        }
+    }
+
     class TestEntity : TaskEntity<TestState>
     {
         readonly bool allowStateDispatch;