@goodandready/dsh-key-rotation 0.8.76 → 0.8.78

package/lib/quota-runway.js

@@ -0,0 +1,222 @@
+// lib/quota-runway.js — Predictive Quota Runway and burn rate estimation (#549).
+// Estimates time to quota exhaustion and preemptively rotates keys before hitting HTTP 429.
+
+export const MINUTE = 60 * 1000;
+export const HOUR = 60 * MINUTE;
+export const DAY = 24 * HOUR;
+
+export const DEFAULT_WINDOW_DURATIONS = {
+  minute: MINUTE,
+  fiveMinute: 5 * MINUTE,
+  hourly: HOUR,
+  daily: DAY,
+  monthly: 30 * DAY,
+};
+
+/**
+ * Bounds percentage value to [0, 100].
+ * @param {number|unknown} value
+ * @returns {number}
+ */
+export function clampPercent(value) {
+  const num = Number(value);
+  return Number.isFinite(num) ? Math.max(0, Math.min(100, num)) : 0;
+}
+
+/**
+ * Estimates when a quota window reaches 100% at the average burn rate observed.
+ *
+ * @param {object|undefined} window
+ * @param {number} window.percentUsed Current percentage used (0..100)
+ * @param {string|number} window.resetsAt ISO timestamp or epoch ms when window resets
+ * @param {number} durationMs Window duration in milliseconds
+ * @param {number} [now] Current epoch timestamp in milliseconds
+ * @returns {{
+ *   etaMs: number,
+ *   resetsAtMs: number,
+ *   exhaustsBeforeReset: boolean,
+ *   burnRatePerSec: number
+ * } | null}
+ */
+export function estimateRunway(window, durationMs, now = Date.now()) {
+  if (!window || window.resetsAt == null || !durationMs || durationMs <= 0) return null;
+
+  const resetsAtMs = typeof window.resetsAt === 'number' ? window.resetsAt : Date.parse(window.resetsAt);
+  if (!Number.isFinite(resetsAtMs) || resetsAtMs <= now) return null;
+
+  const used = clampPercent(window.percentUsed);
+  const timeToReset = resetsAtMs - now;
+
+  // Already 100% or more: exhausts immediately
+  if (used >= 100) {
+    return {
+      etaMs: 0,
+      resetsAtMs,
+      exhaustsBeforeReset: true,
+      burnRatePerSec: 0,
+    };
+  }
+
+  const startedAt = resetsAtMs - durationMs;
+  const elapsed = Math.max(1, now - startedAt);
+  const ratePerMs = used / elapsed;
+
+  if (!Number.isFinite(ratePerMs) || ratePerMs <= 0) {
+    return {
+      etaMs: Number.POSITIVE_INFINITY,
+      resetsAtMs,
+      exhaustsBeforeReset: false,
+      burnRatePerSec: 0,
+    };
+  }
+
+  const remainingPercent = 100 - used;
+  const etaMs = remainingPercent / ratePerMs;
+  const exhaustsBeforeReset = etaMs < timeToReset;
+  const burnRatePerSec = ratePerMs * 1000;
+
+  return {
+    etaMs,
+    resetsAtMs,
+    exhaustsBeforeReset,
+    burnRatePerSec,
+  };
+}
+
+/**
+ * Stateful Quota Runway tracker for credentials and pools.
+ */
+export class QuotaRunway {
+  /**
+   * @param {object} [options]
+   * @param {number} [options.slidingWindowMs]
+   */
+  constructor(options = {}) {
+    this.slidingWindowMs = options.slidingWindowMs ?? (5 * MINUTE);
+    /** @type {Map<string, Array<{ at: number, amount: number }>>} */
+    this.samples = new Map();
+    /** @type {Map<string, { limit: number, remaining: number, resetsAtMs: number, durationMs: number }>} */
+    this.knownLimits = new Map();
+  }
+
+  /**
+   * Records usage tokens or requests for a key.
+   * @param {string} keyRef
+   * @param {number} amount
+   * @param {number} [now]
+   */
+  recordUsage(keyRef, amount, now = Date.now()) {
+    if (!keyRef || !amount || amount <= 0) return;
+    let list = this.samples.get(keyRef);
+    if (!list) {
+      list = [];
+      this.samples.set(keyRef, list);
+    }
+    list.push({ at: now, amount });
+    this._prune(keyRef, now);
+  }
+
+  /**
+   * Updates key limits detected from response headers or provider metadata.
+   * @param {string} keyRef
+   * @param {object} limits
+   * @param {number} [limits.limit]
+   * @param {number} [limits.remaining]
+   * @param {number|string} [limits.resetsAt]
+   * @param {number} [limits.durationMs]
+   * @param {number} [now]
+   */
+  updateLimits(keyRef, limits, now = Date.now()) {
+    if (!keyRef || !limits) return;
+    const limit = Number(limits.tpmLimit ?? limits.rpmLimit ?? limits.limit);
+    const remaining = Number(limits.remainingTokens ?? limits.remainingRequests ?? limits.remaining);
+    const resetSec = Number(limits.resetTokens ?? limits.resetRequests ?? limits.retryAfter);
+
+    let resetsAtMs = NaN;
+    if (Number.isFinite(resetSec)) {
+      resetsAtMs = now + (resetSec * 1000);
+    } else if (limits.resetsAt || limits.resetAt) {
+      const at = limits.resetsAt || limits.resetAt;
+      resetsAtMs = typeof at === 'number' ? at : Date.parse(at);
+    }
+
+    const durationMs = Number(limits.durationMs) || DEFAULT_WINDOW_DURATIONS.minute;
+
+    if (Number.isFinite(limit) && Number.isFinite(remaining) && Number.isFinite(resetsAtMs)) {
+      this.knownLimits.set(keyRef, {
+        limit,
+        remaining,
+        resetsAtMs,
+        durationMs,
+      });
+    }
+  }
+
+  /**
+   * Checks if key is projected to exhaust its quota before the current window resets.
+   * @param {string} keyRef
+   * @param {number} [now]
+   * @returns {boolean}
+   */
+  isExhaustingBeforeReset(keyRef, now = Date.now()) {
+    const assessment = this.assessRunway(keyRef, now);
+    return Boolean(assessment?.exhaustsBeforeReset);
+  }
+
+  /**
+   * Returns complete runway assessment for a key.
+   * @param {string} keyRef
+   * @param {number} [now]
+   * @returns {{
+   *   exhaustsBeforeReset: boolean,
+   *   etaMs: number,
+   *   burnRatePerSec: number,
+   *   timeToResetMs: number
+   * } | null}
+   */
+  assessRunway(keyRef, now = Date.now()) {
+    if (!keyRef) return null;
+
+    const limits = this.knownLimits.get(keyRef);
+    if (!limits || limits.resetsAtMs <= now) {
+      return null;
+    }
+
+    const { limit, remaining, resetsAtMs, durationMs } = limits;
+    if (limit <= 0) return null;
+
+    const used = Math.max(0, limit - remaining);
+    const percentUsed = (used / limit) * 100;
+
+    const estimate = estimateRunway(
+      { percentUsed, resetsAt: resetsAtMs },
+      durationMs,
+      now,
+    );
+
+    if (!estimate) return null;
+
+    return {
+      exhaustsBeforeReset: estimate.exhaustsBeforeReset,
+      etaMs: estimate.etaMs,
+      burnRatePerSec: estimate.burnRatePerSec,
+      timeToResetMs: resetsAtMs - now,
+    };
+  }
+
+  _prune(keyRef, now) {
+    const list = this.samples.get(keyRef);
+    if (!list) return;
+    const cutoff = now - this.slidingWindowMs;
+    while (list.length > 0 && list[0].at < cutoff) {
+      list.shift();
+    }
+  }
+
+  clear() {
+    this.samples.clear();
+    this.knownLimits.clear();
+  }
+}
+
+export const defaultQuotaRunway = new QuotaRunway();

