fix: cap poller batch size at config limit and make cleanup/requestShutdown optional

pkgs/edge-worker/src/core/Worker.ts

@@ -1,30 +1,33 @@
-import type postgres from 'postgres';
 import type { IBatchProcessor, ILifecycle, WorkerBootstrap } from './types.js';
 import type { Logger } from '../platform/types.js';
 
+export interface WorkerOptions {
+  requestShutdown?: () => void;
+  cleanup?: () => Promise<void>;
+}
+
 export class Worker {
   private lifecycle: ILifecycle;
   private logger: Logger;
   private abortController = new AbortController();
   private readonly requestShutdown?: () => void;
+  private readonly cleanup?: () => Promise<void>;
 
   private batchProcessor: IBatchProcessor;
-  private sql: postgres.Sql;
   private mainLoopPromise: Promise<void> | undefined;
   private deprecationLogged = false;
 
   constructor(
     batchProcessor: IBatchProcessor,
     lifecycle: ILifecycle,
-    sql: postgres.Sql,
     logger: Logger,
-    requestShutdown?: () => void
+    options: WorkerOptions = {}
   ) {
-    this.sql = sql;
     this.lifecycle = lifecycle;
     this.batchProcessor = batchProcessor;
     this.logger = logger;
-    this.requestShutdown = requestShutdown;
+    this.requestShutdown = options.requestShutdown;
+    this.cleanup = options.cleanup;
   }
 
   startOnlyOnce(workerBootstrap: WorkerBootstrap) {
@@ -96,8 +99,10 @@ export class Worker {
 
       this.lifecycle.acknowledgeStop();
 
-      this.logger.debug('-> Closing SQL connection...');
-      await this.sql.end();
+      if (this.cleanup) {
+        this.logger.debug('-> Running worker cleanup...');
+        await this.cleanup();
+      }
 
       // Signal graceful stop complete
       this.logger.shutdown('stopped');

pkgs/edge-worker/src/flow/StepTaskPoller.ts

@@ -43,7 +43,9 @@ export class StepTaskPoller<TFlow extends AnyFlow>
     }
 
     const workerId = this.getWorkerId();
-    const batchSize = limit ?? this.config.batchSize;
+    const batchSize = limit === undefined
+      ? this.config.batchSize
+      : Math.min(this.config.batchSize, limit);
     this.logger.debug(
       `Two-phase polling for flow tasks with batch size ${batchSize}, maxPollSeconds: ${this.config.maxPollSeconds}, pollIntervalMs: ${this.config.pollIntervalMs}`
     );

pkgs/edge-worker/src/flow/createFlowWorker.ts

@@ -81,6 +81,7 @@ export function createFlowWorker<
     );
   }
 
