@tstdl/base 0.93.160 → 0.93.162

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

@@ -1,12 +1,20 @@
 import { vertexAI } from '@genkit-ai/google-genai';
 import { GenkitError, modelRef } from 'genkit';
 import { genkitPlugin } from 'genkit/plugin';
-import { shuffle } from '../../utils/array/index.js';
+import { distinct, 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');
+const defaultTokenLimitThreshold = 131_072;
 export function vertexAiMultiLocation(options) {
+    if (options.locations.length == 0) {
+        throw new GenkitError({
+            status: 'INVALID_ARGUMENT',
+            message: 'At least one location must be provided for vertexAiMultiLocation',
+        });
+    }
+    const tokenLimitThreshold = options.tokenLimitThreshold ?? defaultTokenLimitThreshold;
     const locationConfigs = options.locations.map((location) => {
         const circuitBreakerKey = `genkit:vertex-ai:location:${location}`;
         const tokenLimitCircuitBreakerKey = `${circuitBreakerKey}:token-limit`;
@@ -41,16 +49,19 @@ export function vertexAiMultiLocation(options) {
             const shuffledConfigs = shuffle([...locationConfigs]);
             let lastError;
             let isLargeRequest = false;
+            const skippedLocations = [];
             for (const { location, circuitBreaker, tokenLimitCircuitBreaker } of shuffledConfigs) {
                 const check = await circuitBreaker.check();
                 if (!check.allowed) {
                     options.logger.warn(`Location ${location} is currently unhealthy. Skipping...`);
+                    skippedLocations.push({ location, reason: 'unhealthy' });
                     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...`);
+                        skippedLocations.push({ location, reason: 'known to have low token limits' });
                         continue;
                     }
                 }
@@ -65,6 +76,10 @@ export function vertexAiMultiLocation(options) {
                         onChunk: streamingCallback,
                     });
                     await circuitBreaker.recordSuccess();
+                    const isLargeSuccess = isLargeRequest || ((result.usage?.inputTokens ?? 0) > tokenLimitThreshold);
+                    if (isLargeSuccess) {
+                        await tokenLimitCircuitBreaker.recordSuccess();
+                    }
                     return result;
                 }
                 catch (error) {
@@ -88,6 +103,14 @@ export function vertexAiMultiLocation(options) {
                     }
                 }
             }
+            if (isNullOrUndefined(lastError)) {
+                const uniqueReasons = distinct(skippedLocations.map((s) => s.reason));
+                const reasonsString = uniqueReasons.join(' or ');
+                throw new GenkitError({
+                    status: 'UNAVAILABLE',
+                    message: `All locations were skipped because they are ${reasonsString}`,
+                });
+            }
             throw lastError;
         });
     };

package/ai/genkit/tests/multi-region.test.js

@@ -138,4 +138,42 @@ describe('Genkit vertexai-multi-location Plugin Tests', () => {
         });
         expect(response2.text).toBe('success from cb-success');
     });
+    it('should throw if no locations are provided', async () => {
+        expect(() => {
+            vertexAiMultiLocation({
+                locations: [],
+                circuitBreakerProvider: cbProvider,
+                logger,
+            });
+        }).toThrow('At least one location must be provided');
+    });
+    it('should throw if all locations are unhealthy', async () => {
+        const ai2 = genkit({
+            plugins: [
+                vertexAiMultiLocation({
+                    locations: ['unhealthy-1'],
+                    circuitBreakerProvider: cbProvider,
+                    logger,
+                    circuitBreakerConfig: { resetTimeout: 1000000, threshold: 1 },
+                }),
+            ],
+        });
+        // Manually trip the circuit breaker
+        const cb = cbProvider.provide('genkit:vertex-ai:location:unhealthy-1', { threshold: 1, resetTimeout: 1000000 });
+        await cb.recordFailure();
+        ai2.defineModel({
+            name: 'vertexai/gemini-2.5-flash',
+        }, async () => {
+            return {
+                message: {
+                    role: 'model',
+                    content: [{ text: 'success' }],
+                },
+            };
+        });
+        await expect(ai2.generate({
+            model: 'vertexai-multi-location/gemini-2.5-flash',
+            prompt: 'test',
+        })).rejects.toThrow('All locations were skipped because they are unhealthy');
+    });
 });

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

@@ -5,6 +5,7 @@ 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 { timeout } from '../../../utils/timing.js';
 import { vertexAiMultiLocation } from '../multi-region.plugin.js';
 vi.mock('#/utils/array/index.js', async (importOriginal) => {
     const actual = await importOriginal();
@@ -43,6 +44,7 @@ describe('Genkit vertexai-multi-location Token Limit Fallback Tests', () => {
                     circuitBreakerProvider: cbProvider,
                     logger,
                     circuitBreakerConfig: { resetTimeout: 1_000_000, threshold: 1 },
+                    tokenLimitCircuitBreakerConfig: { resetTimeout: 10, threshold: 1 },
                 }),
             ],
         });
@@ -135,4 +137,73 @@ describe('Genkit vertexai-multi-location Token Limit Fallback Tests', () => {
         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
     });
+    it('should record success for token limit breaker when it succeeds after being 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.';
+        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: ${request.config?.location}`);
+        });
+        // Initial shuffle is [region-1, region-2, region-3].
+        // Let's use a very short reset timeout for region-2.
+        const shortConfig = { threshold: 1, resetTimeout: 10 };
+        const tokenLimitCBShort = cbProvider.provide('genkit:vertex-ai:location:region-2:token-limit', shortConfig);
+        await tokenLimitCBShort.recordFailure();
+        // Small delay to ensure it goes to Half-Open
+        await timeout(50);
+        await ai.generate({
+            model: 'vertexai-multi-location/gemini-2.5-flash',
+            prompt: 'test',
+        });
+        expect(region1Called).toBe(true);
+        expect(region2Called).toBe(true); // Should NOT be skipped because it is Half-Open
+        expect(region3Called).toBe(false);
+        const status = await tokenLimitCBShort.check();
+        expect(status.state).toBe(CircuitBreakerState.Closed);
+    });
+    it('should record success for token limit breaker via usage-based learning (proactive success)', async () => {
+        // We want to verify that if a request has many tokens, it records success even if isLargeRequest was false.
+        const tokenLimitThreshold = 131_072;
+        const aiWithThreshold = genkit({
+            plugins: [
+                vertexAiMultiLocation({
+                    locations: ['region-1'],
+                    circuitBreakerProvider: cbProvider,
+                    logger,
+                    tokenLimitThreshold,
+                }),
+            ],
+        });
+        // Trip the breaker first
+        const tokenLimitCB = cbProvider.provide('genkit:vertex-ai:location:region-1:token-limit', { threshold: 1, resetTimeout: 1000000 });
+        await tokenLimitCB.recordFailure();
+        aiWithThreshold.defineModel({ name: 'vertexai/gemini-2.5-flash' }, async () => ({
+            message: { role: 'model', content: [{ text: 'success' }] },
+            usage: { inputTokens: tokenLimitThreshold + 1 },
+        }));
+        // This request is NOT known to be large initially.
+        // But it returns usage > threshold.
+        await aiWithThreshold.generate({
+            model: 'vertexai-multi-location/gemini-2.5-flash',
+            prompt: 'test',
+        });
+        const status = await tokenLimitCB.check();
+        expect(status.state).toBe(CircuitBreakerState.Closed);
+    });
 });

