addressed pr comments

Author: bachuv

client/src/main/java/com/microsoft/durabletask/AbstractTaskEntity.java

@@ -204,6 +204,9 @@ private Object handleImplicitOperations(TaskEntityOperation operation) {
     }
 
     // region Re-entrant self-dispatch
+    // NOTE: Cross-SDK alignment — dispatch() is a Java-specific addition enabling re-entrant
+    // self-calls for operation composition. Not present in the .NET SDK. Consider proposing
+    // for cross-SDK alignment.
 
     /**
      * Dispatches an operation to this entity instance synchronously (re-entrant self-call).

client/src/main/java/com/microsoft/durabletask/CallEntityOptions.java

@@ -8,6 +8,8 @@
 /**
  * Options for calling a durable entity and waiting for a response.
  */
+// NOTE: Cross-SDK alignment — setTimeout(Duration) is a Java-specific addition using a timer-based
+// implementation. Not present in the .NET SDK. Consider proposing for cross-SDK alignment.
 public final class CallEntityOptions {
     private Duration timeout;
 

client/src/main/java/com/microsoft/durabletask/EntityMetadata.java

@@ -115,8 +115,8 @@ public String getSerializedState() {
      * Gets whether this metadata response includes the entity state.
      * <p>
      * Queries can exclude the state of the entity from the metadata that is retrieved.
-     * When this returns {@code false}, {@link #getSerializedState()} and {@link #readStateAs(Class)}
-     * will return {@code null}.
+     * When this returns {@code false}, calling {@link #readStateAs(Class)} will throw
+     * {@link IllegalStateException}. Use this method to check before accessing state.
      *
      * @return {@code true} if state was requested and included in this metadata
      */
@@ -140,10 +140,17 @@ DataConverter getDataConverter() {
      *
      * @param stateType the class to deserialize the state into
      * @param <T>       the target type
-     * @return the deserialized state, or {@code null} if no state is available
+     * @return the deserialized state, or {@code null} if the entity has no state
+     * @throws IllegalStateException if state was not included in this metadata
+     *         (i.e., {@link #isIncludesState()} returns {@code false})
      */
     @Nullable
     public <T> T readStateAs(Class<T> stateType) {
+        if (!this.includesState) {
+            throw new IllegalStateException(
+                    "Entity state was not included in this metadata response. " +
+                    "Check isIncludesState() before accessing state, or request state inclusion when querying.");
+        }
         if (this.serializedState == null) {
             return null;
         }

client/src/main/java/com/microsoft/durabletask/GrpcDurableEntityClient.java

@@ -56,6 +56,12 @@ public void signalEntity(
             builder.setScheduledTime(ts);
         }
 
+        // Capture and propagate distributed trace context (matching .NET SDK pattern)
+        TraceContext traceContext = TracingHelper.getCurrentTraceContext();
+        if (traceContext != null) {
+            builder.setParentTraceContext(traceContext);
+        }
+
         this.sidecarClient.signalEntity(builder.build());
     }
 

client/src/main/java/com/microsoft/durabletask/TaskEntityState.java

@@ -50,6 +50,20 @@ public <T> T getState(Class<T> stateType) {
         return this.dataConverter.deserialize(this.serializedState, stateType);
     }
 
+    /**
+     * Gets the current entity state, deserialized to the specified type,
+     * returning the provided default value if no state exists.
+     *
+     * @param stateType    the class to deserialize the state into
+     * @param defaultValue the value to return if no state exists
+     * @param <T>          the type of the state
+     * @return the deserialized state, or {@code defaultValue} if no state exists
+     */
+    public <T> T getState(Class<T> stateType, T defaultValue) {
+        T state = getState(stateType);
+        return state != null ? state : defaultValue;
+    }
+
     /**
      * Sets the entity state. The state will be serialized using the configured {@link DataConverter}.
      *

client/src/main/java/com/microsoft/durabletask/TypedEntityMetadata.java

@@ -33,7 +33,7 @@ public final class TypedEntityMetadata<T> extends EntityMetadata {
     /**
      * Creates a new {@code TypedEntityMetadata} from an existing {@link EntityMetadata} and a state type.
      * <p>
-     * The state is eagerly deserialized from the metadata's serialized state.
+     * The state is eagerly deserialized from the metadata's serialized state when state is included.
      *
      * @param source    the source metadata to wrap
      * @param stateType the class to deserialize the state into
@@ -48,18 +48,25 @@ public final class TypedEntityMetadata<T> extends EntityMetadata {
                 source.isIncludesState(),
                 source.getDataConverter());
         this.stateType = stateType;
-        this.state = source.readStateAs(stateType);
+        this.state = source.isIncludesState() ? source.readStateAs(stateType) : null;
     }
 
     /**
      * Gets the deserialized entity state.
      * <p>
-     * Returns {@code null} if the entity has no state or if state was not included in the query.
+     * Returns {@code null} if the entity has no state.
      *
      * @return the deserialized state, or {@code null}
+     * @throws IllegalStateException if state was not included in this metadata
+     *         (i.e., {@link #isIncludesState()} returns {@code false})
      */
     @Nullable
     public T getState() {
+        if (!this.isIncludesState()) {
+            throw new IllegalStateException(
+                    "Entity state was not included in this metadata response. " +
+                    "Check isIncludesState() before accessing state, or request state inclusion when querying.");
+        }
         return this.state;
     }
 

client/src/test/java/com/microsoft/durabletask/EntityMetadataTest.java

@@ -60,12 +60,20 @@ void readStateAs_deserializesStringState() {
     @Test
     void readStateAs_nullState_returnsNull() {
         EntityMetadata metadata = new EntityMetadata(
-                "@counter@c1", Instant.EPOCH, 0, null, null, false, dataConverter);
+                "@counter@c1", Instant.EPOCH, 0, null, null, true, dataConverter);
 
         Integer state = metadata.readStateAs(Integer.class);
         assertNull(state);
     }
 
+    @Test
+    void readStateAs_throwsWhenStateNotIncluded() {
+        EntityMetadata metadata = new EntityMetadata(
+                "@counter@c1", Instant.EPOCH, 0, null, null, false, dataConverter);
+
+        assertThrows(IllegalStateException.class, () -> metadata.readStateAs(Integer.class));
+    }
+
     @Test
     void lockedBy_nullWhenNotLocked() {
         EntityMetadata metadata = new EntityMetadata(

client/src/test/java/com/microsoft/durabletask/TaskEntityStateTest.java

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.microsoft.durabletask;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * Unit tests for {@link TaskEntityState}.
+ */
+public class TaskEntityStateTest {
+
+    private static final DataConverter dataConverter = new JacksonDataConverter();
+
+    @Test
+    void getState_returnsDeserializedValue() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "42");
+        assertEquals(42, state.getState(Integer.class));
+    }
+
+    @Test
+    void getState_returnsNullWhenNoState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, null);
+        assertNull(state.getState(Integer.class));
+    }
+
+    @Test
+    void getStateWithDefault_returnsValueWhenStateExists() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "42");
+        assertEquals(42, state.getState(Integer.class, 0));
+    }
+
+    @Test
+    void getStateWithDefault_returnsDefaultWhenNoState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, null);
+        assertEquals(0, state.getState(Integer.class, 0));
+    }
+
+    @Test
+    void getStateWithDefault_returnsDefaultAfterDelete() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "42");
+        state.deleteState();
+        assertEquals(99, state.getState(Integer.class, 99));
+    }
+
+    @Test
+    void hasState_trueWhenStateExists() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "\"hello\"");
+        assertTrue(state.hasState());
+    }
+
+    @Test
+    void hasState_falseWhenNoState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, null);
+        assertFalse(state.hasState());
+    }
+
+    @Test
+    void setState_updatesState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, null);
+        state.setState(100);
+        assertTrue(state.hasState());
+        assertEquals(100, state.getState(Integer.class));
+    }
+
+    @Test
+    void deleteState_removesState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "42");
+        state.deleteState();
+        assertFalse(state.hasState());
+        assertNull(state.getState(Integer.class));
+    }
+
+    @Test
+    void rollback_restoresCommittedState() {
+        TaskEntityState state = new TaskEntityState(dataConverter, "42");
+        state.commit();
+        state.setState(100);
+        assertEquals(100, state.getState(Integer.class));
+        state.rollback();
+        assertEquals(42, state.getState(Integer.class));
+    }
+}

client/src/test/java/com/microsoft/durabletask/TypedEntityMetadataTest.java

@@ -37,12 +37,21 @@ void getState_deserializesStringState() {
     @Test
     void getState_nullWhenNoState() {
         EntityMetadata base = new EntityMetadata(
-                "@counter@c1", Instant.EPOCH, 0, null, null, false, dataConverter);
+                "@counter@c1", Instant.EPOCH, 0, null, null, true, dataConverter);
 
         TypedEntityMetadata<Integer> typed = new TypedEntityMetadata<>(base, Integer.class);
         assertNull(typed.getState());
     }
 
+    @Test
+    void getState_throwsWhenStateNotIncluded() {
+        EntityMetadata base = new EntityMetadata(
+                "@counter@c1", Instant.EPOCH, 0, null, null, false, dataConverter);
+
+        TypedEntityMetadata<Integer> typed = new TypedEntityMetadata<>(base, Integer.class);
+        assertThrows(IllegalStateException.class, typed::getState);
+    }
+
     @Test
     void getStateType_returnsCorrectClass() {
         EntityMetadata base = new EntityMetadata(