+  const ownsSql = !config.sql;
   const sql =
     config.sql ||
     postgres(config.connectionString as string, {
@@ -207,8 +208,10 @@ export function createFlowWorker<
   return new Worker(
     batchProcessor,
     lifecycle,
-    sql,
     createLogger('Worker'),
-    () => platformAdapter.requestShutdown()
+    {
+      requestShutdown: platformAdapter.requestShutdown?.bind(platformAdapter),
+      cleanup: ownsSql ? () => sql.end() : undefined,
+    }
   );
 }

pkgs/edge-worker/src/platform/SupabasePlatformAdapter.ts

@@ -106,11 +106,13 @@ export class SupabasePlatformAdapter implements PlatformAdapter<SupabaseResource
   async stopWorker(): Promise<void> {
     this.requestShutdown();
 
-    if (this.worker) {
-      await this.worker.stop();
+    try {
+      if (this.worker) {
+        await this.worker.stop();
+      }
+    } finally {
+      await this._platformResources.sql.end();
     }
-
-    await this._platformResources.sql.end();
   }
 
   requestShutdown(): void {

pkgs/edge-worker/src/platform/types.ts

@@ -80,7 +80,7 @@ export interface PlatformAdapter<TResources extends Record<string, unknown> = Re
   /**
    * Trigger the shared shutdown signal used by pollers, executors, and contexts.
    */
-  requestShutdown(): void;
+  requestShutdown?(): void;
 
   /**
    * Get the connection string for the database

pkgs/edge-worker/src/queue/ReadWithPollPoller.ts

@@ -28,7 +28,9 @@ export class ReadWithPollPoller<TPayload extends Json> {
       return [];
     }
 
-    const batchSize = limit ?? this.config.batchSize;
+    const batchSize = limit === undefined
+      ? this.config.batchSize
+      : Math.min(this.config.batchSize, limit);
 
     this.logger.debug(`Polling queue '${this.queue.queueName}' with batch size ${batchSize}`);
     const messages = await this.queue.readWithPoll(

pkgs/edge-worker/src/queue/createQueueWorker.ts

@@ -135,6 +135,7 @@ export function createQueueWorker<TPayload extends Json, TResources extends Reco
   const abortSignal = platformAdapter.shutdownSignal;
 
   // Use provided SQL connection if available, otherwise create one from connection string
+  const ownsSql = !config.sql;
   const sql =
     config.sql ||
     postgres(config.connectionString || '', {
@@ -221,8 +222,10 @@ export function createQueueWorker<TPayload extends Json, TResources extends Reco
   return new Worker(
     batchProcessor,
     lifecycle,
-    sql,
     createLogger('Worker'),
-    () => platformAdapter.requestShutdown()
+    {
+      requestShutdown: platformAdapter.requestShutdown?.bind(platformAdapter),
+      cleanup: ownsSql ? () => sql.end() : undefined,
+    }
   );
 }

pkgs/edge-worker/tests/types/platform-adapter.test-d.ts

@@ -0,0 +1,23 @@
+import type { PlatformAdapter } from '../../src/platform/types.ts';
+
+const adapterWithoutRequestShutdown: PlatformAdapter = {
+  async startWorker() {},
+  async stopWorker() {},
+  get connectionString() {
+    return undefined;
+  },
+  get env() {
+    return {};
+  },
+  get shutdownSignal() {
+    return new AbortController().signal;
+  },
+  get platformResources() {
+    return {};
+  },
+  get isLocalEnvironment() {
+    return false;
+  },
+};
+
+void adapterWithoutRequestShutdown;

pkgs/edge-worker/tests/unit/Poller.batchSize.test.ts

@@ -0,0 +1,184 @@
+import { assertEquals } from '@std/assert';
+import { ReadWithPollPoller } from '../../src/queue/ReadWithPollPoller.ts';
+import { StepTaskPoller } from '../../src/flow/StepTaskPoller.ts';
+import { fakeLogger } from '../fakes.ts';
+
+Deno.test('ReadWithPollPoller caps limit at configured batchSize', async () => {
+  let readQty: number | undefined;
+  const queue = {
+    queueName: 'test_queue',
+    readWithPoll: (qty: number) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+  };
+
+  const poller = new ReadWithPollPoller(
+    queue as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+      visibilityTimeout: 10,
+    },
+    fakeLogger
+  );
+
+  await poller.poll(20);
+
+  assertEquals(readQty, 5);
+});
+
+Deno.test('ReadWithPollPoller uses smaller available slot limit', async () => {
+  let readQty: number | undefined;
+  const queue = {
+    queueName: 'test_queue',
+    readWithPoll: (qty: number) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+  };
+
+  const poller = new ReadWithPollPoller(
+    queue as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+      visibilityTimeout: 10,
+    },
+    fakeLogger
+  );
+
+  await poller.poll(2);
+
+  assertEquals(readQty, 2);
+});
+
+Deno.test('ReadWithPollPoller uses configured batchSize without limit', async () => {
+  let readQty: number | undefined;
+  const queue = {
+    queueName: 'test_queue',
+    readWithPoll: (qty: number) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+  };
+
+  const poller = new ReadWithPollPoller(
+    queue as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+      visibilityTimeout: 10,
+    },
+    fakeLogger
+  );
+
+  await poller.poll();
+
+  assertEquals(readQty, 5);
+});
+
+Deno.test('StepTaskPoller caps limit at configured batchSize', async () => {
+  let readQty: number | undefined;
+  const adapter = {
+    readMessages: (
+      _queueName: string,
+      _visibilityTimeout: number,
+      qty: number
+    ) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+    startTasks: () => Promise.resolve([]),
+  };
+
+  const poller = new StepTaskPoller(
+    adapter as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      queueName: 'test_flow',
+      visibilityTimeout: 10,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+    },
+    () => 'worker-id',
+    fakeLogger
+  );
+
+  await poller.poll(20);
+
+  assertEquals(readQty, 5);
+});
+
+Deno.test('StepTaskPoller uses smaller available slot limit', async () => {
+  let readQty: number | undefined;
+  const adapter = {
+    readMessages: (
+      _queueName: string,
+      _visibilityTimeout: number,
+      qty: number
+    ) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+    startTasks: () => Promise.resolve([]),
+  };
+
+  const poller = new StepTaskPoller(
+    adapter as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      queueName: 'test_flow',
+      visibilityTimeout: 10,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+    },
+    () => 'worker-id',
+    fakeLogger
+  );
+
+  await poller.poll(2);
+
+  assertEquals(readQty, 2);
+});
+
+Deno.test('StepTaskPoller uses configured batchSize without limit', async () => {
+  let readQty: number | undefined;
+  const adapter = {
+    readMessages: (
+      _queueName: string,
+      _visibilityTimeout: number,
+      qty: number
+    ) => {
+      readQty = qty;
+      return Promise.resolve([]);
+    },
+    startTasks: () => Promise.resolve([]),
+  };
+
+  const poller = new StepTaskPoller(
+    adapter as never,
+    new AbortController().signal,
+    {
+      batchSize: 5,
+      queueName: 'test_flow',
+      visibilityTimeout: 10,
+      maxPollSeconds: 1,
+      pollIntervalMs: 100,
+    },
+    () => 'worker-id',
+    fakeLogger
+  );
+
+  await poller.poll();
+
+  assertEquals(readQty, 5);
+});