Address PR comments

src/Abstractions/Entities/TaskEntity.cs

@@ -77,10 +77,10 @@ public interface ITaskEntity
 public abstract class TaskEntity<TState> : ITaskEntity
 {
     /// <summary>
-    /// Gets a value indicating whether dispatching operations to <see cref="State"/> is allowed. State dispatch will
-    /// only be attempted if entity-level dispatch does not succeed. Default is <c>true</c>.
+    /// Gets a value indicating whether dispatching operations to <see cref="State"/> is allowed. State dispatch
+    /// will only be attempted if entity-level dispatch does not succeed. Default is <c>false</c>.
     /// </summary>
-    protected bool AllowStateDispatch => true;
+    protected virtual bool AllowStateDispatch => false;
 
     /// <summary>
     /// Gets or sets the state for this entity.
@@ -90,7 +90,7 @@ public abstract class TaskEntity<TState> : ITaskEntity
     /// will be persisted to <see cref="TaskEntityContext.SetState(object?)"/> when the operation completes. Deleting
     /// entity state can be accomplished by setting this to default(<typeparamref name="TState"/>).
     /// </remarks>
-    protected TState State { get; set; } = default!; // leave null-checks to end implementation.
+    protected TState? State { get; set; }
 
     /// <summary>
     /// Gets the entity operation.
@@ -107,13 +107,36 @@ public abstract class TaskEntity<TState> : ITaskEntity
     {
         this.Operation = Check.NotNull(operation);
         object? state = operation.Context.GetState(typeof(TState));
-        this.State = state is null ? default! : (TState)state;
+        this.State = state is null ? this.InitializeState() : (TState)state;
         if (!operation.TryDispatch(this, out object? result, out Type returnType)
-            && (this.AllowStateDispatch && !operation.TryDispatch(this.State, out result, out returnType)))
+            && !this.TryDispatchState(out result, out returnType))
         {
             throw new NotSupportedException($"No suitable method found for entity operation '{operation}'.");
         }
 
         return TaskEntityHelpers.UnwrapAsync(this.Context, () => this.State, result, returnType);
     }
+
+    /// <summary>
+    /// Initializes the entity state.
+    /// </summary>
+    /// <returns>The entity state.</returns>
+    protected virtual TState? InitializeState() => default;
+
+    bool TryDispatchState(out object? result, out Type returnType)
+    {
+        if (!this.AllowStateDispatch)
+        {
+            result = null;
+            returnType = typeof(void);
+            return false;
+        }
+
+        if (this.State is null)
+        {
+            throw new InvalidOperationException("Attempting to dispatch to state, but entity state is null.");
+        }
+
+        return this.Operation.TryDispatch(this.State, out result, out returnType);
+    }
 }

src/Abstractions/Entities/TaskEntityContext.cs

@@ -53,7 +53,8 @@ public abstract void StartOrchestration(
         TaskName name, object? input = null, StartOrchestrationOptions? options = null);
 
     /// <summary>
-    /// Gets the current state for the entity this context is for.
+    /// Gets the current state for the entity this context is for. This will return <c>null</c> if no state is present,
+    /// regardless if <paramref name="type"/> is a value-type or not.
     /// </summary>
     /// <param name="type">The type to retrieve the state as.</param>
     /// <returns>The entity state.</returns>

src/Abstractions/Entities/TaskEntityHelpers.cs

@@ -11,7 +11,7 @@ namespace Microsoft.DurableTask.Entities;
 static class TaskEntityHelpers
 {
     /// <summary>
-    /// Unwraps a dispatched result for a <see cref="TaskEntityOperation"/> into a <see cref="ValueTask{Object}"/>
+    /// Unwraps a dispatched result for a <see cref="TaskEntityOperation"/> into a <see cref="ValueTask{Object}"/>.
     /// </summary>
     /// <param name="context">The entity context.</param>
     /// <param name="state">Delegate to resolve new state for the entity.</param>

src/Shared/Core/TaskExtensions.cs

@@ -24,7 +24,7 @@ static class TaskExtensions
         Type t = task.GetType();
         if (t.IsGenericType)
         {
-            return (T)t.GetProperty("Result", BindingFlags.Public | BindingFlags.Instance).GetValue(task);
+            return (T)t.GetProperty("Result", BindingFlags.Public | BindingFlags.Instance)!.GetValue(task)!;
         }
 
         return default;

test/Abstractions.Tests/Entities/TaskEntityTests.cs → test/Abstractions.Tests/Entities/EntityTaskEntityTests.cs

@@ -6,7 +6,7 @@
 
 namespace Microsoft.DurableTask.Entities.Tests;
 
-public partial class TaskEntityTests
+public class EntityTaskEntityTests
 {
     [Theory]
     [InlineData("doesNotExist")] // method does not exist.
@@ -125,6 +125,8 @@ public async Task DefaultValue_Input_Succeeds()
         result.Should().BeOfType<string>().Which.Should().Be("not-default");
     }
 
+#pragma warning disable CA1822 // Mark members as static
+#pragma warning disable IDE0060 // Remove unused parameter
     class TestEntity : TaskEntity<int>
     {
         public static string StaticMethod() => throw new NotImplementedException();
@@ -258,4 +260,6 @@ int Get(Optional<TaskEntityContext> context)
             return this.State;
         }
     }
+#pragma warning restore IDE0060 // Remove unused parameter
+#pragma warning restore CA1822 // Mark members as static
 }