package/lib/rotate.js

@@ -15,6 +15,7 @@ import {
   getSessionSticky,
   setSessionSticky,
   clearSessionSticky,
+  unstickFailedKey,
 } from './pool.js';
 import { pickCascadeFallback, CASCADE_MAX_DEPTH } from './cascade.js';
 import { createRequestBudget, BUDGET_EXHAUSTED_TIMEOUT, BUDGET_EXHAUSTED_ATTEMPTS } from './request-budget.js';
@@ -210,6 +211,7 @@ export function createRotate(deps) {
       const penalizeRef = (targetRef, errCode, errMsg, rawHeaders) => {
         if (!targetRef) return;
         if (sessionId) clearSessionSticky(pool, sessionId);
+        unstickFailedKey(pool, targetRef);
         if (errCode === LOCAL_QUOTA_CODE) return;
         let headerResetMs = null;
         if (rawHeaders) {
@@ -411,6 +413,7 @@ export function createRotate(deps) {
           let yielded = false;
           let switching = false;
           let inner;
+          const preTokenBuffer = [];
           try {
             // #464 & #461: Record attempt & trace dispatch
             budget.recordAttempt();
@@ -450,10 +453,13 @@ export function createRotate(deps) {
                     if (ok) reqStore.concurrencyRef = fallbackRef;
                   }
                 }
-                // 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')) {
-                  yielded = true;
+                // Pre-First-Token Failover (#527): Lock the stream only upon first content delta.
+                if (chunk && isContentChunk(chunk)) {
+                  if (!yielded) {
+                    yielded = true;
+                    for (const prev of preTokenBuffer) yield prev;
+                    preTokenBuffer.length = 0;
+                  }
                   yield chunk;
                   continue;
                 }
@@ -465,15 +471,16 @@ export function createRotate(deps) {
                   const message = failure?.message ?? '';
                   const effectiveSwitchCodes = pool.switchCodes ?? switchCodes;
                   const switchable = !yielded && kind === 'error' && isSwitchableError(failure, effectiveSwitchCodes);
-                  const activeRef = reqStore.pickedRef ?? pool.state.lastUsed;
+                  const activeRef = reqStore.pickedRef ?? pool.state.lastUsed ?? attemptList[attempt];
                   if (switchable) {
+                    preTokenBuffer.length = 0; // Discard metadata from aborted attempt (#527)
                     penalizeRef(activeRef, code ?? 'UNKNOWN', message, rawHdrs ?? failure?.headers ?? pool.state.detectedLimits?.get?.(activeRef));
                     recordFailureCategory(pool, code ?? 'UNKNOWN', message, activeRef, now());
                     pool.state.switches = (pool.state.switches ?? 0) + 1;
                     pool.state.lastReason = String(code ?? 'UNKNOWN');
                     pool.state.lastSwitchAt = now();
                     lastFailure = chunk;
-                    logWarn(`[dsh-key-rotation] ${options.provider}: key ${String(activeRef ?? '?')} failed (${String(code)} ${String(message).slice(0, 100)}) - next key`);
+                    logWarn(`[dsh-key-rotation] ${options.provider}: key ${String(activeRef ?? '?')} failed before first token (${String(code)} ${String(message).slice(0, 100)}) - failover to next key`);
                     // #216: per-switch webhook (opt-in switchNotify), deduped per provider
                     if (switchNotify && activeRef) {
                       notifySwitch(runtime0, pool, {
@@ -492,6 +499,10 @@ export function createRotate(deps) {
                     if (kind === 'error' && isSwitchableError(failure, effectiveSwitchCodes)) {
                       penalizeRef(activeRef, code ?? 'UNKNOWN', message, rawHdrs ?? failure?.headers ?? pool.state.detectedLimits?.get?.(activeRef));
                     }
+                    if (!yielded) {
+                      for (const prev of preTokenBuffer) yield prev;
+                      preTokenBuffer.length = 0;
+                    }
                     yield chunk;
                     return;
                   }
@@ -539,24 +550,33 @@ export function createRotate(deps) {
                     });
                   }
                   recordLatency(pool, reqStore);
-                  activePermit?.success();
+                  if (typeof activePermit?.success === 'function') activePermit.success();
                   if (isTraceOwner) trace.finish('SUCCESS', activeRef);
+                  if (!yielded) {
+                    for (const prev of preTokenBuffer) yield prev;
+                    preTokenBuffer.length = 0;
+                  }
                   yield chunk;
                   return;
                 }
-                yield chunk;
+                if (yielded) {
+                  yield chunk;
+                } else {
+                  preTokenBuffer.push(chunk);
+                }
               }
             } catch (e) {
               const effectiveSwitchCodes = pool.switchCodes ?? switchCodes;
-              const activeRef = _pickedRef ?? reqStore.pickedRef ?? pool.state.lastUsed;
+              const activeRef = _pickedRef ?? reqStore.pickedRef ?? pool.state.lastUsed ?? attemptList[attempt];
               if (!yielded && isSwitchableError(e, effectiveSwitchCodes)) {
+                preTokenBuffer.length = 0; // Discard metadata from failed attempt (#527)
                 penalizeRef(activeRef, e?.code ?? 'TRANSPORT', String(e?.message ?? e), e?.headers ?? e?.response?.headers ?? pool.state.detectedLimits?.get?.(activeRef));
                 recordFailureCategory(pool, e?.code ?? 'TRANSPORT', String(e?.message ?? e), activeRef, now());
                 pool.state.switches = (pool.state.switches ?? 0) + 1;
                 pool.state.lastReason = String(e?.code ?? 'TRANSPORT');
                 pool.state.lastSwitchAt = now();
                 lastFailure = finishError(e?.code ?? 'TRANSPORT', String(e?.message ?? e));
-                logWarn(`[dsh-key-rotation] ${options.provider}: key ${String(activeRef ?? '?')} stream threw ${String(e?.code ?? e?.message ?? e)} - failover to next key`);
+                logWarn(`[dsh-key-rotation] ${options.provider}: key ${String(activeRef ?? '?')} stream threw before first token ${String(e?.code ?? e?.message ?? e)} - failover to next key`);
                 if (switchNotify && activeRef) {
                   notifySwitch(runtime0, pool, {
                     provider: options.provider,
@@ -570,16 +590,24 @@ export function createRotate(deps) {
               if (yielded && isSwitchableError(e, effectiveSwitchCodes)) {
                 penalizeRef(activeRef, e?.code ?? 'TRANSPORT', String(e?.message ?? e), e?.headers ?? e?.response?.headers ?? pool.state.detectedLimits?.get?.(activeRef));
               }
+              if (!yielded) {
+                for (const prev of preTokenBuffer) yield prev;
+                preTokenBuffer.length = 0;
+              }
               yield finishError(e?.code ?? 'TRANSPORT', String(e?.message ?? e));
               return;
             }
 
             if (switching) continue; // try the next key
+            if (!yielded && preTokenBuffer.length > 0) {
+              for (const prev of preTokenBuffer) yield prev;
+              preTokenBuffer.length = 0;
+            }
             if (yielded) {
               const activeRef = reqStore.pickedRef ?? _pickedRef ?? pool.state.lastUsed;
               if (activeRef) recordSuccess(pool, activeRef, now());
               recordLatency(pool, reqStore);
-              activePermit?.success();
+              if (typeof activePermit?.success === 'function') activePermit.success();
               if (isTraceOwner) trace.finish('SUCCESS', activeRef);
             } else if (isTraceOwner) {
               trace.finish('COMPLETED', reqStore.pickedRef ?? _pickedRef ?? pool.state.lastUsed);
@@ -626,3 +654,25 @@ export function createRotate(deps) {
 
   return rotate;
 }
+
+/**
+ * Whether a chunk carries a model-visible content delta (#527).
+ * @param {Record<string, unknown>} chunk
+ * @returns {boolean}
+ */
+export function isContentChunk(chunk) {
+  if (!chunk || typeof chunk !== 'object') return false;
+  if (chunk.type === 'text-delta' || chunk.type === 'reasoning-delta' || chunk.type === 'tool-call-delta') {
+    return true;
+  }
+  if (chunk.delta !== undefined || chunk.text !== undefined || chunk.content !== undefined || chunk.choices !== undefined) {
+    return true;
+  }
+  if (chunk.type === 'start' || chunk.type === 'block-start' || chunk.type === 'block-stop' ||
+      chunk.type === 'metadata' || chunk.type === 'response-metadata' || chunk.type === 'usage' ||
+      chunk.type === 'finish') {
+    return false;
+  }
+  return false;
+}
+

package/lib/session-router.js

@@ -0,0 +1,265 @@
+// lib/session-router.js — Bounded SessionRouter with LRU eviction and proactive unsticking (#548).
+// Prevents unbounded memory growth and unpins sessions from failed/cooldown keys.
+
+export const MAX_SESSIONS = 1000;
+export const SESSION_TTL_MS = 24 * 60 * 60 * 1000; // 24 hours
+
+export class SessionRouter {
+  /**
+   * @param {{
+   *   maxSessions?: number,
+   *   sessionTtlMs?: number,
+   *   now?: () => number,
+   * }} [options]
+   */
+  constructor(options = {}) {
+    this.maxSessions = options.maxSessions ?? MAX_SESSIONS;
+    this.sessionTtlMs = options.sessionTtlMs ?? SESSION_TTL_MS;
+    this.now = options.now ?? (() => Date.now());
+    // sessionId -> { keyRef: string, assignedAt: number, lastUsedAt: number }
+    this.sessions = new Map();
+  }
+
+  /**
+   * Evict sessions older than TTL and prune LRU entries if capacity exceeded.
+   *
+   * @param {number} [currentTime]
+   * @returns {number} count of evicted sessions
+   */
+  prune(currentTime = this.now()) {
+    let evicted = 0;
+    // 1. Evict by TTL
+    for (const [id, entry] of this.sessions.entries()) {
+      if (currentTime - entry.lastUsedAt > this.sessionTtlMs) {
+        this.sessions.delete(id);
+        evicted++;
+      }
+    }
+
+    // 2. Evict least-recently used if size exceeds maxSessions
+    while (this.sessions.size > this.maxSessions) {
+      let oldestId = null;
+      let oldestTime = Infinity;
+      for (const [id, entry] of this.sessions.entries()) {
+        if (entry.lastUsedAt < oldestTime) {
+          oldestTime = entry.lastUsedAt;
+          oldestId = id;
+        }
+      }
+      if (oldestId !== null) {
+        this.sessions.delete(oldestId);
+        evicted++;
+      } else {
+        break;
+      }
+    }
+
+    return evicted;
+  }
+
+  /**
+   * Count how many active sessions are currently assigned to each key.
+   *
+   * @param {number} [currentTime]
+   * @returns {Map<string, number>}
+   */
+  getKeyAssignmentCounts(currentTime = this.now()) {
+    const counts = new Map();
+    for (const entry of this.sessions.values()) {
+      if (currentTime - entry.lastUsedAt <= this.sessionTtlMs) {
+        counts.set(entry.keyRef, (counts.get(entry.keyRef) ?? 0) + 1);
+      }
+    }
+    return counts;
+  }
+
+  /**
+   * Select candidate key with the least active sessions assigned (Least-Used Initial Assignment).
+   *
+   * @param {string[]} availableRefs List of candidate key refs
+   * @param {number} [currentTime]
+   * @returns {string|null}
+   */
+  selectLeastUsedKey(availableRefs, currentTime = this.now()) {
+    if (!Array.isArray(availableRefs) || availableRefs.length === 0) return null;
+    if (availableRefs.length === 1) return availableRefs[0];
+
+    const counts = this.getKeyAssignmentCounts(currentTime);
+    let minKey = availableRefs[0];
+    let minCount = counts.get(minKey) ?? 0;
+
+    for (let i = 1; i < availableRefs.length; i++) {
+      const key = availableRefs[i];
+      const count = counts.get(key) ?? 0;
+      if (count < minCount) {
+        minCount = count;
+        minKey = key;
+      }
+    }
+
+    return minKey;
+  }
+
+  /**
+   * Resolve pinned key for a session or allocate via Least-Used Initial Assignment.
+   *
+   * @param {string} sessionId
+   * @param {string[]} availableRefs Candidate keys currently healthy/available
+   * @param {{ preferKey?: string }} [options]
+   * @returns {{ keyRef: string, isReassigned: boolean, reason: string } | null}
+   */
+  resolveSessionKey(sessionId, availableRefs, options = {}) {
+    if (!sessionId || !Array.isArray(availableRefs) || availableRefs.length === 0) {
+      const fallback = availableRefs?.[0] ?? null;
+      return fallback ? { keyRef: fallback, isReassigned: false, reason: 'ad_hoc' } : null;
+    }
+
+    const t = this.now();
+    let existing = this.sessions.get(sessionId);
+
+    // Check TTL
+    if (existing && (t - existing.lastUsedAt > this.sessionTtlMs)) {
+      this.sessions.delete(sessionId);
+      existing = null;
+    }
+
+    // If session is already pinned to an available key:
+    if (existing && availableRefs.includes(existing.keyRef)) {
+      existing.lastUsedAt = t;
+      return { keyRef: existing.keyRef, isReassigned: false, reason: 'sticky_pinned' };
+    }
+
+    // Prune before adding new entry
+    if (!existing && this.sessions.size >= this.maxSessions) {
+      this.prune(t);
+    }
+
+    // Select candidate: prefer explicit key or least-used
+    let chosen = options.preferKey && availableRefs.includes(options.preferKey)
+      ? options.preferKey
+      : this.selectLeastUsedKey(availableRefs, t);
+
+    if (!chosen) chosen = availableRefs[0];
+
+    const wasReassigned = Boolean(existing);
+    this.sessions.set(sessionId, {
+      keyRef: chosen,
+      assignedAt: existing ? existing.assignedAt : t,
+      lastUsedAt: t,
+    });
+
+    return {
+      keyRef: chosen,
+      isReassigned: wasReassigned,
+      reason: wasReassigned ? 'mid_session_failover' : 'new_session_assigned',
+    };
+  }
+
+  /**
+   * Pin or update session sticky assignment explicitly.
+   *
+   * @param {string} sessionId
+   * @param {string} keyRef
+   */
+  setSession(sessionId, keyRef) {
+    if (!sessionId || !keyRef) return;
+    const t = this.now();
+    const existing = this.sessions.get(sessionId);
+    this.sessions.set(sessionId, {
+      keyRef,
+      assignedAt: existing ? existing.assignedAt : t,
+      lastUsedAt: t,
+    });
+    if (this.sessions.size > this.maxSessions) {
+      this.prune(t);
+    }
+  }
+
+  /**
+   * Get pinned key for a session if still valid.
+   *
+   * @param {string} sessionId
+   * @returns {string|null}
+   */
+  getSession(sessionId) {
+    if (!sessionId) return null;
+    const entry = this.sessions.get(sessionId);
+    if (!entry) return null;
+    const t = this.now();
+    if (t - entry.lastUsedAt > this.sessionTtlMs) {
+      this.sessions.delete(sessionId);
+      return null;
+    }
+    return entry.keyRef;
+  }
+
+  /**
+   * Unstick a specific session (e.g. after a failure or manual unpin).
+   *
+   * @param {string} sessionId
+   * @returns {boolean} true if session was present and removed
+   */
+  unstick(sessionId) {
+    if (!sessionId) return false;
+    return this.sessions.delete(sessionId);
+  }
+
+  /**
+   * Proactively unstick all sessions pinned to a failed / cooling key (#548).
+   *
+   * @param {string} keyRef Failed key ref
+   * @returns {number} count of sessions unpinned
+   */
+  unstickKey(keyRef) {
+    if (!keyRef) return 0;
+    let count = 0;
+    for (const [id, entry] of this.sessions.entries()) {
+      if (entry.keyRef === keyRef) {
+        this.sessions.delete(id);
+        count++;
+      }
+    }
+    return count;
+  }
+
+  /**
+   * Return router diagnostics.
+   */
+  getStatus() {
+    const t = this.now();
+    const counts = this.getKeyAssignmentCounts(t);
+    return {
+      activeSessionsCount: this.sessions.size,
+      maxSessions: this.maxSessions,
+      sessionTtlMs: this.sessionTtlMs,
+      keyAssignments: Object.fromEntries(counts.entries()),
+    };
+  }
+
+  /**
+   * Reset all sessions (useful for tests).
+   */
+  reset() {
+    this.sessions.clear();
+  }
+}
+
+// Global default router instance for convenience
+export const defaultSessionRouter = new SessionRouter();
+
+/**
+ * Helper to unstick all sessions pinned to keyRef across router or default instance (#548).
+ * @param {string|object} keyRefOrRouter
+ * @param {string} [keyRef]
+ * @returns {number}
+ */
+export function unstickFailedKey(keyRefOrRouter, keyRef) {
+  if (typeof keyRefOrRouter === 'string') {
+    return defaultSessionRouter.unstickKey(keyRefOrRouter);
+  }
+  if (keyRefOrRouter instanceof SessionRouter) {
+    return keyRefOrRouter.unstickKey(keyRef);
+  }
+  return defaultSessionRouter.unstickKey(keyRef);
+}
+

package/lib/stream-metrics.js

@@ -1,4 +1,4 @@
-// lib/stream-metrics.js — streaming success telemetry, usage, cost and token recording (#512).
+// lib/stream-metrics.js — streaming success telemetry, usage, cost and token recording (#512, #546, #549).
 // Extracted from lib/rotate.js to enforce < 600 line modularity standard.
 
 import { recordSuccess, setSessionSticky, recordFailure, extractRateLimit, isRateLimited } from './pool.js';
@@ -7,6 +7,7 @@ import { recordRequestOutcome } from './ops-analytics.js';
 import { tpmRecord } from './bucket.js';
 import { extractUsageTokens, consumeModelTokens } from './model-quota.js';
 import { extractProviderLimits } from './provider-limits-detector.js';
+import { normalizeUsageTokens } from './transports/tokens.js';
 
 export async function* preserveDispatchContext(iterable, storage, store) {
   if (!storage || typeof storage.getStore !== 'function') {
@@ -68,11 +69,15 @@ export function recordStreamSuccessMetrics(pool, activeRef, options, chunk, deps
     cascaded: Boolean(options.__isCascade),
   }, wallNow);
 
-  // 3. TPM sliding window bucket recording
-  const usedTokens = extractUsageTokens(usagePayload);
+  // 3. TPM sliding window bucket recording & token normalization (#546, #549)
+  const normUsage = normalizeUsageTokens(usagePayload);
+  const usedTokens = normUsage.totalTokens > 0 ? normUsage.totalTokens : extractUsageTokens(usagePayload);
   if (usedTokens && usedTokens > 0) {
     if (!pool.state.tpmWindows) pool.state.tpmWindows = new Map();
     tpmRecord(pool.state.tpmWindows, activeRef, usedTokens, wallNow);
+    if (pool.state.quotaRunway && typeof pool.state.quotaRunway.recordUsage === 'function') {
+      pool.state.quotaRunway.recordUsage(activeRef, usedTokens, wallNow);
+    }
   }
 
   // 4. Local model quota consumption
@@ -91,12 +96,15 @@ export function recordStreamSuccessMetrics(pool, activeRef, options, chunk, deps
     }
   }
 
-  // 5. Provider rate limit headers detection (#503)
+  // 5. Provider rate limit headers detection & predictive runway (#503, #549)
   if (rawHdrs) {
     const detected = extractProviderLimits(rawHdrs);
     if (detected) {
       if (!pool.state.detectedLimits) pool.state.detectedLimits = new Map();
       pool.state.detectedLimits.set(activeRef, { ...detected, at: nowMs });
+      if (pool.state.quotaRunway && typeof pool.state.quotaRunway.updateLimits === 'function') {
+        pool.state.quotaRunway.updateLimits(activeRef, detected, wallNow);
+      }
     }
   }
 
@@ -130,4 +138,4 @@ export function recordStreamSuccessMetrics(pool, activeRef, options, chunk, deps
   if (rate && Number.isFinite(rate.remaining) && quotaStore) {
     quotaStore.set(activeRef, { remaining: rate.remaining, limit: rate.limit, reset: rate.reset, at: nowMs });
   }
-}
+}