dsh-plugin-subscription-oauth 0.6.1 → 0.8.0

package/lib/providers/responses-stream.js

@@ -12,8 +12,49 @@
  * @module dsh-plugin-subscription-oauth/providers/responses-stream
  */
 import { EMPTY_RESPONSE_CODE, LlmError } from '@deepseek-ai/dsh-llm';
+import { allowanceScheduler } from '../auth/allowance-scheduler.js';
 import { streamResponses } from '../translate/responses.js';
 import { httpLlmError, idleWatchdog, mapFetchFailure } from './common.js';
+/**
+ * Build the cross-account hooks one route needs to draw on its whole account pool.
+ *
+ * Every Responses-wire route (Kimi, Grok, GLM, Ark) meters per account and can run out;
+ * this wires the three account hooks identically, so a route only supplies how to read
+ * one account's allowance and its provider id. A route with no `readUsage` still rotates
+ * accounts on failure, it just cannot predict exhaustion before the first failure.
+ *
+ * @param provider - the provider route id, the scheduler's key.
+ * @param tokens - the route's token manager, enumerates accounts.
+ * @param readUsage - read one account's metered windows, when wired.
+ */
+export function accountPoolHooks(input) {
+    const { provider, tokens, readUsage } = input;
+    return {
+        accounts: () => tokens.accounts(),
+        // Steer away from an account that has no allowance: read its live figures when a
+        // reader is wired (so a 100 % window is avoided before it fails there), else fall
+        // back to the remembered reset time from a past failure.
+        isAccountSpent: async (account) => {
+            if (readUsage !== undefined) {
+                try {
+                    const windows = await readUsage(account);
+                    const spentWindow = windows.find((window) => (window.usedPercent ?? 0) >= 100);
+                    if (spentWindow !== undefined) {
+                        allowanceScheduler.noteReset(provider, account, spentWindow.resetsAt);
+                        return true;
+                    }
+                    return false;
+                }
+                catch {
+                    // A failed read is not a reason to avoid the account; fall through to memory.
+                }
+            }
+            return allowanceScheduler.isSpentUntilReset(provider, account);
+        },
+        noteAccountSpent: (account, resetsAt) => { allowanceScheduler.noteReset(provider, account, resetsAt); },
+        resetsAt: async (account) => allowanceScheduler.resetAt(provider, account),
+    };
+}
 /**
  * Whether a refusal says "this model cannot serve you now", as opposed to "this
  * request is wrong".
@@ -60,20 +101,92 @@ export async function* streamProviderResponses(options, input) {
         if (models.length === 0) {
             throw new LlmError(`${label} has no enabled model to dispatch to`, EMPTY_RESPONSE_CODE);
         }
-        let token = await tokens.resolve();
-        for (const [index, model] of models.entries()) {
-            try {
-                let response = await request(options, model, token.accessToken, watchdog.signal);
-                if (response.status === 401 || response.status === 403) {
-                    token = await tokens.resolve(token.account, true);
-                    response = await request(options, model, token.accessToken, watchdog.signal);
+        // When the provider meters per account and an account's allowance can run out,
+        // every account is a candidate: a spent one falls through to the next, so one
+        // model id draws on the whole pool's allowance. Accounts are tried in order;
+        // the first that serves the request wins.
+        let accounts = input.accounts === undefined
+            ? undefined
+            : [...await input.accounts()];
+        // Send work straight to an account that has allowance: one already known to be
+        // spent until a future reset is moved to the back, so a fresh request does not
+        // fail against it once to find out. This is a preference, not a skip — a lone
+        // account still serves, and the spent-retry below catches a stale record.
+        if (accounts !== undefined && accounts.length > 1 && input.isAccountSpent !== undefined) {
+            const fresh = [];
+            const tired = [];
+            for (const account of accounts) {
+                if (await input.isAccountSpent(account))
+                    tired.push(account);
+                else
+                    fresh.push(account);
+            }
+            if (tired.length > 0 && fresh.length > 0) {
+                accounts = [...fresh, ...tired];
+            }
+        }
+        const tryAccounts = accounts === undefined || accounts.length === 0
+            ? [undefined]
+            : accounts;
+        for (const accountKey of tryAccounts) {
+            let token = await tokens.resolve(accountKey);
+            // Track which accounts a spent-allowance already moved past, so a later model
+            // in the chain does not re-try an account known to be exhausted.
+            const spent = new Set();
+            for (const [index, model] of models.entries()) {
+                try {
+                    if (spent.has(token.account))
+                        continue;
+                    let response = await request(options, model, token.accessToken, watchdog.signal);
                     if (response.status === 401 || response.status === 403) {
-                        input.onCredentialRejected?.(token.account);
+                        token = await tokens.resolve(token.account, true);
+                        response = await request(options, model, token.accessToken, watchdog.signal);
+                        if (response.status === 401 || response.status === 403) {
+                            input.onCredentialRejected?.(token.account);
+                        }
+                    }
+                    // A spent allowance moves to the next account (when the provider supports
+                    // it), not to the next model: the model is fine, the account is not.
+                    if (!response.ok
+                        && input.isAllowanceExhausted !== undefined
+                        && input.accounts !== undefined
+                        && input.isAllowanceExhausted(response.status)
+                        && tryAccounts.length > 1) {
+                        spent.add(token.account);
+                        // Remember when this account's allowance refills, so later requests go
+                        // straight to another account until then instead of failing here again.
+                        input.noteAccountSpent?.(token.account, await input.resetsAt?.(token.account));
+                        const next = tryAccounts.find((candidate) => candidate !== undefined && !spent.has(candidate));
+                        if (next !== undefined) {
+                            lastError = await httpLlmError(response, label);
+                            input.onAccountSwitch?.(token.account, next, reasonOf(lastError));
+                            token = await tokens.resolve(next);
+                            // Re-try this same model on the fresh account.
+                            response = await request(options, model, token.accessToken, watchdog.signal);
+                        }
+                    }
+                    if (!response.ok) {
+                        const error = await httpLlmError(response, label);
+                        if (models.length > 1 && isModelFallbackWorthy(response.status) && index < models.length - 1) {
+                            const next = models[index + 1];
+                            if (next !== undefined) {
+                                lastError = error;
+                                input.onFallback?.(model, next, reasonOf(error));
+                                continue;
+                            }
+                        }
+                        throw error;
                     }
+                    if (response.body === null) {
+                        throw new LlmError(`${label} returned no response body`, EMPTY_RESPONSE_CODE);
+                    }
+                    yield* streamResponses(response.body, () => { watchdog.pulse(); });
+                    return;
                 }
-                if (!response.ok) {
-                    const error = await httpLlmError(response, label);
-                    if (models.length > 1 && isModelFallbackWorthy(response.status) && index < models.length - 1) {
+                catch (error) {
+                    // A transport failure on a chain is also worth another model.
+                    const transport = error instanceof TypeError || error.name === 'AbortError';
+                    if (transport && options.signal?.aborted !== true && index < models.length - 1) {
                         const next = models[index + 1];
                         if (next !== undefined) {
                             lastError = error;
@@ -81,26 +194,8 @@ export async function* streamProviderResponses(options, input) {
                             continue;
                         }
                     }
-                    throw error;
-                }
-                if (response.body === null) {
-                    throw new LlmError(`${label} returned no response body`, EMPTY_RESPONSE_CODE);
-                }
-                yield* streamResponses(response.body, () => { watchdog.pulse(); });
-                return;
-            }
-            catch (error) {
-                // A transport failure on a chain is also worth another model.
-                const transport = error instanceof TypeError || error.name === 'AbortError';
-                if (transport && options.signal?.aborted !== true && index < models.length - 1) {
-                    const next = models[index + 1];
-                    if (next !== undefined) {
-                        lastError = error;
-                        input.onFallback?.(model, next, reasonOf(error));
-                        continue;
-                    }
+                    throw input.describeFailure === undefined ? error : input.describeFailure(error);
                 }
-                throw input.describeFailure === undefined ? error : input.describeFailure(error);
             }
         }
         throw lastError ?? new LlmError(`${label} served none of the requested models`, EMPTY_RESPONSE_CODE);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "dsh-plugin-subscription-oauth",
-  "version": "0.6.1",
+  "version": "0.8.0",
   "description": "Use subscription accounts as DeepSeek Harness LLM providers — Kimi Code (OAuth device flow), Grok / X Premium (OIDC + PKCE), Google Antigravity (Google OAuth + Cloud Code Assist) and GLM / BigModel (ZCode's browser sign-in) — with per-model visibility, drag-to-reorder Auto fallback and plan quotas from the web Settings page",
   "license": "MIT",
   "author": "EricFL1998",