make getAllEntities stateType param functional with TypedEntityQueryPageable

Author: bachuv

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

@@ -183,30 +183,29 @@ public EntityQueryPageable getAllEntities() {
 
     /**
      * Returns an auto-paginating iterable over entity instances matching the specified filter criteria,
-     * with state included for typed access.
+     * with typed state access.
      * <p>
-     * This convenience overload ensures that entity state is fetched, matching the .NET SDK's
-     * {@code GetAllEntitiesAsync<T>()} pattern. Use {@link EntityMetadata#readStateAs(Class)} on
-     * each result to access the typed state.
+     * This mirrors the .NET SDK's {@code GetAllEntitiesAsync<T>()} pattern. Entity state is always
+     * included in the results and eagerly deserialized into the specified type. Each item is a
+     * {@link TypedEntityMetadata} with a {@link TypedEntityMetadata#getState()} accessor.
      * <p>
-     * Note: The provided query's {@code includeState} setting is preserved. A copy of the query
-     * is made with {@code includeState} set to {@code true} so the original query is not modified.
+     * Note: A copy of the query is made with {@code includeState} set to {@code true} so the
+     * original query is not modified.
      *
      * <pre>{@code
      * EntityQuery query = new EntityQuery().setInstanceIdStartsWith("counter");
-     * for (EntityMetadata entity : client.getEntities().getAllEntities(query, Integer.class)) {
-     *     Integer state = entity.readStateAs(Integer.class);
+     * for (TypedEntityMetadata<Integer> entity : client.getEntities().getAllEntities(query, Integer.class)) {
+     *     Integer state = entity.getState();
      *     System.out.println("Counter value: " + state);
      * }
      * }</pre>
      *
      * @param query     the query filter criteria
-     * @param stateType the expected type of the entity's state, used with
-     *                  {@link EntityMetadata#readStateAs(Class)} for deserialization
+     * @param stateType the class to deserialize each entity's state into
      * @param <T>       the entity state type
-     * @return a pageable iterable over all matching entities with state included
+     * @return a pageable iterable over all matching entities with typed state
      */
-    public <T> EntityQueryPageable getAllEntities(EntityQuery query, Class<T> stateType) {
+    public <T> TypedEntityQueryPageable<T> getAllEntities(EntityQuery query, Class<T> stateType) {
         // Create a copy with includeState=true so we don't mutate the caller's query
         EntityQuery typedQuery = new EntityQuery()
                 .setInstanceIdStartsWith(query.getInstanceIdStartsWith())
@@ -216,7 +215,8 @@ public <T> EntityQueryPageable getAllEntities(EntityQuery query, Class<T> stateT
                 .setIncludeTransient(query.isIncludeTransient())
                 .setPageSize(query.getPageSize())
                 .setContinuationToken(query.getContinuationToken());
-        return new EntityQueryPageable(typedQuery, this::queryEntities);
+        EntityQueryPageable inner = new EntityQueryPageable(typedQuery, this::queryEntities);
+        return new TypedEntityQueryPageable<>(inner, stateType);
     }
 
     /**

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

@@ -0,0 +1,74 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.microsoft.durabletask;
+
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+
+/**
+ * An auto-paginating iterable over entity query results with typed state access.
+ * <p>
+ * This class wraps an {@link EntityQueryPageable} and yields {@link TypedEntityMetadata} items
+ * with eagerly deserialized state, mirroring the .NET SDK's {@code AsyncPageable<EntityMetadata<TState>>}
+ * returned by {@code GetAllEntitiesAsync<T>()}.
+ * <p>
+ * Use {@link DurableEntityClient#getAllEntities(EntityQuery, Class)} to obtain an instance.
+ *
+ * <h3>Example:</h3>
+ * <pre>{@code
+ * EntityQuery query = new EntityQuery().setInstanceIdStartsWith("counter");
+ * for (TypedEntityMetadata<Integer> entity : client.getEntities().getAllEntities(query, Integer.class)) {
+ *     Integer state = entity.getState();
+ *     System.out.println("Counter value: " + state);
+ * }
+ * }</pre>
+ *
+ * @param <T> the entity state type
+ */
+public final class TypedEntityQueryPageable<T> implements Iterable<TypedEntityMetadata<T>> {
+    private final EntityQueryPageable inner;
+    private final Class<T> stateType;
+
+    /**
+     * Creates a new {@code TypedEntityQueryPageable}.
+     *
+     * @param inner     the underlying pageable that fetches raw entity metadata
+     * @param stateType the class to deserialize each entity's state into
+     */
+    TypedEntityQueryPageable(EntityQueryPageable inner, Class<T> stateType) {
+        this.inner = inner;
+        this.stateType = stateType;
+    }
+
+    /**
+     * Returns an iterator over individual {@link TypedEntityMetadata} items with eagerly
+     * deserialized state, automatically fetching subsequent pages as needed.
+     *
+     * @return an iterator over all matching entities with typed state
+     */
+    @Override
+    public Iterator<TypedEntityMetadata<T>> iterator() {
+        return new TypedEntityItemIterator(inner.iterator());
+    }
+
+    private class TypedEntityItemIterator implements Iterator<TypedEntityMetadata<T>> {
+        private final Iterator<EntityMetadata> delegate;
+
+        TypedEntityItemIterator(Iterator<EntityMetadata> delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return delegate.hasNext();
+        }
+
+        @Override
+        public TypedEntityMetadata<T> next() {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
+            }
+            return new TypedEntityMetadata<>(delegate.next(), stateType);
+        }
+    }
+}