@tstdl/base 0.93.153 → 0.93.155

package/ai/genkit/multi-region.plugin.js

@@ -3,9 +3,27 @@ import { GenkitError, modelRef } from 'genkit';
 import { genkitPlugin } from 'genkit/plugin';
 import { shuffle } from '../../utils/array/index.js';
 import { isInstanceOf, isNullOrUndefined } from '../../utils/type-guards.js';
+import { millisecondsPerMinute, millisecondsPerSecond } from '../../utils/units.js';
 const pluginKey = 'vertexai-multi-location';
 const geminiModelReference = vertexAI.model('gemini-2.5-flash');
 export function vertexAiMultiLocation(options) {
+    const locationConfigs = options.locations.map((location) => {
+        const circuitBreakerKey = `genkit:vertex-ai:location:${location}`;
+        const tokenLimitCircuitBreakerKey = `${circuitBreakerKey}:token-limit`;
+        return {
+            location,
+            circuitBreaker: options.circuitBreakerProvider.provide(circuitBreakerKey, {
+                threshold: 1,
+                resetTimeout: 30 * millisecondsPerSecond,
+                ...options.circuitBreakerConfig,
+            }),
+            tokenLimitCircuitBreaker: options.circuitBreakerProvider.provide(tokenLimitCircuitBreakerKey, {
+                threshold: 1,
+                resetTimeout: 15 * millisecondsPerMinute,
+                ...options.tokenLimitCircuitBreakerConfig,
+            }),
+        };
+    });
     const createVirtualizedModelAction = async (ai, modelName) => {
         const baseModelName = `vertexai/${modelName}`;
         const target = modelName;
@@ -20,20 +38,22 @@ export function vertexAiMultiLocation(options) {
             configSchema: baseModelAction.__action.inputSchema?.shape?.config,
             label: `${baseModelAction.__action.description ?? baseModelAction.__action.name} (Multi-Location Routing)`,
         }, async (request, streamingCallback) => {
-            const shuffledLocations = shuffle([...options.locations]);
+            const shuffledConfigs = shuffle([...locationConfigs]);
             let lastError;
-            for (const location of shuffledLocations) {
-                const circuitBreakerKey = `genkit:vertex-ai:location:${location}`;
-                const circuitBreaker = options.circuitBreakerProvider.provide(circuitBreakerKey, {
-                    threshold: 1, // Aggressive for 429
-                    resetTimeout: options.circuitBreakerConfig?.resetTimeout ?? 30000,
-                    ...options.circuitBreakerConfig,
-                });
+            let isLargeRequest = false;
+            for (const { location, circuitBreaker, tokenLimitCircuitBreaker } of shuffledConfigs) {
                 const check = await circuitBreaker.check();
                 if (!check.allowed) {
                     options.logger.warn(`Location ${location} is currently unhealthy. Skipping...`);
                     continue;
                 }
+                if (isLargeRequest) {
+                    const tokenCheck = await tokenLimitCircuitBreaker.check();
+                    if (!tokenCheck.allowed) {
+                        options.logger.warn(`Location ${location} is known to have a low token limit. Skipping for this large request...`);
+                        continue;
+                    }
+                }
                 try {
                     const result = await baseModelAction({
                         ...request,
@@ -52,12 +72,20 @@ export function vertexAiMultiLocation(options) {
                     if (!isInstanceOf(error, GenkitError)) {
                         throw error;
                     }
-                    const isRetryable = ((error.status == 'RESOURCE_EXHAUSTED') || (error.status == 'UNAVAILABLE') || error.message.includes('quota'));
+                    const isTokenLimitError = (error.status == 'INVALID_ARGUMENT') && error.message.includes('input token count') && error.message.includes('model only supports up to');
+                    const isRetryable = isTokenLimitError || ((error.status == 'RESOURCE_EXHAUSTED') || (error.status == 'UNAVAILABLE') || error.message.includes('quota'));
                     if (!isRetryable) {
                         throw error;
                     }
-                    options.logger.warn(`Location ${location} responded with ${error.status}. Tripping circuit breaker and trying next location...`);
-                    await circuitBreaker.recordFailure();
+                    if (isTokenLimitError) {
+                        options.logger.warn(`Location ${location} responded with token limit error. Trying next location...`);
+                        isLargeRequest = true;
+                        await tokenLimitCircuitBreaker.recordFailure();
+                    }
+                    else {
+                        options.logger.warn(`Location ${location} responded with ${error.status}. Tripping circuit breaker and trying next location...`);
+                        await circuitBreaker.recordFailure();
+                    }
                 }
             }
             throw lastError;

package/ai/genkit/tests/token-limit-fallback.test.d.ts

@@ -0,0 +1,2 @@
+/** biome-ignore-all lint/suspicious/useAwait: defineModel requires async */
+export {};

package/ai/genkit/tests/token-limit-fallback.test.js

@@ -0,0 +1,138 @@
+/** biome-ignore-all lint/suspicious/useAwait: defineModel requires async */
+import { genkit, GenkitError, z } from 'genkit';
+import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import { CircuitBreakerState } from '../../../circuit-breaker/index.js';
+import { CircuitBreakerProvider } from '../../../circuit-breaker/provider.js';
+import { Logger } from '../../../logger/logger.js';
+import { setupIntegrationTest } from '../../../testing/index.js';
+import { vertexAiMultiLocation } from '../multi-region.plugin.js';
+vi.mock('#/utils/array/index.js', async (importOriginal) => {
+    const actual = await importOriginal();
+    return {
+        ...actual,
+        shuffle: vi.fn((items) => [...items]),
+    };
+});
+vi.mock('@genkit-ai/google-genai', () => ({
+    // biome-ignore lint/style/useNamingConvention: given
+    vertexAI: {
+        model: vi.fn((name) => ({
+            name: `vertexai/${name}`,
+            info: { label: 'mock' },
+            configSchema: z.object({}),
+        })),
+    },
+    // biome-ignore lint/style/useNamingConvention: given
+    googleAI: vi.fn(),
+}));
+describe('Genkit vertexai-multi-location Token Limit Fallback Tests', () => {
+    let ai;
+    let cbProvider;
+    let logger;
+    beforeAll(async () => {
+        const { injector } = await setupIntegrationTest({ modules: { circuitBreaker: true } });
+        cbProvider = injector.resolve(CircuitBreakerProvider);
+        logger = injector.resolve(Logger, 'Test');
+    });
+    beforeEach(async () => {
+        vi.clearAllMocks();
+        ai = genkit({
+            plugins: [
+                vertexAiMultiLocation({
+                    locations: ['region-1', 'region-2', 'region-3'],
+                    circuitBreakerProvider: cbProvider,
+                    logger,
+                    circuitBreakerConfig: { resetTimeout: 1_000_000, threshold: 1 },
+                }),
+            ],
+        });
+        const config = { threshold: 1, resetTimeout: 1_000_000 };
+        await cbProvider.provide('genkit:vertex-ai:location:region-1', config).recordSuccess();
+        await cbProvider.provide('genkit:vertex-ai:location:region-2', config).recordSuccess();
+        await cbProvider.provide('genkit:vertex-ai:location:region-3', config).recordSuccess();
+        await cbProvider.provide('genkit:vertex-ai:location:region-1:token-limit', config).recordSuccess();
+        await cbProvider.provide('genkit:vertex-ai:location:region-2:token-limit', config).recordSuccess();
+        await cbProvider.provide('genkit:vertex-ai:location:region-3:token-limit', config).recordSuccess();
+    });
+    it('should fallback on token limit error but NOT trip main circuit breaker', async () => {
+        const tokenLimitErrorMessage = 'Unable to submit request because the input token count is 135224 but model only supports up to 131072.';
+        let region1Called = false;
+        let region2Called = false;
+        ai.defineModel({
+            name: 'vertexai/gemini-2.5-flash',
+        }, async (request) => {
+            if (request.config?.location === 'region-1') {
+                region1Called = true;
+                throw new GenkitError({
+                    status: 'INVALID_ARGUMENT',
+                    message: tokenLimitErrorMessage,
+                });
+            }
+            if (request.config?.location === 'region-2') {
+                region2Called = true;
+                return {
+                    message: {
+                        role: 'model',
+                        content: [{ text: 'success from region-2' }],
+                    },
+                };
+            }
+            throw new Error('Unexpected location');
+        });
+        const response = await ai.generate({
+            model: 'vertexai-multi-location/gemini-2.5-flash',
+            prompt: 'test',
+        });
+        expect(response.text).toBe('success from region-2');
+        expect(region1Called).toBe(true);
+        expect(region2Called).toBe(true);
+        // Verify main circuit breaker for region-1 is still CLOSED (allowed)
+        const cb = cbProvider.provide('genkit:vertex-ai:location:region-1', { threshold: 1, resetTimeout: 1000000 });
+        const status = await cb.check();
+        expect(status.state).toBe(CircuitBreakerState.Closed);
+        expect(status.allowed).toBe(true);
+    });
+    it('should skip locations with known token limits within the same request once it is known to be large', async () => {
+        const tokenLimitErrorMessage = 'Unable to submit request because the input token count is 135224 but model only supports up to 131072.';
+        // First, trip the token limit breaker for region-2
+        const tokenLimitCB2 = cbProvider.provide('genkit:vertex-ai:location:region-2:token-limit', { threshold: 1, resetTimeout: 1000000 });
+        await tokenLimitCB2.recordFailure();
+        let region1Called = false;
+        let region2Called = false;
+        let region3Called = false;
+        ai.defineModel({
+            name: 'vertexai/gemini-2.5-flash',
+        }, async (request) => {
+            if (request.config?.location === 'region-1') {
+                region1Called = true;
+                throw new GenkitError({
+                    status: 'INVALID_ARGUMENT',
+                    message: tokenLimitErrorMessage,
+                });
+            }
+            if (request.config?.location === 'region-2') {
+                region2Called = true;
+                return { message: { role: 'model', content: [{ text: 'success from region-2' }] } };
+            }
+            if (request.config?.location === 'region-3') {
+                region3Called = true;
+                return {
+                    message: {
+                        role: 'model',
+                        content: [{ text: 'success from region-3' }],
+                    },
+                };
+            }
+            throw new Error('Unexpected location');
+        });
+        // Initial shuffle is mocked to [region-1, region-2, region-3]
+        const response = await ai.generate({
+            model: 'vertexai-multi-location/gemini-2.5-flash',
+            prompt: 'test',
+        });
+        expect(response.text).toBe('success from region-3');
+        expect(region1Called).toBe(true); // Fails, makes request "known to be large"
+        expect(region2Called).toBe(false); // Should be skipped because it is known to have a low limit
+        expect(region3Called).toBe(true); // Should be tried as it is not known to be limited
+    });
+});

package/ai/genkit/types.d.ts

@@ -1,5 +1,5 @@
 import type { CircuitBreakerConfig } from '../../circuit-breaker/circuit-breaker.js';
-export interface VertexAiMultiLocationOptions {
+export declare abstract class VertexAiMultiLocationOptions {
     /** The Google Cloud locations to use for routing. */
     locations: string[];
     /**
@@ -7,4 +7,9 @@ export interface VertexAiMultiLocationOptions {
      * By default, a threshold of 1 is used for 429 errors.
      */
     circuitBreakerConfig?: Partial<CircuitBreakerConfig>;
+    /**
+     * Optional token limit circuit breaker configuration.
+     * By default, a threshold of 1 and a reset timeout of 15 minutes is used.
+     */
+    tokenLimitCircuitBreakerConfig?: Partial<CircuitBreakerConfig>;
 }

package/ai/genkit/types.js

@@ -1 +1,14 @@
-export {};
+export class VertexAiMultiLocationOptions {
+    /** The Google Cloud locations to use for routing. */
+    locations;
+    /**
+     * Optional circuit breaker configuration.
+     * By default, a threshold of 1 is used for 429 errors.
+     */
+    circuitBreakerConfig;
+    /**
+     * Optional token limit circuit breaker configuration.
+     * By default, a threshold of 1 and a reset timeout of 15 minutes is used.
+     */
+    tokenLimitCircuitBreakerConfig;
+}

package/examples/document-management/main.js

@@ -34,7 +34,7 @@ import { TstdlCategoryParents, TstdlDocumentCategoryLabels, TstdlDocumentPropert
 const config = {
     database: {
         host: string('DATABASE_HOST', '127.0.0.1'),
-        port: positiveInteger('DATABASE_PORT', 5432),
+        port: positiveInteger('DATABASE_PORT', 15433),
         user: string('DATABASE_USER', 'tstdl'),
         pass: string('DATABASE_PASS', 'wf7rq6glrk5jykne'),
         database: string('DATABASE_NAME', 'tstdl'),

package/examples/rate-limit/basic-usage.d.ts

@@ -0,0 +1 @@
+export {};

package/examples/rate-limit/basic-usage.js

@@ -0,0 +1,52 @@
+import { configurePostgresRateLimiter, migratePostgresRateLimiterSchema } from '../../rate-limit/postgres/index.js';
+import { RateLimiterProvider } from '../../rate-limit/index.js';
+import { Injector, runInInjectionContext } from '../../injector/index.js';
+import { configureOrm } from '../../orm/server/index.js';
+import { ConsoleLogTransport, LogFormatter, Logger, LogTransport, PrettyPrintLogFormatter } from '../../logger/index.js';
+import * as configParser from '../../utils/config-parser.js';
+import { timeout } from '../../utils/timing.js';
+async function main() {
+    const injector = new Injector('ExampleInjector');
+    // 1. Configure Logging
+    injector.register(LogFormatter, { useToken: PrettyPrintLogFormatter });
+    injector.register(LogTransport, { useToken: ConsoleLogTransport });
+    const logger = injector.resolve(Logger);
+    // 2. Configure Database
+    configureOrm({
+        connection: {
+            host: configParser.string('DATABASE_HOST', '127.0.0.1'),
+            port: configParser.positiveInteger('DATABASE_PORT', 15433),
+            user: configParser.string('DATABASE_USER', 'tstdl'),
+            password: configParser.string('DATABASE_PASS', 'wf7rq6glrk5jykne'),
+            database: configParser.string('DATABASE_NAME', 'tstdl'),
+        },
+    });
+    // 3. Configure Rate Limiter
+    configurePostgresRateLimiter();
+    // 4. Run Migrations (for setup)
+    logger.info('Running migrations...');
+    await runInInjectionContext(injector, migratePostgresRateLimiterSchema);
+    // 5. Get a Rate Limiter Instance
+    const provider = injector.resolve(RateLimiterProvider);
+    const limiter = provider.get('api-limiter', {
+        burstCapacity: 10,
+        refillInterval: 1000, // 10 tokens per second
+    });
+    const resource = 'user-123';
+    // 6. Simulate Traffic
+    logger.info('Starting simulation...');
+    for (let i = 0; i < 15; i++) {
+        const success = await limiter.tryAcquire(resource);
+        if (success) {
+            logger.info(`Request ${i + 1}: Allowed`);
+        }
+        else {
+            logger.warn(`Request ${i + 1}: Throttled`);
+        }
+        // Small delay to simulate some processing/network time
+        await timeout(50);
+    }
+    // 7. Cleanup
+    await injector.dispose();
+}
+void main();

package/orm/sqls/sqls.d.ts

@@ -4,7 +4,7 @@
  * simplifying common SQL operations like generating UUIDs, working with intervals,
  * and aggregating data.
  */
-import { Column, SQL, type AnyColumn, type SQLChunk, type SQLWrapper } from 'drizzle-orm';
+import { Column, type AnyColumn, type SQL, type SQLChunk, type SQLWrapper } from 'drizzle-orm';
 import type { GetSelectTableSelection, SelectResultField, TableLike } from 'drizzle-orm/query-builders/select.types';
 import type { EnumerationObject, EnumerationValue, Record } from '../../types/types.js';
 import { type PgEnumFromEnumeration } from '../enums.js';
@@ -125,7 +125,7 @@ export declare function autoAlias<T>(column: AnyColumn<{
  * @param unit - The unit of the interval (e.g., 'day', 'hour').
  * @returns A Drizzle SQL object representing the interval.
  */
-export declare function interval(value: number, unit: IntervalUnit): SQL;
+export declare function interval(value: number | SQL<number>, unit: IntervalUnit): SQL;
 /**
  * Creates a PostgreSQL `array_agg` aggregate function call.
  * Aggregates values from a column into a PostgreSQL array.
@@ -183,6 +183,7 @@ export declare function greatest<T extends (Column | SQL | SQL.Aliased | number)
     [P in keyof T]: T[P] extends number ? Exclude<T[P], number> | SQL<number> : T[P];
 }[number]>>;
 export declare function greatest<T>(...values: T[]): SQL<SelectResultField<T>>;
+export declare function power(base: number | SQLChunk, exponent: number | SQLChunk): SQL<number>;
 export declare function unnest<T>(array: SQL<readonly T[]> | SQL.Aliased<readonly T[]> | Column): SQL<T>;
 /**
  * Creates a PostgreSQL array contains operator expression (@>).

package/orm/sqls/sqls.js

@@ -4,7 +4,7 @@
  * simplifying common SQL operations like generating UUIDs, working with intervals,
  * and aggregating data.
  */
-import { and, Column, eq, isSQLWrapper, sql, SQL, isNotNull as sqlIsNotNull, isNull as sqlIsNull, Table } from 'drizzle-orm';
+import { and, Column, eq, isSQLWrapper, sql, isNotNull as sqlIsNotNull, isNull as sqlIsNull, Table } from 'drizzle-orm';
 import { match, P } from 'ts-pattern';
 import { distinct, toArray } from '../../utils/array/array.js';
 import { objectEntries, objectValues } from '../../utils/object/object.js';
@@ -146,7 +146,7 @@ export function autoAlias(column) {
  * @returns A Drizzle SQL object representing the interval.
  */
 export function interval(value, unit) {
-    return sql `(${value} ||' ${sql.raw(unit)}')::interval`;
+    return sql `(${value} || ' ${sql.raw(unit)}')::interval`;
 }
 /**
  * Creates a PostgreSQL `array_agg` aggregate function call.
@@ -224,6 +224,9 @@ export function greatest(...values) {
     const sqlValues = values.map((value) => isNumber(value) ? sql.raw(String(value)) : value);
     return sql `greatest(${sql.join(sqlValues, sql.raw(', '))})`;
 }
+export function power(base, exponent) {
+    return sql `power(${base}, ${exponent})`;
+}
 export function unnest(array) {
     return sql `unnest(${array})`;
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tstdl/base",
-  "version": "0.93.153",
+  "version": "0.93.155",
   "author": "Patrick Hein",
   "publishConfig": {
     "access": "public"
@@ -152,8 +152,8 @@
     "type-fest": "^5.4"
   },
   "peerDependencies": {
-    "@aws-sdk/client-s3": "^3.1000",
-    "@aws-sdk/s3-request-presigner": "^3.1000",
+    "@aws-sdk/client-s3": "^3.1002",
+    "@aws-sdk/s3-request-presigner": "^3.1002",
     "@genkit-ai/google-genai": "^1.29",
     "@google-cloud/storage": "^7.19",
     "@toon-format/toon": "^2.1.0",
@@ -168,7 +168,7 @@
     "handlebars": "^4.7",
     "mjml": "^4.18",
     "nodemailer": "^8.0",
-    "pg": "^8.19",
+    "pg": "^8.20",
     "playwright": "^1.58",
     "preact": "^10.28",
     "preact-render-to-string": "^6.6",

package/task-queue/postgres/task-queue.js

@@ -60,15 +60,14 @@ import { aliasedTable, and, asc, count, eq, gt, gte, inArray, lt, lte, notExists
 import { filter, merge, throttleTime } from 'rxjs';
 import { CancellationSignal } from '../../cancellation/index.js';
 import { CircuitBreaker, CircuitBreakerState } from '../../circuit-breaker/index.js';
-import { serializeError, TimeoutError } from '../../errors/index.js';
+import { NotFoundError, serializeError, TimeoutError } from '../../errors/index.js';
 import { afterResolve, inject, provide, Singleton } from '../../injector/index.js';
 import { Logger } from '../../logger/index.js';
 import { MessageBus } from '../../message-bus/index.js';
-import { arrayOverlaps, caseWhen, coalesce, enumValue, getEntityIds, greatest, interval, jsonbBuildObject, RANDOM_UUID_V4, TRANSACTION_TIMESTAMP } from '../../orm/index.js';
-import { Database, DatabaseConfig, injectRepository } from '../../orm/server/index.js';
+import { arrayOverlaps, caseWhen, coalesce, enumValue, greatest, interval, jsonbBuildObject, least, power, RANDOM_UUID_V4, TRANSACTION_TIMESTAMP } from '../../orm/index.js';
+import { DatabaseConfig, injectRepository } from '../../orm/server/index.js';
 import { RateLimiter } from '../../rate-limit/index.js';
 import { distinct, toArray } from '../../utils/array/array.js';
-import { currentTimestamp } from '../../utils/date-time.js';
 import { Timer } from '../../utils/timer.js';
 import { cancelableTimeout } from '../../utils/timing.js';
 import { isArray, isDefined, isNotNull, isNull, isNumber, isString, isUndefined } from '../../utils/type-guards.js';
@@ -76,10 +75,9 @@ import { millisecondsPerMinute, millisecondsPerSecond } from '../../utils/units.
 import { defaultQueueConfig, queueableOrWaitableStatuses, queueableStatuses, TaskDependencyType, TaskQueue, TaskStatus, terminalStatuses } from '../task-queue.js';
 import { ensureTaskError } from '../task.error.js';
 import { PostgresTaskQueueModuleConfig } from './module.js';
-import { taskArchive as taskArchiveTable, taskDependency as taskDependencyTable, taskDependencyType, taskStatus, task as taskTable } from './schemas.js';
+import { taskArchive as taskArchiveTable, taskDependency as taskDependencyTable, taskStatus, task as taskTable } from './schemas.js';
 import { PostgresTask, PostgresTaskArchive } from './task.model.js';
 let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
-    #database = inject(Database);
     #repository = injectRepository(PostgresTask);
     #archiveRepository = injectRepository(PostgresTaskArchive);
     #config = this.config;
@@ -399,13 +397,17 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
             .set({
             unresolvedScheduleDependencies: sql `${taskTable.unresolvedScheduleDependencies} + ${updates.scheduleIncrement}`,
             unresolvedCompleteDependencies: sql `${taskTable.unresolvedCompleteDependencies} + ${updates.completeIncrement}`,
-            status: caseWhen(and(eq(taskTable.status, TaskStatus.Pending), gt(sql `${taskTable.unresolvedScheduleDependencies} + ${updates.scheduleIncrement}`, 0)), enumValue(TaskStatus, taskStatus, TaskStatus.Waiting)).else(taskTable.status),
+            status: caseWhen(and(eq(taskTable.status, TaskStatus.Pending), gt(sql `${taskTable.unresolvedScheduleDependencies} + ${updates.scheduleIncrement}`, 0)), TaskStatus.Waiting).else(taskTable.status),
         })
             .from(updates)
             .where(eq(taskTable.id, updates.taskId))
             .returning({ id: taskTable.id, status: taskTable.status, namespace: taskTable.namespace });
+        const notifiedNamespaces = new Set();
         for (const row of updatedRows) {
-            this.notify(row.namespace);
+            if (!notifiedNamespaces.has(row.namespace)) {
+                this.notify(row.namespace);
+                notifiedNamespaces.add(row.namespace);
+            }
         }
     }
     async has(id, options) {
@@ -535,23 +537,29 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         await this.cancelMany([id], options);
     }
     async cancelMany(ids, options) {
+        if (ids.length == 0) {
+            return;
+        }
         await this.#repository.useTransaction(options?.transaction, async (tx) => {
-            const tree = await this.getTree(ids, { transaction: tx });
-            const treeIds = tree.map((task) => task.id);
-            if (treeIds.length == 0) {
-                return;
-            }
-            const cancelledRows = await tx.pgTransaction
-                .update(taskTable)
-                .set({
-                status: TaskStatus.Cancelled,
-                token: null,
-                completeTimestamp: TRANSACTION_TIMESTAMP,
-            })
-                .where(and(inArray(taskTable.id, treeIds), notInArray(taskTable.status, terminalStatuses)))
-                .returning({ id: taskTable.id, namespace: taskTable.namespace });
-            if (cancelledRows.length > 0) {
-                await this.resolveDependenciesMany(cancelledRows.map((row) => ({ id: row.id, status: TaskStatus.Cancelled, namespace: row.namespace })), { transaction: tx });
+            const cancelledRows = await tx.pgTransaction.execute(sql `
+        WITH RECURSIVE task_tree AS (
+          SELECT id FROM ${taskTable} WHERE ${inArray(taskTable.id, ids)}
+          UNION ALL
+          SELECT child.id FROM ${taskTable} child JOIN task_tree parent ON child.parent_id = parent.id
+        )
+        UPDATE ${taskTable}
+        SET
+          status = ${enumValue(TaskStatus, taskStatus, TaskStatus.Cancelled)},
+          token = NULL,
+          complete_timestamp = ${TRANSACTION_TIMESTAMP}
+        FROM task_tree
+        WHERE
+          ${taskTable.id} = task_tree.id
+          AND ${taskTable.status} NOT IN (${sql.join(terminalStatuses.map((s) => enumValue(TaskStatus, taskStatus, s)), sql `, `)})
+        RETURNING ${taskTable.id} as id, ${taskTable.namespace} as namespace
+      `);
+            if (cancelledRows.rows.length > 0) {
+                await this.resolveDependenciesMany(cancelledRows.rows.map((row) => ({ id: row.id, status: TaskStatus.Cancelled, namespace: row.namespace })), { transaction: tx });
             }
         });
     }
@@ -705,42 +713,29 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         if (isNull(task.token)) {
             return undefined;
         }
-        return await this.#repository.useTransaction(options?.transaction, async (tx) => {
-            const update = {
-                visibilityDeadline: sql `${TRANSACTION_TIMESTAMP} + ${interval(this.visibilityTimeout, 'milliseconds')}`,
-            };
-            if (isDefined(options?.progress)) {
-                update.progress = options.progress;
-            }
-            if (isDefined(options?.state)) {
-                update.state = options.state;
-            }
-            const [updatedRow] = await tx.pgTransaction
-                .update(taskTable)
-                .set(update)
-                .where(and(eq(taskTable.id, task.id), eq(taskTable.status, TaskStatus.Running), isNull(task.token) ? sqlIsNull(taskTable.token) : eq(taskTable.token, task.token), gt(taskTable.startTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.maxExecutionTime, 'milliseconds')}`)))
-                .returning();
-            if (isDefined(updatedRow)) {
-                return await this.#repository.mapToEntity(updatedRow);
-            }
-            const [existingRow] = await tx.pgTransaction
-                .select({ startTimestamp: taskTable.startTimestamp })
-                .from(taskTable)
-                .where(and(eq(taskTable.id, task.id), eq(taskTable.status, TaskStatus.Running), isNull(task.token) ? sqlIsNull(taskTable.token) : eq(taskTable.token, task.token)));
-            if (isDefined(existingRow) && isNotNull(existingRow.startTimestamp) && (currentTimestamp() - existingRow.startTimestamp) > this.maxExecutionTime) {
-                await tx.pgTransaction
-                    .update(taskTable)
-                    .set({
-                    status: TaskStatus.TimedOut,
-                    completeTimestamp: TRANSACTION_TIMESTAMP,
-                    error: { code: 'MaxTimeExceeded', message: 'Hard Execution Timeout' },
-                })
-                    .where(eq(taskTable.id, task.id));
-                await this.resolveDependenciesMany([{ id: task.id, status: TaskStatus.TimedOut, namespace: task.namespace }], { transaction: tx });
-                this.notify();
-            }
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
+        const exceededMaxExecutionTime = lt(taskTable.startTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.maxExecutionTime, 'milliseconds')}`);
+        const [updatedRow] = await session
+            .update(taskTable)
+            .set({
+            status: caseWhen(exceededMaxExecutionTime, enumValue(TaskStatus, taskStatus, TaskStatus.TimedOut)).else(taskTable.status),
+            visibilityDeadline: caseWhen(exceededMaxExecutionTime, null).else(sql `${TRANSACTION_TIMESTAMP} + ${interval(this.visibilityTimeout, 'milliseconds')}`),
+            completeTimestamp: caseWhen(exceededMaxExecutionTime, TRANSACTION_TIMESTAMP).else(taskTable.completeTimestamp),
+            error: caseWhen(exceededMaxExecutionTime, jsonbBuildObject({ code: 'MaxTimeExceeded', message: 'Hard Execution Timeout' })).else(taskTable.error),
+            progress: caseWhen(exceededMaxExecutionTime, taskTable.progress).else(isDefined(options?.progress) ? options.progress : taskTable.progress),
+            state: caseWhen(exceededMaxExecutionTime, taskTable.state).else(isDefined(options?.state) ? options.state : taskTable.state),
+        })
+            .where(and(eq(taskTable.id, task.id), eq(taskTable.status, TaskStatus.Running), isNull(task.token) ? sqlIsNull(taskTable.token) : eq(taskTable.token, task.token)))
+            .returning();
+        if (isUndefined(updatedRow)) {
             return undefined;
-        });
+        }
+        if (updatedRow.status == TaskStatus.TimedOut) {
+            await this.resolveDependencies(task.id, TaskStatus.TimedOut, { namespace: task.namespace, transaction: options?.transaction });
+            this.notify();
+            return undefined;
+        }
+        return await this.#repository.mapToEntity(updatedRow);
     }
     async touchMany(tasks, options) {
         if (tasks.length == 0) {
@@ -778,34 +773,24 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
     }
     async complete(task, options) {
         await this.#repository.useTransaction(options?.transaction, async (tx) => {
-            const [freshTask] = await tx.pgTransaction
-                .select({ unresolvedCompleteDependencies: taskTable.unresolvedCompleteDependencies })
-                .from(taskTable)
-                .where(eq(taskTable.id, task.id))
-                .for('update');
-            if (isUndefined(freshTask)) {
-                return;
-            }
-            const hasActiveChildren = freshTask.unresolvedCompleteDependencies > 0;
-            const nextStatus = hasActiveChildren ? TaskStatus.WaitingChildren : TaskStatus.Completed;
             const [updatedTask] = await tx.pgTransaction.update(taskTable)
                 .set({
-                status: nextStatus,
+                status: caseWhen(gt(taskTable.unresolvedCompleteDependencies, 0), enumValue(TaskStatus, taskStatus, TaskStatus.WaitingChildren)).else(enumValue(TaskStatus, taskStatus, TaskStatus.Completed)),
                 token: null,
                 result: options?.result,
-                progress: hasActiveChildren ? task.progress : 1,
-                completeTimestamp: (nextStatus == TaskStatus.Completed) ? TRANSACTION_TIMESTAMP : null,
+                progress: caseWhen(gt(taskTable.unresolvedCompleteDependencies, 0), task.progress).else(sql.raw('1')),
+                completeTimestamp: caseWhen(gt(taskTable.unresolvedCompleteDependencies, 0), null).else(TRANSACTION_TIMESTAMP),
                 visibilityDeadline: null,
             })
                 .where(and(eq(taskTable.id, task.id), isNull(task.token) ? sqlIsNull(taskTable.token) : eq(taskTable.token, task.token)))
-                .returning({ id: taskTable.id });
+                .returning({ id: taskTable.id, status: taskTable.status });
             if (isUndefined(updatedTask)) {
                 return;
             }
-            if (nextStatus == TaskStatus.Completed) {
+            if (updatedTask.status == TaskStatus.Completed) {
                 await this.#circuitBreaker.recordSuccess();
             }
-            await this.resolveDependencies(task.id, nextStatus, { namespace: task.namespace, transaction: tx });
+            await this.resolveDependencies(task.id, updatedTask.status, { namespace: task.namespace, transaction: tx });
         });
     }
     async completeMany(tasks, options) {
@@ -855,9 +840,9 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         const isRetryable = (options?.fatal != true) && (task.tries < this.maxTries);
         const nextStatus = isRetryable ? TaskStatus.Retrying : TaskStatus.Dead;
         const delay = isRetryable
-            ? Math.min(this.retryDelayMaximum, this.retryDelayMinimum * (this.retryDelayGrowth ** task.tries))
+            ? Math.min(this.retryDelayMaximum, this.retryDelayMinimum * (this.retryDelayGrowth ** (task.tries - 1)))
             : 0;
-        const nextSchedule = currentTimestamp() + delay;
+        const nextSchedule = sql `${TRANSACTION_TIMESTAMP} + ${interval(delay, 'milliseconds')}`;
         await this.#repository.useTransaction(options?.transaction, async (tx) => {
             const [updatedRow] = await tx.pgTransaction
                 .update(taskTable)
@@ -866,7 +851,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
                 token: null,
                 error: serializeError(error),
                 visibilityDeadline: null,
-                scheduleTimestamp: nextSchedule,
+                scheduleTimestamp: isRetryable ? nextSchedule : taskTable.scheduleTimestamp,
                 startTimestamp: null,
                 completeTimestamp: (nextStatus == TaskStatus.Dead) ? TRANSACTION_TIMESTAMP : null,
             })
@@ -889,31 +874,31 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
                 const isRetryable = (task.tries < this.maxTries);
                 const nextStatus = isRetryable ? TaskStatus.Retrying : TaskStatus.Dead;
                 const delay = isRetryable
-                    ? Math.min(this.retryDelayMaximum, this.retryDelayMinimum * (this.retryDelayGrowth ** task.tries))
+                    ? Math.min(this.retryDelayMaximum, this.retryDelayMinimum * (this.retryDelayGrowth ** (task.tries - 1)))
                     : 0;
-                const nextSchedule = new Date(currentTimestamp() + delay);
-                const completeTimestamp = (nextStatus == TaskStatus.Dead) ? new Date() : null;
-                return sql `(${task.id}::uuid, ${task.token}::uuid, ${task.tries}::int, ${nextStatus}::text, ${serializeError(error)}::jsonb, ${nextSchedule}::timestamptz, ${completeTimestamp}::timestamptz)`;
+                const nextSchedule = sql `(${TRANSACTION_TIMESTAMP} + ${interval(delay, 'milliseconds')})`;
+                const completeTimestamp = (nextStatus == TaskStatus.Dead) ? TRANSACTION_TIMESTAMP : null;
+                return sql `(${task.id}::uuid, ${task.token}::uuid, ${task.tries}::int, ${nextStatus}::${taskStatus}, ${serializeError(error)}::jsonb, ${nextSchedule}::timestamptz, ${completeTimestamp}::timestamptz)`;
             });
             const updates = tx.pgTransaction.$with('updates').as((qb) => qb
                 .select({
-                updateId: sql `(id)::uuid`.as('update_id'),
-                updateToken: sql `(token)::uuid`.as('update_token'),
-                updateTries: sql `(tries)::int`.as('update_tries'),
-                updateStatus: sql `(status)::text`.as('update_status'),
-                updateError: sql `(error)::jsonb`.as('update_error'),
-                updateSchedule: sql `(schedule_timestamp)::timestamptz`.as('update_schedule'),
-                updateComplete: sql `(complete_timestamp)::timestamptz`.as('update_complete'),
+                updateId: sql `(id)`.as('update_id'),
+                updateToken: sql `(token)`.as('update_token'),
+                updateTries: sql `(tries)`.as('update_tries'),
+                updateStatus: sql `(status)`.as('update_status'),
+                updateError: sql `(error)`.as('update_error'),
+                updateSchedule: sql `(schedule_timestamp)`.as('update_schedule'),
+                updateComplete: sql `(complete_timestamp)`.as('update_complete'),
             })
                 .from(sql `(VALUES ${sql.join(rows, sql `, `)}) AS t(id, token, tries, status, error, schedule_timestamp, complete_timestamp)`));
             const updated = tx.pgTransaction.$with('updated').as(() => tx.pgTransaction
                 .update(taskTable)
                 .set({
-                status: sql `${updates.updateStatus}::${taskStatus}`,
+                status: sql `${updates.updateStatus}`,
                 token: null,
                 error: sql `${updates.updateError}`,
                 visibilityDeadline: null,
-                scheduleTimestamp: sql `${updates.updateSchedule}`,
+                scheduleTimestamp: caseWhen(eq(updates.updateStatus, TaskStatus.Retrying), updates.updateSchedule).else(taskTable.scheduleTimestamp),
                 startTimestamp: null,
                 completeTimestamp: sql `${updates.updateComplete}`,
             })
@@ -934,41 +919,41 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         await this.resolveDependenciesMany([{ id, status, namespace: options?.namespace }], options);
     }
     async resolveDependenciesMany(tasks, options) {
-        if (tasks.length == 0) {
+        const tasksToResolve = tasks.filter((t) => terminalStatuses.includes(t.status));
+        if (tasksToResolve.length == 0) {
             return;
         }
-        const taskStatusMap = new Map(tasks.map((t) => [t.id, t.status]));
         const notifiedNamespaces = new Set();
         await this.#repository.useTransaction(options?.transaction, async (tx) => {
-            const taskIds = tasks.map((t) => t.id);
-            const dependents = await tx.pgTransaction
+            const taskValues = tasksToResolve.map((t) => sql `(${t.id}::uuid, ${t.status}::${taskStatus})`);
+            // 1. CTE: Load the incoming terminal tasks into a memory table
+            const resolvedTasks = tx.pgTransaction.$with('resolved_tasks').as((qb) => qb
                 .select({
-                taskId: taskDependencyTable.taskId,
-                dependencyTaskId: taskDependencyTable.dependencyTaskId,
-                type: taskDependencyTable.type,
-                requiredStatuses: taskDependencyTable.requiredStatuses,
+                resolvedId: sql `(id)`.as('resolved_id'),
+                resolvedStatus: sql `(status)`.as('resolved_status'),
             })
-                .from(taskDependencyTable)
-                .where(inArray(taskDependencyTable.dependencyTaskId, taskIds));
-            if (dependents.length == 0) {
-                return;
-            }
-            const resolvedEdges = [];
-            const abortOnDependencyFailureTaskIds = new Set();
-            for (const dep of dependents) {
-                const status = taskStatusMap.get(dep.dependencyTaskId);
-                const isMatched = dep.requiredStatuses.includes(status);
-                const isTerminal = terminalStatuses.includes(status);
-                if (isMatched || isTerminal) {
-                    resolvedEdges.push(dep);
-                    if (!isMatched) {
-                        abortOnDependencyFailureTaskIds.add(dep.taskId);
-                    }
-                }
-            }
+                .from(sql `(VALUES ${sql.join(taskValues, sql `, `)}) AS t(id, status)`));
+            // 2. CTE: Atomically delete all edges pointing to these terminal tasks and return them
+            const deletedEdges = tx.pgTransaction.$with('deleted_edges').as(() => tx.pgTransaction
+                .delete(taskDependencyTable)
+                .where(inArray(taskDependencyTable.dependencyTaskId, tx.pgTransaction.select({ id: resolvedTasks.resolvedId }).from(resolvedTasks)))
+                .returning());
+            // 3. Execute: Join deleted edges with their resolving status to determine if they matched the required status
+            const resolvedEdges = await tx.pgTransaction
+                .with(resolvedTasks, deletedEdges)
+                .select({
+                taskId: deletedEdges.taskId,
+                dependencyTaskId: deletedEdges.dependencyTaskId,
+                type: deletedEdges.type,
+                isMatched: sql `${resolvedTasks.resolvedStatus} = ANY(${deletedEdges.requiredStatuses})`.as('is_matched'),
+            })
+                .from(deletedEdges)
+                .innerJoin(resolvedTasks, eq(deletedEdges.dependencyTaskId, resolvedTasks.resolvedId));
             if (resolvedEdges.length == 0) {
                 return;
             }
+            // Extract skipped dependencies (terminal status but not a matched status)
+            const abortOnDependencyFailureTaskIds = distinct(resolvedEdges.filter((d) => !d.isMatched).map((d) => d.taskId));
             const sortedResolvedEdges = resolvedEdges.toSorted((a, b) => {
                 const idCompare = a.taskId.localeCompare(b.taskId);
                 if (idCompare != 0) {
@@ -980,18 +965,10 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
                 }
                 return a.type.localeCompare(b.type);
             });
-            const edgeValues = sortedResolvedEdges.map((e) => sql `(${e.taskId}::uuid, ${e.dependencyTaskId}::uuid, ${e.type}::text)`);
-            await tx.pgTransaction.execute(sql `
-        DELETE FROM ${taskDependencyTable}
-        WHERE (task_id, dependency_task_id, type) IN (
-          SELECT t.task_id, t.dependency_task_id, t.type::${taskDependencyType}
-          FROM (VALUES ${sql.join(edgeValues, sql `, `)}) AS t(task_id, dependency_task_id, type)
-        )
-      `);
             const terminalTasks = [];
             const skippedTaskIds = new Set();
-            if (abortOnDependencyFailureTaskIds.size > 0) {
-                const sortedAbortIds = [...abortOnDependencyFailureTaskIds].toSorted();
+            if (abortOnDependencyFailureTaskIds.length > 0) {
+                const sortedAbortIds = abortOnDependencyFailureTaskIds.toSorted();
                 const dependentTasks = await tx.pgTransaction
                     .select({ id: taskTable.id, namespace: taskTable.namespace, abortOnDependencyFailure: taskTable.abortOnDependencyFailure, status: taskTable.status })
                     .from(taskTable)
@@ -1054,7 +1031,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
                     .set({
                     unresolvedScheduleDependencies: greatest(0, sql `${taskTable.unresolvedScheduleDependencies} - ${updates.scheduleDecrement}`),
                     unresolvedCompleteDependencies: greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`),
-                    status: caseWhen(and(eq(taskTable.status, TaskStatus.Waiting), eq(greatest(0, sql `${taskTable.unresolvedScheduleDependencies} - ${updates.scheduleDecrement}`), 0)), enumValue(TaskStatus, taskStatus, TaskStatus.Pending)).else(caseWhen(and(eq(taskTable.status, TaskStatus.WaitingChildren), eq(greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`), 0)), enumValue(TaskStatus, taskStatus, TaskStatus.Completed)).else(taskTable.status)),
+                    status: caseWhen(and(eq(taskTable.status, TaskStatus.Waiting), eq(greatest(0, sql `${taskTable.unresolvedScheduleDependencies} - ${updates.scheduleDecrement}`), 0)), TaskStatus.Pending).else(caseWhen(and(eq(taskTable.status, TaskStatus.WaitingChildren), eq(greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`), 0)), TaskStatus.Completed).else(taskTable.status)),
                     progress: caseWhen(and(eq(taskTable.status, TaskStatus.WaitingChildren), eq(greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`), 0)), 1).else(taskTable.progress),
                     completeTimestamp: caseWhen(and(eq(taskTable.status, TaskStatus.WaitingChildren), eq(greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`), 0)), TRANSACTION_TIMESTAMP).else(taskTable.completeTimestamp),
                     token: caseWhen(and(eq(taskTable.status, TaskStatus.WaitingChildren), eq(greatest(0, sql `${taskTable.unresolvedCompleteDependencies} - ${updates.completeDecrement}`), 0)), null).else(taskTable.token),
@@ -1086,47 +1063,50 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async maintenance(options) {
-        await Promise.allSettled([
-            this.processExpirations(options),
-            this.processZombieRetries(options),
-            this.processZombieExhaustions(options),
-            this.processHardTimeouts(options),
-            this.processPriorityAging(options),
-        ]);
+        await this.processExpirations(options);
+        await this.processZombieRetries(options);
+        await this.processZombieExhaustions(options);
+        await this.processHardTimeouts(options);
+        await this.processPriorityAging(options);
         await this.performArchival(options);
         await this.performArchivePurge(options);
     }
     async performArchival(options) {
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
         while (true) {
-            const archivedCount = await this.#repository.useTransaction(options?.transaction, async (tx) => {
-                const childTaskTable = aliasedTable(taskTable, 'childTask');
-                const rowsToArchive = await tx.pgTransaction
-                    .select()
-                    .from(taskTable)
-                    .where(and(eq(taskTable.namespace, this.#namespace), inArray(taskTable.status, terminalStatuses), lte(taskTable.completeTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.retention, 'milliseconds')}`), notExists(tx.pgTransaction
-                    .select({ id: childTaskTable.id })
-                    .from(childTaskTable)
-                    .where(eq(childTaskTable.parentId, taskTable.id))), notExists(tx.pgTransaction
-                    .select({ taskId: taskDependencyTable.taskId })
-                    .from(taskDependencyTable)
-                    .where(eq(taskDependencyTable.dependencyTaskId, taskTable.id)))))
-                    .limit(1000)
-                    .for('update', { skipLocked: true });
-                if (rowsToArchive.length > 0) {
-                    const rowsToArchiveIds = getEntityIds(rowsToArchive);
-                    await tx.pgTransaction.insert(taskArchiveTable).values(rowsToArchive);
-                    await tx.pgTransaction.delete(taskTable).where(inArray(taskTable.id, rowsToArchiveIds));
-                }
-                return rowsToArchive.length;
-            });
-            if (archivedCount < 1000) {
+            const childTaskTable = aliasedTable(taskTable, 'childTask');
+            const selection = session.$with('selection').as((qb) => qb
+                .select({ id: taskTable.id })
+                .from(taskTable)
+                .where(and(eq(taskTable.namespace, this.#namespace), inArray(taskTable.status, terminalStatuses), lte(taskTable.completeTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.retention, 'milliseconds')}`), notExists(session
+                .select({ id: childTaskTable.id })
+                .from(childTaskTable)
+                .where(eq(childTaskTable.parentId, taskTable.id))), notExists(session
+                .select({ taskId: taskDependencyTable.taskId })
+                .from(taskDependencyTable)
+                .where(eq(taskDependencyTable.dependencyTaskId, taskTable.id)))))
+                .limit(1000)
+                .for('update', { skipLocked: true }));
+            const deleted = session.$with('deleted').as(() => session
+                .delete(taskTable)
+                .where(inArray(taskTable.id, session.select().from(selection)))
+                .returning());
+            const inserted = session.$with('inserted').as(() => session
+                .insert(taskArchiveTable)
+                .select(session.select().from(deleted))
+                .returning({ id: taskArchiveTable.id }));
+            const [result] = await session
+                .with(selection, deleted, inserted)
+                .select({ count: count() })
+                .from(inserted);
+            if ((result?.count ?? 0) < 1000) {
                 break;
             }
         }
     }
     async performArchivePurge(options) {
-        const session = options?.transaction?.pgTransaction ?? this.#database;
-        const selection = session.$with('archive_purge_selection').as((qb) => qb
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
+        const selection = session.$with('selection').as((qb) => qb
             .select({ id: taskArchiveTable.id })
             .from(taskArchiveTable)
             .where(and(eq(taskArchiveTable.namespace, this.#namespace), lte(taskArchiveTable.completeTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.archiveRetention, 'milliseconds')}`)))
@@ -1143,7 +1123,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async processExpirations(options) {
-        const expiredSelection = this.#database.$with('expired_selection').as((qb) => qb
+        const expiredSelection = this.#repository.session.$with('expired_selection').as((qb) => qb
             .select({ id: taskTable.id })
             .from(taskTable)
             .where(and(eq(taskTable.namespace, this.#namespace), inArray(taskTable.status, queueableOrWaitableStatuses), lt(taskTable.timeToLive, TRANSACTION_TIMESTAMP)))
@@ -1173,7 +1153,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async processZombieRetries(options) {
-        const session = options?.transaction?.pgTransaction ?? this.#database;
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
         const zombieRetrySelection = session.$with('zombie_retry_selection').as((qb) => qb
             .select({ id: taskTable.id })
             .from(taskTable)
@@ -1189,7 +1169,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
                 token: null,
                 visibilityDeadline: null,
                 startTimestamp: null,
-                scheduleTimestamp: sql `${TRANSACTION_TIMESTAMP} + ${interval(this.retryDelayMinimum, 'milliseconds')}`, // Simple backoff for zombies
+                scheduleTimestamp: sql `${TRANSACTION_TIMESTAMP} + ${interval(least(this.retryDelayMaximum, sql `${this.retryDelayMinimum} * ${power(this.retryDelayGrowth, sql `${taskTable.tries} - 1`)}`), 'milliseconds')}`,
                 error: jsonbBuildObject({ code: 'VisibilityTimeout', message: 'Worker Lost', lastError: taskTable.error }),
             })
                 .where(inArray(taskTable.id, session.select().from(zombieRetrySelection)))
@@ -1200,7 +1180,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async processZombieExhaustions(options) {
-        const zombieExhaustionSelection = this.#database.$with('zombie_exhaustion_selection').as((qb) => qb
+        const zombieExhaustionSelection = this.#repository.session.$with('selection').as((qb) => qb
             .select({ id: taskTable.id })
             .from(taskTable)
             .where(and(eq(taskTable.namespace, this.#namespace), eq(taskTable.status, TaskStatus.Running), lt(taskTable.visibilityDeadline, TRANSACTION_TIMESTAMP), gte(taskTable.tries, this.maxTries)))
@@ -1231,7 +1211,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async processHardTimeouts(options) {
-        const timeoutSelection = this.#database.$with('timeout_selection').as((qb) => qb
+        const timeoutSelection = this.#repository.session.$with('selection').as((qb) => qb
             .select({ id: taskTable.id })
             .from(taskTable)
             .where(and(eq(taskTable.namespace, this.#namespace), eq(taskTable.status, TaskStatus.Running), lt(taskTable.startTimestamp, sql `${TRANSACTION_TIMESTAMP} - ${interval(this.maxExecutionTime, 'milliseconds')}`)))
@@ -1262,7 +1242,7 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async processPriorityAging(options) {
-        const session = options?.transaction?.pgTransaction ?? this.#database;
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
         const agingSelection = session.$with('aging_selection').as((qb) => qb
             .select({ id: taskTable.id })
             .from(taskTable)
@@ -1285,8 +1265,8 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
         }
     }
     async restart(id, options) {
-        const repository = this.#repository.withOptionalTransaction(options?.transaction);
-        await repository.session
+        const session = options?.transaction?.pgTransaction ?? this.#repository.session;
+        const [updatedTask] = await session
             .update(taskTable)
             .set({
             status: TaskStatus.Pending,
@@ -1301,7 +1281,12 @@ let PostgresTaskQueue = class PostgresTaskQueue extends TaskQueue {
             priorityAgeTimestamp: TRANSACTION_TIMESTAMP,
             state: (options?.resetState == true) ? null : undefined,
         })
-            .where(and(eq(taskTable.id, id), or(inArray(taskTable.status, queueableStatuses), inArray(taskTable.status, terminalStatuses), lt(taskTable.visibilityDeadline, TRANSACTION_TIMESTAMP))));
+            .where(and(eq(taskTable.id, id), or(inArray(taskTable.status, queueableStatuses), inArray(taskTable.status, terminalStatuses), lt(taskTable.visibilityDeadline, TRANSACTION_TIMESTAMP))))
+            .returning();
+        if (isUndefined(updatedTask)) {
+            throw new NotFoundError('Task not found or not in a restartable state.');
+        }
+        this.notify(updatedTask.namespace);
     }
     notify(namespace = this.#namespace) {
         this.#messageBus.publishAndForget(namespace);

package/task-queue/tests/optimization-edge-cases.test.d.ts

@@ -0,0 +1 @@
+export {};

package/task-queue/tests/optimization-edge-cases.test.js

@@ -0,0 +1,127 @@
+import { and, eq, sql } from 'drizzle-orm';
+import { beforeAll, describe, expect, vi } from 'vitest';
+import { inject } from '../../injector/index.js';
+import { TRANSACTION_TIMESTAMP } from '../../orm/index.js';
+import { injectRepository } from '../../orm/server/index.js';
+import { setupIntegrationTest, testInInjector } from '../../testing/index.js';
+import { timeout } from '../../utils/timing.js';
+import { taskArchive as taskArchiveTable, taskDependency as taskDependencyTable, task as taskTable } from '../postgres/schemas.js';
+import { PostgresTaskQueue } from '../postgres/task-queue.js';
+import { PostgresTask, PostgresTaskArchive } from '../postgres/task.model.js';
+import { TaskDependencyType, TaskStatus } from '../task-queue.js';
+describe('Task Queue Optimization Edge Cases', () => {
+    let context;
+    beforeAll(async () => {
+        context = await setupIntegrationTest({ modules: { taskQueue: true } });
+    });
+    testInInjector('should notify unique namespaces exactly once in incrementCounters', () => context.injector, async () => {
+        const q1 = inject(PostgresTaskQueue, 'ns-1');
+        const q2 = inject(PostgresTaskQueue, 'ns-2');
+        const notifySpy = vi.spyOn(PostgresTaskQueue.prototype, 'notify');
+        const [t1] = await q1.enqueueMany([{ type: 't1', data: {} }], { returnTasks: true });
+        const [t2] = await q2.enqueueMany([{ type: 't2', data: {} }], { returnTasks: true });
+        notifySpy.mockClear();
+        await q1.incrementCounters([
+            { taskId: t1.id, dependencyTaskId: 'some-dep', type: TaskDependencyType.Schedule },
+            { taskId: t2.id, dependencyTaskId: 'some-dep', type: TaskDependencyType.Schedule },
+            { taskId: t1.id, dependencyTaskId: 'other-dep', type: TaskDependencyType.Schedule },
+        ]);
+        const notifiedNamespaces = notifySpy.mock.calls.map(call => call[0]);
+        expect(notifiedNamespaces).toContain('ns-1');
+        expect(notifiedNamespaces).toContain('ns-2');
+        expect(notifiedNamespaces.filter(n => n == 'ns-1').length).toBe(1);
+        expect(notifiedNamespaces.filter(n => n == 'ns-2').length).toBe(1);
+    });
+    testInInjector('should resolve edge and abort on unmatched terminal status in resolveDependenciesMany', () => context.injector, async () => {
+        const queue = inject(PostgresTaskQueue, 'test-namespace');
+        const repository = injectRepository(PostgresTask);
+        const [parent] = await queue.enqueueMany([{ type: 'parent', data: {}, abortOnDependencyFailure: true }], { returnTasks: true });
+        const [child] = await queue.enqueueMany([{ type: 'child', data: {} }], { returnTasks: true });
+        await repository.session.insert(taskDependencyTable).values({
+            taskId: parent.id,
+            dependencyTaskId: child.id,
+            type: TaskDependencyType.Schedule,
+            requiredStatuses: [TaskStatus.Completed],
+        });
+        await repository.session.update(taskTable).set({ unresolvedScheduleDependencies: 1, status: TaskStatus.Waiting }).where(eq(taskTable.id, parent.id));
+        await queue.resolveDependenciesMany([{ id: child.id, status: TaskStatus.Dead, namespace: 'test-namespace' }]);
+        const updatedParent = await queue.getTask(parent.id);
+        expect(updatedParent.status).toBe(TaskStatus.Skipped);
+    });
+    testInInjector('should handle hard timeout during touch', () => context.injector, async () => {
+        // Configure with small maxExecutionTime
+        const queue = inject(PostgresTaskQueue, { namespace: 'timeout-test', maxExecutionTime: 10 });
+        const [task] = await queue.enqueueMany([{ type: 'test', data: {} }], { returnTasks: true });
+        const [runningTask] = await queue.dequeueMany(1);
+        await timeout(50);
+        const result = await queue.touch(runningTask);
+        expect(result).toBeUndefined();
+        const updated = await queue.getTask(runningTask.id);
+        expect(updated.status).toBe(TaskStatus.TimedOut);
+    });
+    testInInjector('should handle non-existent tasks in complete and fail', () => context.injector, async () => {
+        const queue = inject(PostgresTaskQueue, 'missing-test');
+        const fakeTask = { id: crypto.randomUUID(), token: crypto.randomUUID(), tries: 0 };
+        await expect(queue.complete(fakeTask)).resolves.toBeUndefined();
+        await expect(queue.fail(fakeTask, new Error('fail'))).resolves.toBeUndefined();
+    });
+    testInInjector('should handle terminal tasks with no dependents in resolveDependenciesMany', () => context.injector, async () => {
+        const queue = inject(PostgresTaskQueue, 'no-deps-test');
+        const [task] = await queue.enqueueMany([{ type: 'test', data: {} }], { returnTasks: true });
+        await queue.resolveDependenciesMany([{ id: task.id, status: TaskStatus.Completed }]);
+    });
+    testInInjector('should handle archival and purge in maintenance', () => context.injector, async () => {
+        const namespace = `archival-test-${crypto.randomUUID()}`;
+        // Configure with small retention
+        const queue = inject(PostgresTaskQueue, { namespace, retention: 1000, archiveRetention: 1000 });
+        const repository = injectRepository(PostgresTask);
+        const archiveRepository = injectRepository(PostgresTaskArchive);
+        const [task] = await queue.enqueueMany([{ type: 'test', data: {} }], { returnTasks: true });
+        await repository.session.update(taskTable)
+            .set({
+            status: TaskStatus.Completed,
+            completeTimestamp: sql `${TRANSACTION_TIMESTAMP} - interval '2 seconds'`,
+        })
+            .where(and(eq(taskTable.id, task.id), eq(taskTable.namespace, namespace)));
+        await queue.performArchival();
+        const archived = await archiveRepository.load(task.id);
+        expect(archived).toBeDefined();
+        await repository.session.update(taskArchiveTable)
+            .set({ completeTimestamp: sql `${TRANSACTION_TIMESTAMP} - interval '2 seconds'` })
+            .where(and(eq(taskArchiveTable.id, task.id), eq(taskArchiveTable.namespace, namespace)));
+        await queue.performArchivePurge();
+        const purged = await archiveRepository.load(task.id).catch(() => undefined);
+        expect(purged).toBeUndefined();
+    });
+    testInInjector('should notify on restart', () => context.injector, async () => {
+        const namespace = `restart-test-${crypto.randomUUID()}`;
+        const queue = inject(PostgresTaskQueue, namespace);
+        const notifySpy = vi.spyOn(PostgresTaskQueue.prototype, 'notify');
+        const [task] = await queue.enqueueMany([{ type: 'test', data: {} }], { returnTasks: true });
+        await queue.cancelMany([task.id]);
+        notifySpy.mockClear();
+        await queue.restart(task.id);
+        expect(notifySpy).toHaveBeenCalledWith(namespace);
+    });
+    testInInjector('should use exponential backoff for zombies', () => context.injector, async () => {
+        const namespace = `zombie-backoff-${crypto.randomUUID()}`;
+        // Configure with standard growth
+        const queue = inject(PostgresTaskQueue, { namespace, retryDelayMinimum: 1000, retryDelayGrowth: 2 });
+        const repository = injectRepository(PostgresTask);
+        const [task] = await queue.enqueueMany([{ type: 'test', data: {} }], { returnTasks: true });
+        await queue.dequeueMany(1);
+        await repository.session.update(taskTable)
+            .set({
+            status: TaskStatus.Running,
+            visibilityDeadline: sql `${TRANSACTION_TIMESTAMP} - interval '1 minute'`,
+            tries: 2,
+        })
+            .where(and(eq(taskTable.id, task.id), eq(taskTable.namespace, namespace)));
+        await queue.processZombieRetries();
+        const updated = await queue.getTask(task.id);
+        expect(updated.status).toBe(TaskStatus.Retrying);
+        const delay = updated.scheduleTimestamp - Date.now();
+        expect(delay).toBeGreaterThan(1000);
+        expect(delay).toBeLessThan(3000);
+    });
+});

package/test1.js

@@ -14,7 +14,7 @@ import { assert } from './utils/type-guards.js';
 const config = {
     database: {
         host: configParser.string('DATABASE_HOST', '127.0.0.1'),
-        port: configParser.positiveInteger('DATABASE_PORT', 5432),
+        port: configParser.positiveInteger('DATABASE_PORT', 15433),
         user: configParser.string('DATABASE_USER', 'tstdl'),
         pass: configParser.string('DATABASE_PASS', 'wf7rq6glrk5jykne'),
         database: configParser.string('DATABASE_NAME', 'tstdl'),

package/test4.js

@@ -8,7 +8,7 @@ import { boolean, positiveInteger, string } from './utils/config-parser.js';
 const config = {
     database: {
         host: string('DATABASE_HOST', '127.0.0.1'),
-        port: positiveInteger('DATABASE_PORT', 5432),
+        port: positiveInteger('DATABASE_PORT', 15433),
         user: string('DATABASE_USER', 'tstdl'),
         pass: string('DATABASE_PASS', 'wf7rq6glrk5jykne'),
         database: string('DATABASE_NAME', 'tstdl'),

package/test5.js

@@ -3,12 +3,18 @@ import { Application } from './application/application.js';
 import { provideModule, provideSignalHandler } from './application/index.js';
 import { PrettyPrintLogFormatter } from './logger/index.js';
 import { provideConsoleLogTransport } from './logger/transports/console.js';
+import { TaskQueue } from './task-queue/task-queue.js';
+import { setupIntegrationTest } from './testing/integration-setup.js';
+import { createArray } from './utils/array/array.js';
+import { timedBenchmarkAsync } from './utils/benchmark.js';
 async function main(_cancellationSignal) {
-    const arr = [1, 2, 3, 4, 5];
-    for (const item of arr) {
-        console.log(item);
-        arr.push(arr.at(-1) + 1);
-    }
+    const { injector } = await setupIntegrationTest({ modules: { taskQueue: true } });
+    const queue1 = injector.resolve(TaskQueue, 'namespace-1');
+    const batch = createArray(1000, (i) => ({ type: 'test', data: { index: i } }));
+    const enqueueResult = await timedBenchmarkAsync(1000, async () => {
+        await queue1.enqueueMany(batch);
+    });
+    console.log(enqueueResult.operationsPerMillisecond * batch.length * 1000, 'items/s');
 }
 Application.run('Test', [
     provideConsoleLogTransport(PrettyPrintLogFormatter),

package/testing/integration-setup.js

@@ -48,7 +48,7 @@ export async function setupIntegrationTest(options = {}) {
     // 2. Database Config
     const dbConfig = {
         host: configParser.string('DATABASE_HOST', '127.0.0.1'),
-        port: configParser.positiveInteger('DATABASE_PORT', 5432),
+        port: configParser.positiveInteger('DATABASE_PORT', 15433),
         user: configParser.string('DATABASE_USER', 'tstdl'),
         password: configParser.string('DATABASE_PASS', 'wf7rq6glrk5jykne'),
         database: configParser.string('DATABASE_NAME', 'tstdl'),