@jsonstudio/llms 0.6.1354 → 0.6.1399

package/dist/router/virtual-router/engine-selection/alias-selection.js

@@ -1,8 +1,9 @@
+const COOLDOWN_EMPTY_THRESHOLD_MS = 30_000;
 // Default provider-level strategy table.
 // This is a data-only default; callers can override via `loadBalancing.aliasSelection.providers`.
 export const DEFAULT_PROVIDER_ALIAS_SELECTION = {
-    // Antigravity: upstream gateway may reject rapid cross-key switching; stick to one alias until error.
-    antigravity: 'sticky-queue'
+    // Antigravity: prefer the alias with highest remaining quota; fall back to sticky-queue when quota is unknown.
+    antigravity: 'best-quota'
 };
 export function resolveAliasSelectionStrategy(providerId, cfg) {
     if (!providerId)
@@ -11,11 +12,11 @@ export function resolveAliasSelectionStrategy(providerId, cfg) {
         return 'none';
     const overrides = cfg?.providers ?? {};
     const override = overrides[providerId];
-    if (override === 'none' || override === 'sticky-queue') {
+    if (override === 'none' || override === 'sticky-queue' || override === 'best-quota') {
         return override;
     }
     const def = cfg?.defaultStrategy;
-    if (def === 'none' || def === 'sticky-queue') {
+    if (def === 'none' || def === 'sticky-queue' || def === 'best-quota') {
         return def;
     }
     const table = DEFAULT_PROVIDER_ALIAS_SELECTION[providerId];
@@ -100,6 +101,86 @@ export function pinCandidatesByAliasQueue(opts) {
     const selectedSet = new Set(selectedKeys);
     return candidates.filter((key) => selectedSet.has(key));
 }
+export function pinCandidatesByBestQuota(opts) {
+    const { providerId, modelId, candidates, orderedTargets, aliasOfKey, modelIdOfKey, quotaView, now } = opts;
+    if (!quotaView)
+        return null;
+    if (!providerId || !modelId)
+        return null;
+    if (!Array.isArray(candidates) || candidates.length < 2)
+        return null;
+    const aliasBuckets = new Map();
+    const aliasOrder = new Map();
+    let order = 0;
+    for (const key of candidates) {
+        if (!key || typeof key !== 'string')
+            continue;
+        if (!key.startsWith(`${providerId}.`))
+            return null;
+        const m = modelIdOfKey(key);
+        if (!m || m !== modelId)
+            return null;
+        const alias = aliasOfKey(key);
+        if (!alias)
+            return null;
+        const list = aliasBuckets.get(alias) ?? [];
+        list.push(key);
+        aliasBuckets.set(alias, list);
+        if (!aliasOrder.has(alias)) {
+            aliasOrder.set(alias, order++);
+        }
+    }
+    if (aliasBuckets.size <= 1)
+        return null;
+    const preferredOrder = resolveAliasOrderFromTargets({
+        orderedTargets,
+        providerId,
+        modelId,
+        aliasOfKey,
+        modelIdOfKey,
+        allowedAliases: new Set(aliasBuckets.keys())
+    });
+    for (const alias of preferredOrder) {
+        if (!aliasOrder.has(alias)) {
+            aliasOrder.set(alias, order++);
+        }
+    }
+    const eligible = [];
+    for (const [alias, keys] of aliasBuckets.entries()) {
+        const entry = quotaView(keys[0] ?? '');
+        if (!entry || entry.inPool === false) {
+            continue;
+        }
+        if (entry.blacklistUntil && entry.blacklistUntil > now) {
+            continue;
+        }
+        if (entry.cooldownUntil && entry.cooldownUntil - now >= COOLDOWN_EMPTY_THRESHOLD_MS) {
+            continue;
+        }
+        const remainingRaw = entry.remainingFraction;
+        const remaining = typeof remainingRaw === 'number' && Number.isFinite(remainingRaw) ? remainingRaw : 0;
+        if (remaining <= 0) {
+            continue;
+        }
+        eligible.push({
+            alias,
+            score: remaining,
+            order: aliasOrder.get(alias) ?? Number.MAX_SAFE_INTEGER
+        });
+    }
+    if (!eligible.length) {
+        return null;
+    }
+    eligible.sort((a, b) => (b.score - a.score) || (a.order - b.order));
+    const selectedAlias = eligible[0]?.alias;
+    if (!selectedAlias)
+        return null;
+    const selectedKeys = aliasBuckets.get(selectedAlias) ?? [];
+    if (!selectedKeys.length)
+        return null;
+    const selectedSet = new Set(selectedKeys);
+    return candidates.filter((key) => selectedSet.has(key));
+}
 function resolveAliasOrderFromTargets(opts) {
     const { orderedTargets, providerId, modelId, aliasOfKey, modelIdOfKey, allowedAliases } = opts;
     if (!Array.isArray(orderedTargets) || orderedTargets.length === 0) {

package/dist/router/virtual-router/engine-selection/tier-selection-select.js

@@ -1,11 +1,12 @@
 import { computeContextMultiplier } from '../context-weighted.js';
 import { computeHealthWeight } from '../health-weighted.js';
-import { pinCandidatesByAliasQueue, resolveAliasSelectionStrategy } from './alias-selection.js';
+import { pinCandidatesByAliasQueue, pinCandidatesByBestQuota, resolveAliasSelectionStrategy } from './alias-selection.js';
 import { computeContextWeightMultipliers } from './context-weight-multipliers.js';
 import { extractKeyAlias, extractProviderId, getProviderModelId } from './key-parsing.js';
 import { pickPriorityGroup } from './tier-priority.js';
+const ANTIGRAVITY_COOLDOWN_ALIAS_THRESHOLD_MS = 30_000;
 function applyAliasStickyQueuePinning(opts) {
-    const { candidates, orderedTargets, deps, excludedKeys } = opts;
+    const { candidates, orderedTargets, deps, excludedKeys, now } = opts;
     if (!Array.isArray(candidates) || candidates.length < 2) {
         return candidates;
     }
@@ -26,7 +27,7 @@ function applyAliasStickyQueuePinning(opts) {
             continue;
         }
         const strategy = resolveAliasSelectionStrategy(providerId, deps.loadBalancer.getPolicy().aliasSelection);
-        if (strategy !== 'sticky-queue') {
+        if (strategy !== 'sticky-queue' && strategy !== 'best-quota') {
             continue;
         }
         const modelId = getProviderModelId(key, deps.providerRegistry) ?? '';
@@ -58,17 +59,33 @@ function applyAliasStickyQueuePinning(opts) {
         if (aliases.size < 2) {
             continue;
         }
-        const pinned = pinCandidatesByAliasQueue({
-            queueStore: store,
-            providerId: group.providerId,
-            modelId: group.modelId,
-            candidates: group.keys,
-            orderedTargets,
-            excludedProviderKeys: excludedKeys,
-            aliasOfKey: extractKeyAlias,
-            modelIdOfKey: (key) => getProviderModelId(key, deps.providerRegistry),
-            availabilityCheck: (key) => deps.healthManager.isAvailable(key)
-        });
+        const strategy = resolveAliasSelectionStrategy(group.providerId, deps.loadBalancer.getPolicy().aliasSelection);
+        let pinned = null;
+        if (strategy === 'best-quota') {
+            pinned = pinCandidatesByBestQuota({
+                providerId: group.providerId,
+                modelId: group.modelId,
+                candidates: group.keys,
+                orderedTargets,
+                aliasOfKey: extractKeyAlias,
+                modelIdOfKey: (key) => getProviderModelId(key, deps.providerRegistry),
+                quotaView: deps.quotaView,
+                now
+            });
+        }
+        if (!pinned) {
+            pinned = pinCandidatesByAliasQueue({
+                queueStore: store,
+                providerId: group.providerId,
+                modelId: group.modelId,
+                candidates: group.keys,
+                orderedTargets,
+                excludedProviderKeys: excludedKeys,
+                aliasOfKey: extractKeyAlias,
+                modelIdOfKey: (key) => getProviderModelId(key, deps.providerRegistry),
+                availabilityCheck: (key) => deps.healthManager.isAvailable(key)
+            });
+        }
         if (pinned && pinned.length) {
             pinnedByGroup.set(groupId, new Set(pinned));
         }
@@ -161,7 +178,8 @@ export function selectProviderKeyFromCandidatePool(opts) {
             candidates: retryPreferredCandidates,
             orderedTargets: tier.targets,
             deps,
-            excludedKeys
+            excludedKeys,
+            now
         });
         if (tier.mode === 'priority') {
             if (isRecoveryAttempt) {
@@ -240,7 +258,11 @@ export function selectProviderKeyFromCandidatePool(opts) {
             continue;
         }
         if (entry.cooldownUntil && entry.cooldownUntil > now) {
-            continue;
+            const providerId = extractProviderId(key) ?? '';
+            const cooldownMs = entry.cooldownUntil - now;
+            if (providerId !== 'antigravity' || cooldownMs >= ANTIGRAVITY_COOLDOWN_ALIAS_THRESHOLD_MS) {
+                continue;
+            }
         }
         if (entry.blacklistUntil && entry.blacklistUntil > now) {
             continue;
@@ -273,7 +295,8 @@ export function selectProviderKeyFromCandidatePool(opts) {
             candidates: bucketCandidates,
             orderedTargets: tier.targets,
             deps,
-            excludedKeys
+            excludedKeys,
+            now
         });
         const bucketPenaltyMap = {};
         for (const item of bucket) {

package/dist/router/virtual-router/engine-selection/tier-selection.js

@@ -123,7 +123,7 @@ export function trySelectFromTier(routeName, tier, stickyKey, estimatedTokens, f
         return { providerKey: null, poolTargets: [], tierId: tier.id, failureHint: `${routeName}:${tier.id}:empty` };
     }
     const contextResult = deps.contextAdvisor.classify(targets, estimatedTokens, (key) => deps.providerRegistry.get(key));
-    const prioritizedPools = buildContextCandidatePools(contextResult);
+    const prioritizedPools = buildContextCandidatePools(contextResult, routeName);
     const quotaView = deps.quotaView;
     const now = quotaView ? Date.now() : 0;
     const healthWeightedCfg = resolveHealthWeightedConfig(deps.loadBalancer.getPolicy().healthWeighted);
@@ -200,7 +200,7 @@ function recordAliasQueueFailuresFromExcludedKeys(excludedKeys, orderedTargets,
         }
     }
 }
-function buildContextCandidatePools(result) {
+function buildContextCandidatePools(result, routeName) {
     const ordered = [];
     if (result.safe.length) {
         ordered.push(result.safe);
@@ -208,6 +208,9 @@ function buildContextCandidatePools(result) {
     if (result.risky.length) {
         ordered.push(result.risky);
     }
+    if (routeName === 'longcontext' && result.overflow.length) {
+        ordered.push(result.overflow);
+    }
     return ordered;
 }
 function describeAttempt(routeName, poolId, result) {

package/dist/router/virtual-router/engine.js

@@ -12,6 +12,7 @@ import { buildHitReason, formatVirtualRouterHit } from './engine-logging.js';
 import { selectDirectProviderModel, selectFromStickyPool, selectProviderImpl } from './engine-selection.js';
 import { applyQuotaDepletedImpl, applyQuotaRecoveryImpl, handleProviderFailureImpl, mapProviderErrorImpl } from './engine-health.js';
 import { mergeStopMessageFromPersisted } from './stop-message-state-sync.js';
+const ANTIGRAVITY_COOLDOWN_ALIAS_THRESHOLD_MS = 30_000;
 export class VirtualRouterEngine {
     routing = {};
     providerRegistry = new ProviderRegistry();
@@ -1468,10 +1469,17 @@ export class VirtualRouterEngine {
         if (!expiry) {
             return false;
         }
-        if (Date.now() >= expiry) {
+        const now = Date.now();
+        if (now >= expiry) {
             this.providerCooldowns.delete(providerKey);
             return false;
         }
+        if (providerKey.startsWith('antigravity.')) {
+            const remaining = expiry - now;
+            if (remaining < ANTIGRAVITY_COOLDOWN_ALIAS_THRESHOLD_MS) {
+                return false;
+            }
+        }
         return true;
     }
     restoreHealthFromStore() {

package/dist/router/virtual-router/types.d.ts

@@ -152,7 +152,7 @@ export interface HealthWeightedLoadBalancingConfig {
      */
     recoverToBestOnRetry?: boolean;
 }
-export type AliasSelectionStrategy = 'none' | 'sticky-queue';
+export type AliasSelectionStrategy = 'none' | 'sticky-queue' | 'best-quota';
 export interface AliasSelectionConfig {
     /**
      * Global on/off switch. When false, no alias-level selection is applied.
@@ -514,6 +514,19 @@ export interface ProviderQuotaViewEntry {
     inPool: boolean;
     reason?: string;
     priorityTier?: number;
+    /**
+     * Optional remaining quota fraction for the provider key (0..1).
+     * Used by alias-selection strategies that prefer higher remaining quota.
+     */
+    remainingFraction?: number | null;
+    /**
+     * Optional quota reset timestamp (ms since epoch) for the provider key.
+     */
+    quotaResetAtMs?: number | null;
+    /**
+     * Optional quota fetch timestamp (ms since epoch) for the provider key.
+     */
+    quotaFetchedAtMs?: number | null;
     /**
      * Optional soft penalty hint for selection ordering.
      * - 0 / undefined means no penalty

package/dist/servertool/engine.js

@@ -93,10 +93,8 @@ function coerceFollowupPayloadStream(payload, stream) {
     if (!payload || typeof payload !== 'object') {
         return payload;
     }
-    // ServerTool followup requests must be non-streaming to keep parsing deterministic and avoid
-    // provider-side SSE wrappers leaking into internal reenter calls.
-    if (stream === false) {
-        payload.stream = false;
+    if (typeof stream === 'boolean') {
+        payload.stream = stream;
     }
     return payload;
 }
@@ -440,7 +438,6 @@ export async function runServerToolOrchestration(options) {
         };
     }
     const metadata = {
-        stream: false,
         ...(engineResult.execution.followup.metadata ?? {})
     };
     const rt = ensureRuntimeMetadata(metadata);
@@ -476,7 +473,10 @@ export async function runServerToolOrchestration(options) {
     const retryEmptyFollowupOnce = isStopMessageFlow || isGeminiEmptyReplyContinue;
     const maxAttempts = retryEmptyFollowupOnce ? 2 : 1;
     const followupRequestId = buildFollowupRequestId(options.requestId, engineResult.execution.followup.requestIdSuffix);
-    let followupPayload = coerceFollowupPayloadStream(followupSeedPayload, metadata.stream === true);
+    const followupStream = typeof metadata.stream === 'boolean'
+        ? (metadata.stream)
+        : undefined;
+    let followupPayload = coerceFollowupPayloadStream(followupSeedPayload, followupStream);
     if (isGeminiEmptyReplyContinue) {
         // For gemini_empty_reply_continue, the goal is to recover text output from an empty/malformed reply.
         // Force the followup to be non-tool-calling to avoid repeated MALFORMED_FUNCTION_CALL loops.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jsonstudio/llms",
-  "version": "0.6.1354",
+  "version": "0.6.1399",
   "type": "module",
   "main": "dist/index.js",
   "module": "dist/index.js",