@goodandready/dsh-key-rotation 0.8.29 → 0.8.30

package/lib/rotate.js

@@ -16,6 +16,34 @@ import {
 } from './pool.js';
 import { pickCascadeFallback, CASCADE_MAX_DEPTH } from './cascade.js';
 import { nextQuotaReset } from './quota-window.js';
+import {
+  isModelQuotaAvailable,
+  hasModelQuotaConfig,
+  anyModelQuotaAvailable,
+  consumeModelTokens,
+  extractUsageTokens,
+  formatModelQuotaExhaustion,
+} from './model-quota.js';
+
+export const LOCAL_QUOTA_CODE = 'LOCAL_MODEL_QUOTA_EXHAUSTED';
+
+export async function* preserveDispatchContext(iterable, storage, store) {
+  if (!storage || typeof storage.getStore !== 'function') {
+    yield* iterable;
+    return;
+  }
+  const iterator = iterable[Symbol.asyncIterator]();
+  const step = (method, arg) => storage.run(store, () => iterator[method](arg));
+  try {
+    while (true) {
+      const next = await step('next');
+      if (next.done) return next.value;
+      yield next.value;
+    }
+  } finally {
+    if (typeof iterator.return === 'function') await step('return');
+  }
+}
 
 /**
  * Create the rotate(options, pool) async generator used by llm/stream.
@@ -53,8 +81,44 @@ export function createRotate(deps) {
       let activePermit = null;
       const reqStore = { pool, pickedRef: undefined, startMs: now() };
       setRotateStartMs(reqStore.startMs);
+      const wallNow = Date.now();
+      const quotaManaged = hasModelQuotaConfig(pool);
       let attemptList = (pool.weightedRefs ?? pool.refs).slice();
       attemptList = attemptList.filter((r) => !isKeyPaused(pool, r) && !isKeyRevoked(pool, r));
+      if (quotaManaged) {
+        attemptList = attemptList.filter((ref) => isModelQuotaAvailable(pool, ref, wallNow));
+        if (attemptList.length === 0) {
+          pool.state.lastExhaustionAt = now();
+          pool.state.exhaustionCount = (pool.state.exhaustionCount ?? 0) + 1;
+          logWarn(`[dsh-key-rotation] ${options.provider}/${options.model ?? '-'}: local model token budget exhausted for all ${pool.refs.length} key(s) — no upstream request sent`);
+          const runtimeQ = buildRuntime();
+          notifyExhaustion(runtimeQ, pool, { provider: options.provider });
+          const maxDepth = runtimeQ.cascadeMaxDepth ?? CASCADE_MAX_DEPTH;
+          const currentDepth = options.__cascadeDepth ?? (options.__isCascade ? 1 : 0);
+          if (currentDepth < maxDepth && Array.isArray(runtimeQ.cascade) && runtimeQ.cascade.length > 0) {
+            const fb = pickCascadeFallback(options.provider, runtimeQ, runtimeQ.providerToPool, options.model);
+            if (fb && fb.pool && fb.pool !== pool) {
+              pool.state.lastReason = 'CASCADE';
+              pool.state.lastSwitchAt = now();
+              pool.state.cascadeCount = (pool.state.cascadeCount ?? 0) + 1;
+              pushEvent(pool, reqStore.pickedRef ?? 'pool', 'CASCADE_TRIGGERED', 0);
+              const innerCascade = rotate({
+                ...options,
+                provider: fb.provider,
+                model: fb.model || options.model,
+                __isCascade: true,
+                __cascadeDepth: currentDepth + 1,
+              }, fb.pool);
+              for await (const chunk of innerCascade) {
+                yield chunk;
+              }
+              return;
+            }
+          }
+          yield finishError(LOCAL_QUOTA_CODE, formatModelQuotaExhaustion(options.provider, pool));
+          return;
+        }
+      }
       const strategy = pool.routingStrategy ?? runtime0.routingStrategy ?? 'round-robin';
       if (strategy !== 'round-robin') {
         attemptList = sortAttemptList(attemptList, strategy, { latencyHistogram, concurrencyTracker, weights: pool.weightsMap });
@@ -80,6 +144,7 @@ export function createRotate(deps) {
 
       const penalizeRef = (targetRef, errCode, errMsg) => {
         if (!targetRef) return;
+        if (errCode === LOCAL_QUOTA_CODE) return;
         const _retry = parseRetryAfter(errMsg);
         const _base = pool.cooldownMs ?? cooldownMs;
         const _max = pool.maxCooldownMs ?? maxCooldownMs;
@@ -144,6 +209,11 @@ export function createRotate(deps) {
               return llm.stream({ ...options, [MARKER]: true });
             });
           } catch (e) {
+            if (e?.localQuota || e?.code === LOCAL_QUOTA_CODE) {
+              lastFailure = finishError(LOCAL_QUOTA_CODE, String(e?.message ?? ''));
+              logWarn(`[dsh-key-rotation] ${options.provider}/${options.model ?? '-'}: local model token budget blocks dispatch — no upstream request sent`);
+              break;
+            }
             const curRef = reqStore.pickedRef ?? pool.state.lastUsed;
             penalizeRef(curRef, e?.code ?? 'TRANSPORT', String(e?.message ?? ''));
             lastFailure = finishError(e?.code ?? 'TRANSPORT',
@@ -156,7 +226,7 @@ export function createRotate(deps) {
           const acquired = (_pickedRef && runtime0.concurrencyLimit > 0) ? concurrencyTracker.acquire(_pickedRef) : false;
           try {
             try {
-              for await (const chunk of inner) {
+              for await (const chunk of preserveDispatchContext(inner, dispatchStorage, reqStore)) {
                 // Only actual content deltas lock the stream (no more rotation).
                 // Structural/metadata chunks (block-start/end, usage) do not.
                 if (chunk && (chunk.type === 'text-delta' || chunk.type === 'reasoning-delta' || chunk.type === 'tool-call-delta')) {
@@ -232,6 +302,20 @@ export function createRotate(deps) {
                   if (activeRef) {
                     recordSuccess(pool, activeRef, now());
                   }
+                  if (activeRef && quotaManaged) {
+                    const charged = consumeModelTokens(
+                      pool, activeRef, extractUsageTokens(chunk.usage),
+                      runtime0.quotaResetWindow, Date.now(),
+                    );
+                    if (charged) {
+                      if (typeof schedulePersist === 'function') {
+                        try { schedulePersist(); } catch (_) { /* best-effort persistence */ }
+                      }
+                      if (charged.exhausted) {
+                        logWarn(`[dsh-key-rotation] ${options.provider}/${options.model ?? '-'}: key ${activeRef} reached its local token budget (${charged.used}/${charged.limit}) — excluded from this model only`);
+                      }
+                    }
+                  }
 
                   // Proactive rate-limit (#115, #303): if response headers say this key is near
                   // its quota or has retry-after, cool it down so the NEXT request starts on a different key.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@goodandready/dsh-key-rotation",
-  "version": "0.8.29",
+  "version": "0.8.30",
   "packageManager": "pnpm@10.33.2",
   "description": "Per-provider API key rotation for DeepSeek Harness: a key pool per provider, auto-created clone routes, and switching to the next key on quota/rate-limit errors. Includes a Settings section (Key Rotation) to edit the key pools, cooldown and switch codes.",
   "keywords": [