package/ai/genkit/types.d.ts

@@ -12,4 +12,10 @@ export declare abstract class VertexAiMultiLocationOptions {
      * By default, a threshold of 1 and a reset timeout of 15 minutes is used.
      */
     tokenLimitCircuitBreakerConfig?: Partial<CircuitBreakerConfig>;
+    /**
+     * Threshold of input tokens after which a request is considered large.
+     * A successful request with more tokens than this threshold will record a success for the token limit circuit breaker.
+     * Defaults to 131,072.
+     */
+    tokenLimitThreshold?: number;
 }

package/ai/genkit/types.js

@@ -11,4 +11,10 @@ export class VertexAiMultiLocationOptions {
      * By default, a threshold of 1 and a reset timeout of 15 minutes is used.
      */
     tokenLimitCircuitBreakerConfig;
+    /**
+     * Threshold of input tokens after which a request is considered large.
+     * A successful request with more tokens than this threshold will record a success for the token limit circuit breaker.
+     * Defaults to 131,072.
+     */
+    tokenLimitThreshold;
 }

package/circuit-breaker/circuit-breaker.d.ts

@@ -13,7 +13,7 @@ export type CircuitBreakerConfig = {
 export type CircuitBreakerCheckResult = {
     allowed: boolean;
     state: CircuitBreakerState;
-    isProbe?: boolean;
+    isProbe: boolean;
 };
 export type CircuitBreakerArgument = string | (CircuitBreakerConfig & {
     key: string;

package/circuit-breaker/postgres/circuit-breaker.d.ts

@@ -1,8 +1,12 @@
+import { afterResolve } from '../../injector/index.js';
 import { CircuitBreaker, type CircuitBreakerCheckResult } from '../circuit-breaker.js';
 export declare class PostgresCircuitBreakerService extends CircuitBreaker {
     #private;
+    private static checkStatement;
+    [afterResolve](): void;
     check(): Promise<CircuitBreakerCheckResult>;
     recordSuccess(): Promise<void>;
     recordFailure(): Promise<void>;
     recordFailures(count: number): Promise<void>;
+    private getPreparedCheckStatement;
 }

package/circuit-breaker/postgres/circuit-breaker.js

@@ -4,40 +4,42 @@ var __decorate = (this && this.__decorate) || function (decorators, target, key,
     else for (var i = decorators.length - 1; i >= 0; i--) if (d = decorators[i]) r = (c < 3 ? d(r) : c > 3 ? d(target, key, r) : d(target, key)) || r;
     return c > 3 && r && Object.defineProperty(target, key, r), r;
 };
-import { sql } from 'drizzle-orm';
-import { injectArgument, provide, Singleton } from '../../injector/index.js';
-import { interval, TRANSACTION_TIMESTAMP } from '../../orm/index.js';
+var PostgresCircuitBreakerService_1;
+import { and, eq, lte, sql, isNotNull as sqlIsNotNull } from 'drizzle-orm';
+import { afterResolve, injectArgument, provide, Singleton } from '../../injector/index.js';
+import { coalesce, interval, TRANSACTION_TIMESTAMP } from '../../orm/index.js';
 import { DatabaseConfig, injectRepository } from '../../orm/server/index.js';
-import { currentTimestamp } from '../../utils/date-time.js';
-import { isDefined, isString, isUndefined } from '../../utils/type-guards.js';
+import { isString, isUndefined } from '../../utils/type-guards.js';
 import { millisecondsPerSecond } from '../../utils/units.js';
 import { CircuitBreaker, CircuitBreakerState } from '../circuit-breaker.js';
 import { PostgresCircuitBreaker } from './model.js';
 import { PostgresCircuitBreakerModuleConfig } from './module.js';
+import { circuitBreaker } from './schemas.js';
 let PostgresCircuitBreakerService = class PostgresCircuitBreakerService extends CircuitBreaker {
+    static { PostgresCircuitBreakerService_1 = this; }
+    static checkStatement;
     #repository = injectRepository(PostgresCircuitBreaker);
     #arg = injectArgument(this);
     #key = isString(this.#arg) ? this.#arg : this.#arg.key;
     #threshold = (isString(this.#arg) ? undefined : this.#arg.threshold) ?? 5;
     #resetTimeout = (isString(this.#arg) ? undefined : this.#arg.resetTimeout) ?? 30 * millisecondsPerSecond;
+    [afterResolve]() {
+        PostgresCircuitBreakerService_1.checkStatement ??= this.getPreparedCheckStatement();
+    }
     async check() {
-        return await this.#repository.transaction(async (tx) => {
-            const breaker = await this.#repository.withTransaction(tx).tryLoadByQuery({ key: this.#key });
-            if (isUndefined(breaker) || breaker.state === CircuitBreakerState.Closed) {
-                return { allowed: true, state: CircuitBreakerState.Closed };
-            }
-            if (breaker.state === CircuitBreakerState.HalfOpen) {
-                return { allowed: true, state: CircuitBreakerState.HalfOpen, isProbe: false };
-            }
-            // State is Open
-            if (currentTimestamp() < (breaker.resetTimestamp ?? 0)) {
-                return { allowed: false, state: CircuitBreakerState.Open };
-            }
-            // Atomic transition from Open -> HalfOpen (The Probe)
-            const updated = await this.#repository.withTransaction(tx).tryUpdateByQuery({ key: this.#key, state: CircuitBreakerState.Open }, { state: CircuitBreakerState.HalfOpen });
-            const isProbe = isDefined(updated);
-            return { allowed: isProbe, state: isProbe ? CircuitBreakerState.HalfOpen : CircuitBreakerState.Open, isProbe };
-        });
+        const [result] = await PostgresCircuitBreakerService_1.checkStatement.execute({ key: this.#key });
+        // 1. Breaker doesn't exist or is Closed
+        if (isUndefined(result) || (result.state === CircuitBreakerState.Closed)) {
+            return { allowed: true, state: CircuitBreakerState.Closed, isProbe: false };
+        }
+        // 2. Atomic transition to HalfOpen succeeded. This request is the chosen probe!
+        if (result.isProbe) {
+            return { allowed: true, state: CircuitBreakerState.HalfOpen, isProbe: true };
+        }
+        // 3. Fallback: Catch-all for failed transitions.
+        // - If state is HalfOpen, someone else is probing. Reject.
+        // - If state is Open, timeout hasn't expired. Reject
+        return { allowed: false, state: result.state, isProbe: false };
     }
     async recordSuccess() {
         await this.#repository.tryDeleteByQuery({ key: this.#key });
@@ -46,7 +48,6 @@ let PostgresCircuitBreakerService = class PostgresCircuitBreakerService extends
         await this.recordFailures(1);
     }
     async recordFailures(count) {
-        const table = this.#repository.table;
         const initialTrip = count >= this.#threshold;
         const initialState = initialTrip ? CircuitBreakerState.Open : CircuitBreakerState.Closed;
         const initialResetTimestamp = initialTrip
@@ -58,19 +59,40 @@ let PostgresCircuitBreakerService = class PostgresCircuitBreakerService extends
             failureCount: count,
             resetTimestamp: initialResetTimestamp,
         }, {
-            failureCount: sql `${table.failureCount} + ${count}`,
+            failureCount: sql `${circuitBreaker.failureCount} + ${count}`,
             state: sql `CASE
-          WHEN ${table.failureCount} + ${count} >= ${this.#threshold} THEN ${CircuitBreakerState.Open}
-          ELSE ${table.state}
+          WHEN ${circuitBreaker.failureCount} + ${count} >= ${this.#threshold} THEN ${CircuitBreakerState.Open}
+          ELSE ${circuitBreaker.state}
         END`,
             resetTimestamp: sql `CASE
-          WHEN ${table.failureCount} + ${count} >= ${this.#threshold} THEN ${TRANSACTION_TIMESTAMP} + ${interval(this.#resetTimeout, 'milliseconds')}
-          ELSE ${table.resetTimestamp}
+          WHEN ${circuitBreaker.failureCount} + ${count} >= ${this.#threshold} THEN ${TRANSACTION_TIMESTAMP} + ${interval(this.#resetTimeout, 'milliseconds')}
+          ELSE ${circuitBreaker.resetTimestamp}
         END`,
         });
     }
+    getPreparedCheckStatement() {
+        const session = this.#repository.session;
+        const attemptUpdate = session.$with('attempt_update').as(() => session
+            .update(circuitBreaker)
+            .set({ state: CircuitBreakerState.HalfOpen })
+            .where(and(eq(circuitBreaker.key, sql.placeholder('key')), eq(circuitBreaker.state, CircuitBreakerState.Open), lte(circuitBreaker.resetTimestamp, TRANSACTION_TIMESTAMP)))
+            .returning({
+            key: circuitBreaker.key,
+            state: circuitBreaker.state,
+        }));
+        return session
+            .with(attemptUpdate)
+            .select({
+            state: coalesce(attemptUpdate.state, circuitBreaker.state),
+            isProbe: sqlIsNotNull(attemptUpdate.key),
+        })
+            .from(circuitBreaker)
+            .leftJoin(attemptUpdate, eq(circuitBreaker.key, attemptUpdate.key))
+            .where(eq(circuitBreaker.key, sql.placeholder('key')))
+            .prepare('circuit_breaker_check');
+    }
 };
-PostgresCircuitBreakerService = __decorate([
+PostgresCircuitBreakerService = PostgresCircuitBreakerService_1 = __decorate([
     Singleton({
         argumentIdentityProvider: (arg) => isString(arg) ? arg : arg.key,
         providers: [

package/circuit-breaker/tests/circuit-breaker.test.js

@@ -58,9 +58,7 @@ describe('Circuit Breaker (Standalone) Tests', () => {
         expect(probe.isProbe).toBe(true);
         // Subsequent check should be denied (Half-Open wait)
         const subsequent = await breaker.check();
-        expect(subsequent.allowed).toBe(true); // check() allows HalfOpen, but queue logic restricts it.
-        // Wait, let's verify PostgresCircuitBreakerService logic:
-        // if (breaker.state === CircuitBreakerState.HalfOpen) return { allowed: true, state: CircuitBreakerState.HalfOpen, isProbe: false };
+        expect(subsequent.allowed).toBe(false);
         expect(subsequent.state).toBe(CircuitBreakerState.HalfOpen);
         expect(subsequent.isProbe).toBe(false);
     });

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tstdl/base",
-  "version": "0.93.160",
+  "version": "0.93.162",
   "author": "Patrick Hein",
   "publishConfig": {
     "access": "public"
@@ -152,8 +152,8 @@
     "type-fest": "^5.4"
   },
   "peerDependencies": {
-    "@aws-sdk/client-s3": "^3.1002",
-    "@aws-sdk/s3-request-presigner": "^3.1002",
+    "@aws-sdk/client-s3": "^3.1003",
+    "@aws-sdk/s3-request-presigner": "^3.1003",
     "@genkit-ai/google-genai": "^1.29",
     "@google-cloud/storage": "^7.19",
     "@toon-format/toon": "^2.1.0",
@@ -184,7 +184,7 @@
   },
   "devDependencies": {
     "@biomejs/biome": "2.4",
-    "@stylistic/eslint-plugin": "5.9",
+    "@stylistic/eslint-plugin": "5.10",
     "@types/koa__router": "12.0",
     "@types/luxon": "3.7",
     "@types/mjml": "4.7",