fix: normalize entity query prefix in client.getEntities() (#286)

Author: YunchuWang

packages/durabletask-js/src/client/client.ts

@@ -31,7 +31,7 @@ import { mapToRecord } from "../utils/tags.util";
 import { populateTagsMap } from "../utils/pb-helper.util";
 import { EntityInstanceId } from "../entities/entity-instance-id";
 import { EntityMetadata, createEntityMetadata, createEntityMetadataWithoutState } from "../entities/entity-metadata";
-import { EntityQuery } from "../entities/entity-query";
+import { EntityQuery, normalizeInstanceIdPrefix } from "../entities/entity-query";
 import { SignalEntityOptions } from "../entities/signal-entity-options";
 import {
   CleanEntityStorageRequest,
@@ -1103,9 +1103,10 @@ export class TaskHubGrpcClient {
       const req = new pb.QueryEntitiesRequest();
       const protoQuery = new pb.EntityQuery();
 
-      if (query?.instanceIdStartsWith) {
+      const normalizedPrefix = normalizeInstanceIdPrefix(query?.instanceIdStartsWith);
+      if (normalizedPrefix) {
         const prefix = new StringValue();
-        prefix.setValue(query.instanceIdStartsWith);
+        prefix.setValue(normalizedPrefix);
         protoQuery.setInstanceidstartswith(prefix);
       }
 

packages/durabletask-js/test/entity-client.spec.ts

@@ -395,3 +395,96 @@ describe("EntityInstanceId.fromString", () => {
     expect(() => EntityInstanceId.fromString("@onlyname")).toThrow();
   });
 });
+
+describe("getEntities query normalization", () => {
+  const { TaskHubGrpcClient } = require("../src");
+
+  function getStub(client: InstanceType<typeof TaskHubGrpcClient>): any {
+    return (client as any)._stub;
+  }
+
+  function mockQueryEntities(
+    client: InstanceType<typeof TaskHubGrpcClient>,
+    captureRequest: (req: pb.QueryEntitiesRequest) => void,
+  ): void {
+    getStub(client).queryEntities = (req: pb.QueryEntitiesRequest, _metadata: any, callback: any) => {
+      captureRequest(req);
+      const res = new pb.QueryEntitiesResponse();
+      callback(null, res);
+      return {};
+    };
+  }
+
+  it("should normalize mixed-case instanceIdStartsWith prefix", async () => {
+    const client = new TaskHubGrpcClient({ hostAddress: "localhost:4001" });
+    let capturedReq: pb.QueryEntitiesRequest | undefined;
+
+    mockQueryEntities(client, (req) => { capturedReq = req; });
+
+    const query: EntityQuery = { instanceIdStartsWith: "Counter" };
+    const pageable = client.getEntities(query);
+
+    // Consume the first page to trigger the gRPC call
+    for await (const _entity of pageable) {
+      // no-op
+    }
+
+    expect(capturedReq).toBeDefined();
+    const protoPrefix = capturedReq!.getQuery()?.getInstanceidstartswith()?.getValue();
+    expect(protoPrefix).toBe("@counter");
+  });
+
+  it("should normalize prefix with name and key separator", async () => {
+    const client = new TaskHubGrpcClient({ hostAddress: "localhost:4001" });
+    let capturedReq: pb.QueryEntitiesRequest | undefined;
+
+    mockQueryEntities(client, (req) => { capturedReq = req; });
+
+    const query: EntityQuery = { instanceIdStartsWith: "Counter@User-123" };
+    const pageable = client.getEntities(query);
+
+    for await (const _entity of pageable) {
+      // no-op
+    }
+
+    expect(capturedReq).toBeDefined();
+    const protoPrefix = capturedReq!.getQuery()?.getInstanceidstartswith()?.getValue();
+    expect(protoPrefix).toBe("@counter@User-123");
+  });
+
+  it("should not set prefix when instanceIdStartsWith is undefined", async () => {
+    const client = new TaskHubGrpcClient({ hostAddress: "localhost:4001" });
+    let capturedReq: pb.QueryEntitiesRequest | undefined;
+
+    mockQueryEntities(client, (req) => { capturedReq = req; });
+
+    const query: EntityQuery = {};
+    const pageable = client.getEntities(query);
+
+    for await (const _entity of pageable) {
+      // no-op
+    }
+
+    expect(capturedReq).toBeDefined();
+    const protoPrefix = capturedReq!.getQuery()?.getInstanceidstartswith();
+    expect(protoPrefix).toBeUndefined();
+  });
+
+  it("should preserve already-normalized prefix", async () => {
+    const client = new TaskHubGrpcClient({ hostAddress: "localhost:4001" });
+    let capturedReq: pb.QueryEntitiesRequest | undefined;
+
+    mockQueryEntities(client, (req) => { capturedReq = req; });
+
+    const query: EntityQuery = { instanceIdStartsWith: "@counter@" };
+    const pageable = client.getEntities(query);
+
+    for await (const _entity of pageable) {
+      // no-op
+    }
+
+    expect(capturedReq).toBeDefined();
+    const protoPrefix = capturedReq!.getQuery()?.getInstanceidstartswith()?.getValue();
+    expect(protoPrefix).toBe("@counter@");
+  });
+});