@goodandready/dsh-key-rotation 0.7.10 → 0.7.11

package/lib/index.js

@@ -46,6 +46,8 @@ const CONFIG_PATH = '/dsh-key-rotation/config';
 const STATUS_PATH = '/dsh-key-rotation/status';
 const KEY_PATH = '/dsh-key-rotation/key';
 const RESET_PATH = '/dsh-key-rotation/reset';
+const IMPORT_PATH = '/dsh-key-rotation/import';
+const HEALTH_PATH = '/dsh-key-rotation/health';
 const TEST_PATH = '/dsh-key-rotation/test';
 
 /** The llm-pi-ai namespace whose provider profiles map providers to pools. */
@@ -81,10 +83,17 @@ export const Config = Schema.object({
   maxCooldownMs: Schema.number(),
   notifyWebhook: Schema.string().default(''),
   notifyThreshold: Schema.number().default(3),
+  backupDir: Schema.string().default(''),
+  backupIntervalMs: Schema.number().default(86400000),
+  backupKeep: Schema.number().default(7),
   providers: Schema.array(Schema.object({
     provider: Schema.string().required(),
     keys: Schema.array(Schema.string()).default([]),
     weights: Schema.array(Schema.number()).default([]),
+    models: Schema.dict(Schema.object({
+      keys: Schema.array(Schema.string()).default([]),
+      weights: Schema.array(Schema.number()).default([]),
+    })).default({}),
     cooldownMs: Schema.number(),
     maxCooldownMs: Schema.number(),
   })).default([...DEFAULT_PROVIDERS]),
@@ -247,6 +256,68 @@ export function apply(ctx, config = {}) {
   // ── key-pool state, persisted across config reloads ──
   // base provider -> { failedUntil: Map<ref, epochMs>, pointer: number, lastUsed: ref }
   const poolState = new Map();
+  // Periodic backup of pools config
+  ctx.effect(() => {
+    const { backupDir, backupIntervalMs, backupKeep } = buildRuntime();
+    if (!backupDir) return;
+    const id = setInterval(() => {
+      try {
+        const fs = require('node:fs');
+        const path = require('node:path');
+        const dir = backupDir;
+        fs.mkdirSync(dir, { recursive: true });
+        const now = new Date();
+        const dateStr = now.toISOString().slice(0,10).replace(/-/g,'');
+        const file = path.join(dir, 'pools-' + dateStr + '.json');
+        const data = JSON.stringify({ backup: now.toISOString(), providers: getConfig()?.providers ?? [] }, null, 2);
+        fs.writeFileSync(file, data, 'utf8');
+        // prune old backups
+        const keep = backupKeep || 7;
+        const files = fs.readdirSync(dir).filter((f) => f.startsWith('pools-') && f.endsWith('.json')).sort();
+        while (files.length > keep) {
+          const old = files.shift();
+          fs.unlinkSync(path.join(dir, old));
+        }
+      } catch (e) {
+        console.warn('[dsh-key-rotation] backup failed:', String(e?.message ?? e));
+      }
+    }, backupIntervalMs || 86400000);
+    return () => clearInterval(id);
+  }, 'dsh-key-rotation: backup pools');
+  // Periodic save of usage/cost stats to file
+  ctx.effect(() => {
+    const { backupDir } = buildRuntime();
+    if (!backupDir) return;
+    try {
+      const fs = require('node:fs');
+      const path = require('node:path');
+      const statsFile = path.join(backupDir, 'stats.json');
+      // Load existing stats at startup
+      try {
+        if (fs.existsSync(statsFile)) {
+          const saved = JSON.parse(fs.readFileSync(statsFile, 'utf8'));
+          for (const st of poolState.values()) {
+            if (saved.usageCounts && st.usageCounts) { for (const [k, v] of Object.entries(saved.usageCounts)) st.usageCounts.set(k, (st.usageCounts.get(k) ?? 0) + v); }
+            if (saved.costPerKey && st.costPerKey) { for (const [k, v] of Object.entries(saved.costPerKey)) st.costPerKey.set(k, (st.costPerKey.get(k) ?? 0) + v); }
+            if (saved.lastUsedAt && st.lastUsedAt) { for (const [k, v] of Object.entries(saved.lastUsedAt)) { if (!st.lastUsedAt.has(k) || v > st.lastUsedAt.get(k)) st.lastUsedAt.set(k, v); } }
+          }
+        }
+      } catch {}
+      // Periodic save
+      const id = setInterval(() => {
+        try {
+          const usageCounts = {}; const costPerKey = {}; const lastUsedAt = {};
+          for (const [base, st] of poolState) {
+            if (st.usageCounts) for (const [k, v] of st.usageCounts) usageCounts[k] = v;
+            if (st.costPerKey) for (const [k, v] of st.costPerKey) costPerKey[k] = v;
+            if (st.lastUsedAt) for (const [k, v] of st.lastUsedAt) lastUsedAt[k] = v;
+          }
+          fs.writeFileSync(statsFile, JSON.stringify({ t: Date.now(), usageCounts, costPerKey, lastUsedAt }), 'utf8');
+        } catch {}
+      }, 60000);
+      return () => clearInterval(id);
+    } catch { return () => {}; }
+  }, 'dsh-key-rotation: persist stats');
   // Periodic sweep of expired cooldowns — keeps health probe cheap and avoids waiting for next user request
   ctx.effect(() => {
     const id = setInterval(() => {
@@ -283,42 +354,55 @@ export function apply(ctx, config = {}) {
     const poolByRef = new Map();
     // provider route (from llm-pi-ai profiles) -> its key pool
     const providerToPool = new Map();
+    // per-model key pools: provider -> Map<model, pool>
+    const modelPoolByProvider = new Map();
     // clone route ids (for the settings dropdown filter)
     const cloneIds = new Set();
 
-    for (const p of cfg.providers ?? []) {
-      const refs = (p.keys ?? []).filter((ref) => typeof ref === 'string' && ref.length > 0);
-      if (refs.length === 0) continue;
-      let state = poolState.get(p.provider);
-      if (!state) {
-        state = {
-          failedUntil: new Map(),
-          failCounts: new Map(),
-          pointer: 0,
-          lastUsed: undefined,
-          // Счётчики для карточки настроек: без них о работе ротации можно было
-          // судить только по console.warn на сервере.
-          switches: 0,
-          lastReason: undefined,
-          lastSwitchAt: undefined,
-          lastExhaustionAt: undefined,
-          exhaustionCount: 0,
-          events: [],
-          usageCounts: new Map(),
+    const makeState = (base) => {
+      let st = poolState.get(base);
+      if (!st) {
+        st = {
+          failedUntil: new Map(), failCounts: new Map(), pointer: 0, lastUsed: undefined,
+          switches: 0, lastReason: undefined, lastSwitchAt: undefined,
+          lastExhaustionAt: undefined, exhaustionCount: 0, events: [], usageCounts: new Map(),
         };
-        poolState.set(p.provider, state);
+        poolState.set(base, st);
       }
-      const poolCooldown = typeof p.cooldownMs === 'number' ? p.cooldownMs : (cfg.cooldownMs ?? 60000);
-      const poolMax = typeof p.maxCooldownMs === 'number' ? p.maxCooldownMs : (cfg.maxCooldownMs ?? undefined);
-      const weights = Array.isArray(p.weights) ? p.weights : [];
+      return st;
+    };
+    const buildPool = (base, keys, weights, poolCooldown, poolMax) => {
+      const refs = (keys ?? []).filter((ref) => typeof ref === 'string' && ref.length > 0);
+      if (refs.length === 0) return null;
+      const w = Array.isArray(weights) ? weights : [];
       const weightedRefs = [];
       for (let i = 0; i < refs.length; i++) {
-        const w = typeof weights[i] === 'number' && weights[i] > 0 ? Math.floor(weights[i]) : 1;
-        for (let k = 0; k < w; k++) weightedRefs.push(refs[i]);
+        const ww = typeof w[i] === 'number' && w[i] > 0 ? Math.floor(w[i]) : 1;
+        for (let k = 0; k < ww; k++) weightedRefs.push(refs[i]);
+      }
+      return { base, refs, weightedRefs: weightedRefs.length > 0 ? weightedRefs : refs,
+               state: makeState(base), cooldownMs: poolCooldown, maxCooldownMs: poolMax };
+    };
+    for (const p of cfg.providers ?? []) {
+      const poolCooldown = typeof p.cooldownMs === 'number' ? p.cooldownMs : (cfg.cooldownMs ?? 60000);
+      const poolMax = typeof p.maxCooldownMs === 'number' ? p.maxCooldownMs : (cfg.maxCooldownMs ?? undefined);
+      // base provider pool (fallback)
+      const pool = buildPool(p.provider, p.keys, p.weights, poolCooldown, poolMax);
+      if (pool) {
+        for (const ref of pool.refs) poolByRef.set(ref, pool);
+        for (let i = 1; i < pool.refs.length; i++) cloneIds.add(`${p.provider}-${i + 1}`);
+      }
+      // per-model pools
+      const models = p.models ?? {};
+      const byModel = new Map();
+      for (const [model, mp] of Object.entries(models)) {
+        const mpool = buildPool(`${p.provider}::${model}`, mp.keys, mp.weights, poolCooldown, poolMax);
+        if (mpool) {
+          byModel.set(model, mpool);
+          for (const ref of mpool.refs) poolByRef.set(ref, mpool);
+        }
       }
-      const pool = { base: p.provider, refs, weightedRefs: weightedRefs.length > 0 ? weightedRefs : refs, state, cooldownMs: poolCooldown, maxCooldownMs: poolMax };
-      for (const ref of refs) poolByRef.set(ref, pool);
-      for (let i = 1; i < refs.length; i++) cloneIds.add(`${p.provider}-${i + 1}`);
+      if (byModel.size > 0) modelPoolByProvider.set(p.provider, byModel);
     }
 
     let profiles = {};
@@ -337,7 +421,7 @@ export function apply(ctx, config = {}) {
     for (const key of [...poolState.keys()]) {
       if (![...poolByRef.values()].some((p) => p.base === key)) poolState.delete(key);
     }
-    return { switchCodes, cooldownMs, maxCooldownMs, notifyWebhook, notifyThreshold, poolByRef, providerToPool, cloneIds };
+    return { switchCodes, cooldownMs, maxCooldownMs, notifyWebhook, notifyThreshold, backupDir, backupIntervalMs, backupKeep, poolByRef, providerToPool, modelPoolByProvider, cloneIds };
   }
 
   // ── patch credentials.resolve: pool refs resolve to the next healthy key ──
@@ -704,6 +788,68 @@ export function apply(ctx, config = {}) {
     },
   }), 'dsh-key-rotation: reset route');
 
+  ctx.effect(() => ctx.webServer.register({
+    kind: 'exact',
+    path: IMPORT_PATH,
+    handler: async (req, res) => {
+      if (req.method !== 'POST') { json(res, 405, { error: { code: 'method', message: 'POST only' } }); return; }
+      if (!isTrustedBridgeRequest(req)) { json(res, 403, { error: { code: 'forbidden', message: 'dsh-key-rotation: import is local-only' } }); return; }
+      let body; try { body = await readJson(req); } catch (e) { json(res, 400, { error: { code: 'bad-request', message: String(e?.message ?? e) } }); return; }
+      const url = typeof body?.url === 'string' ? body.url.trim() : '';
+      if (!url || !url.startsWith('https://')) { json(res, 400, { error: { code: 'bad-url', message: 'dsh-key-rotation: only HTTPS URLs are allowed' } }); return; }
+      try {
+        const resp = await fetch(url);
+        if (!resp.ok) { json(res, 400, { error: { code: 'fetch-failed', message: 'dsh-key-rotation: fetch returned ' + resp.status } }); return; }
+        const data = await resp.json();
+        if (!Array.isArray(data)) { json(res, 400, { error: { code: 'bad-format', message: 'dsh-key-rotation: expected JSON array of providers' } }); return; }
+        const settings = ctx.get('settings');
+        if (!settings) { json(res, 503, { error: { code: 'settings-rejected', message: 'dsh-key-rotation: no settings provider' } }); return; }
+        const desc = settings.describe({ redactSecrets: true }).find((c) => c.ns === NS);
+        const cur = desc?.value?.providers ?? [];
+        const merged = new Map();
+        for (const p of cur) if (p && p.provider) merged.set(p.provider, p);
+        for (const p of data) if (p && p.provider && typeof p.provider === 'string') merged.set(p.provider, p);
+        const mergedArr = [...merged.values()];
+        await settings.replace(NS, { ...(desc?.value ?? {}), providers: mergedArr }, desc?.revision);
+        json(res, 200, { ok: true, providersImported: data.length, total: mergedArr.length });
+      } catch (e) { json(res, 400, { error: { code: 'import-failed', message: String(e?.message ?? e) } }); }
+    },
+  }), 'dsh-key-rotation: import route');
+
+  // Health for external panels (Beszel/Uptime)
+  ctx.effect(() => ctx.webServer.register({
+    kind: 'exact',
+    path: HEALTH_PATH,
+    handler: async (req, res) => {
+      if (!isTrustedBridgeRequest(req) && req.socket?.remoteAddress !== '127.0.0.1' && req.socket?.remoteAddress !== '::1') { } // allow same-origin already checked
+      if (!isTrustedBridgeRequest(req)) {
+        // also allow plain loopback without Origin
+        if (!isLoopbackAddress(req.socket?.remoteAddress)) { res.writeHead(403); res.end(); return; }
+        if (req.headers['sec-fetch-site'] === 'cross-site') { res.writeHead(403); res.end(); return; }
+      }
+      if (req.method !== 'GET') { json(res, 405, { error: { code: 'method', message: 'GET only' } }); return; }
+      const now = Date.now();
+      const pools = {};
+      let exhaustedAny = false;
+      const { poolByRef: pr } = buildRuntime();
+      const seenH = new Set();
+      for (const pool of pr.values()) {
+        if (seenH.has(pool.base)) continue;
+        seenH.add(pool.base);
+        let healthy = 0;
+        for (const ref of pool.refs) {
+          const until = pool.state.failedUntil.get(ref);
+          if (!(until !== undefined && until > now)) healthy++;
+        }
+        const total = pool.refs.length;
+        const exhausted = healthy === 0 && total > 0;
+        if (exhausted) exhaustedAny = true;
+        pools[pool.base] = { healthy, total, exhausted };
+      }
+      json(res, 200, { status: exhaustedAny ? 'degraded' : 'ok', pools, exhaustedAny });
+    },
+  }), 'dsh-key-rotation: health');
+
   // ── test route: dry-run a single key without rotation ──
   ctx.effect(() => ctx.webServer.register({
     kind: 'exact',
@@ -738,8 +884,9 @@ export function apply(ctx, config = {}) {
   // straight through.
   ctx.on('llm/stream', (options, next) => {
     if (options[MARKER]) return next();
-    const { providerToPool } = buildRuntime();
-    const pool = providerToPool.get(options.provider);
+    const { providerToPool, modelPoolByProvider } = buildRuntime();
+    const byModel = modelPoolByProvider.get(options.provider);
+    const pool = (byModel && byModel.get(options.model)) || providerToPool.get(options.provider);
     if (!pool) return next();
     console.warn(`[dsh-key-rotation] rotating ${options.provider}/${options.model} across ${(pool.weightedRefs ?? pool.refs).length} slots (${pool.refs.length} keys)`);
     return rotate(options, pool);
@@ -752,8 +899,10 @@ export function apply(ctx, config = {}) {
   ctx.on('agent/request-error', async (payload, next) => {
     const provider = payload?.provider ?? payload?.failure?.provider ?? '';
     if (!provider) return next();
-    const { providerToPool, switchCodes } = buildRuntime();
-    const pool = providerToPool.get(provider);
+    const { providerToPool, modelPoolByProvider, switchCodes } = buildRuntime();
+    const model = payload?.model || payload?.failure?.model || '';
+    const byModel = modelPoolByProvider.get(provider);
+    const pool = (byModel && byModel.get(model)) || providerToPool.get(provider);
     if (!pool) return next();
     const code = String(payload?.failure?.code ?? payload?.code ?? '');
     const message = String(payload?.failure?.message ?? payload?.message ?? '');

package/lib/pool.js

@@ -176,3 +176,10 @@ export function parseRetryAfter(value) {
   }
   return undefined;
 }
+
+/** Pick a key pool for a (provider, model) pair. Model sub-pools win over the
+ *  provider base pool; falls back to the base pool when no sub-pool matches. */
+export function selectPool(modelPoolByProvider, providerToPool, provider, model) {
+  const byModel = modelPoolByProvider && modelPoolByProvider.get(provider);
+  return (byModel && byModel.get(model)) || (providerToPool && providerToPool.get(provider)) || null;
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@goodandready/dsh-key-rotation",
-  "version": "0.7.10",
+  "version": "0.7.11",
   "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": [
     "deepseek-harness",
@@ -54,8 +54,5 @@
   },
   "scripts": {
     "test": "node --test test/*.test.js"
-  },
-  "publishConfig": {
-    "access": "public"
   }
 }