feat(restate): Add a var to control whether to expose core services or not

2 files changed, 11 insertions, 2 deletions

diff --git a/packages/plugins/queue-restate/src/env.ts b/packages/plugins/queue-restate/src/env.ts
index 01175e86..abcfbe51 100644
--- a/packages/plugins/queue-restate/src/env.ts
+++ b/packages/plugins/queue-restate/src/env.ts
@@ -1,5 +1,12 @@
 import { z } from "zod";
 
+const stringBool = (defaultValue: string) =>
+  z
+    .string()
+    .default(defaultValue)
+    .refine((s) => s === "true" || s === "false")
+    .transform((s) => s === "true");
+
 export const envConfig = z
   .object({
     RESTATE_LISTEN_PORT: z.coerce.number().optional(),
@@ -9,5 +16,6 @@ export const envConfig = z
       .default("http://localhost:8080"),
     RESTATE_ADMIN_ADDR: z.string().optional().default("http://localhost:9070"),
     RESTATE_PUB_KEY: z.string().optional(),
+    RESTATE_EXPOSE_CORE_SERVICES: stringBool("true"),
   })
   .parse(process.env);
diff --git a/packages/plugins/queue-restate/src/index.ts b/packages/plugins/queue-restate/src/index.ts
index fa636f87..07082955 100644
--- a/packages/plugins/queue-restate/src/index.ts
+++ b/packages/plugins/queue-restate/src/index.ts
@@ -134,8 +134,9 @@ class RestateQueueClient implements QueueClient {
       port: envConfig.RESTATE_LISTEN_PORT ?? 0,
       services: [
         ...[...this.services.values()].map((svc) => svc.def),
-        semaphore,
-        idProvider,
+        ...(envConfig.RESTATE_EXPOSE_CORE_SERVICES
+          ? [semaphore, idProvider]
+          : []),
       ],
       identityKeys: envConfig.RESTATE_PUB_KEY
         ? [envConfig.RESTATE_PUB_KEY]