@moguiyu/dsh-tavily 0.3.2 → 0.3.4

package/src/index.js

@@ -4,7 +4,8 @@
  * Both halves live here, so there is no cross-package contract left to skew:
  *
  * - the model tools (`tools.js`) — `tavily_search`, `tavily_extract`,
- *   `tavily_map`, `tavily_crawl`, with key rotation and failover on 401/429;
+ *   `tavily_map`, `tavily_crawl`, with key rotation and failover on 401/429
+ *   and on Tavily's per-key 432/433 usage-limit refusals;
  * - key and usage management (`backend.js`) — `/api/tavily-usage` and
  *   `/api/tavily-manager`;
  * - the browser card (`client.js`) — keys, usage, and strategy.

package/src/lib.js

@@ -3,12 +3,130 @@
  */
 import { readFileSync } from 'node:fs'
 
-export const STRATEGIES = ['rotate', 'low-usage-first', 'high-usage-first']
+export const STRATEGIES = ['rotate', 'low-usage-first', 'high-usage-first', 'load-balance']
 
 export function isValidStrategy(value) {
   return STRATEGIES.includes(value)
 }
 
+/**
+ * Concurrent requests allowed per Tavily account in `load-balance` mode. One is
+ * the safe default: Tavily's published limits are per environment/key
+ * (100 requests/minute on a free Development key) and every key of one account
+ * shares that budget and its credit pool.
+ */
+export const DEFAULT_ACCOUNT_CAP = 1
+export const MAX_ACCOUNT_CAP = 8
+
+export function clampAccountCap(value) {
+  const n = Number(value)
+  if (!Number.isFinite(n)) return DEFAULT_ACCOUNT_CAP
+  return Math.min(MAX_ACCOUNT_CAP, Math.max(DEFAULT_ACCOUNT_CAP, Math.round(n)))
+}
+
+/** The `/usage` `account` fields that identify an account; order is canonical. */
+const ACCOUNT_FIELDS = [
+  'current_plan', 'plan', 'plan_limit', 'plan_usage', 'paygo_limit', 'paygo_usage',
+  'search_usage', 'extract_usage', 'map_usage', 'crawl_usage', 'research_usage',
+]
+
+/**
+ * Fingerprint a Tavily `/usage` `account` block. Tavily publishes no account id,
+ * but every key of one account reports the identical account block at the same
+ * moment (observed live: two keys both at search_usage 846 / extract 129 / map
+ * 23), so an equal fingerprint is the only account-level signal available.
+ * Returns `null` when the block carries no recognizable field, which means
+ * "makes no claim about the account" rather than "its own account".
+ */
+export function accountFingerprint(account) {
+  if (account === null || account === undefined || typeof account !== 'object' || Array.isArray(account)) return null
+  const parts = []
+  for (const field of ACCOUNT_FIELDS) {
+    const value = account[field]
+    if (value === undefined || value === null) continue
+    parts.push(field + '=' + String(value))
+  }
+  return parts.length === 0 ? null : parts.join('|')
+}
+
+/**
+ * Assign stable account group ids to the stored keys.
+ *
+ * `entries` is `[{ masked, fingerprint }]` in credential-list order; `previous`
+ * is the persisted `accounts` map. A `manual: true` entry is kept verbatim —
+ * that is the user's override — and offers its group id to keys that turn out
+ * to share its account. Inferred ids survive a re-run: a cluster carries over
+ * the id its own members already had instead of being renumbered, so usage
+ * drifting between fetches cannot reshuffle the groups.
+ */
+export function groupAccounts(entries, previous) {
+  const prior = previous !== null && previous !== undefined && typeof previous === 'object' && !Array.isArray(previous)
+    ? previous
+    : {}
+  const assigned = {}
+  const used = new Set()
+  const manualGroupByFingerprint = new Map()
+
+  for (const entry of entries) {
+    const before = prior[entry.masked]
+    if (before === null || before === undefined || typeof before !== 'object') continue
+    if (before.manual !== true) continue
+    if (typeof before.group !== 'string' || before.group.length === 0) continue
+    assigned[entry.masked] = { group: before.group, manual: true }
+    used.add(before.group)
+    if (typeof entry.fingerprint === 'string' && entry.fingerprint.length > 0 && !manualGroupByFingerprint.has(entry.fingerprint)) {
+      manualGroupByFingerprint.set(entry.fingerprint, before.group)
+    }
+  }
+
+  const clusters = new Map()
+  for (const entry of entries) {
+    if (assigned[entry.masked] !== undefined) continue
+    const fingerprint = typeof entry.fingerprint === 'string' && entry.fingerprint.length > 0 ? entry.fingerprint : null
+    if (fingerprint === null) continue
+    const members = clusters.get(fingerprint)
+    if (members === undefined) clusters.set(fingerprint, [entry.masked])
+    else members.push(entry.masked)
+  }
+
+  let next = 1
+  for (const [fingerprint, members] of clusters) {
+    const carried = new Map()
+    for (const masked of members) {
+      const before = prior[masked]
+      if (before !== null && before !== undefined && typeof before === 'object' &&
+        typeof before.group === 'string' && before.group.length > 0) {
+        carried.set(before.group, (carried.get(before.group) ?? 0) + 1)
+      }
+    }
+    let group = manualGroupByFingerprint.get(fingerprint)
+    if (group === undefined) {
+      let best = null
+      let bestCount = 0
+      for (const [id, count] of carried) {
+        if (count > bestCount) {
+          best = id
+          bestCount = count
+        }
+      }
+      if (best !== null) group = best
+    }
+    if (group === undefined) {
+      while (used.has('acct-' + next)) next += 1
+      group = 'acct-' + next
+      next += 1
+    }
+    used.add(group)
+    for (const masked of members) assigned[masked] = { group, manual: false }
+  }
+
+  const out = {}
+  for (const entry of entries) {
+    if (assigned[entry.masked] !== undefined) out[entry.masked] = assigned[entry.masked]
+  }
+  return out
+}
+
 /**
  * Mask a key for display: `first12…last4`; comma-separated lists are masked
  * per part. Never reveals more than a fingerprint.
@@ -39,10 +157,12 @@ export function parseKeyList(value) {
 
 /**
  * Order keys according to a strategy. `usageOf(key)` returns a number or null
- * (unknown usage is treated as 0). `rotate` keeps the given order.
+ * (unknown usage is treated as 0). `rotate` keeps the given order, and so does
+ * `load-balance`: that mode is about how many requests may be in flight per
+ * account, not about which key comes first.
  */
 export function orderKeys(values, strategy, usageOf) {
-  if (strategy === 'rotate') return [...values]
+  if (strategy === 'rotate' || strategy === 'load-balance') return [...values]
   const usageOfKey = (value) => {
     const usage = usageOf(value)
     return usage !== null && typeof usage === 'number' ? usage : 0

package/src/tools.js

@@ -2,7 +2,8 @@
  * Tool half of `@moguiyu/dsh-tavily`: the advanced `tavily_search` model tool
  * plus the extra Tavily direct tools (`tavily_extract`, `tavily_map`,
  * `tavily_crawl`). Keys resolve from the `TAVILY_API_KEYS` credential per
- * call; rotation is round-robin with failover on 401/429.
+ * call; rotation is round-robin with failover on every per-key refusal
+ * (401/429, plus Tavily's 432 plan-limit and 433 pay-as-you-go-limit codes).
  *
  * The tools are an OPTION for search, never a replacement: they are
  * independent of the built-in `web_search` provider, and registering them
@@ -15,10 +16,130 @@
  * (0.1.6 row restart, then the 0.2.x dependency-range skew). Removing the
  * capability removes the failure mode; see `docs/agents/verification.md`.
  */
+import { join } from 'node:path'
+import { resolveDshHome } from '@deepseek-ai/dsh-home-paths'
 import { defineTool } from '@deepseek-ai/dsh-tools'
+import { clampAccountCap, maskValue, readJsonFile } from './lib.js'
 
 const TAVILY_BASE_URL = 'https://api.tavily.com'
 
+/**
+ * Statuses that mean "this key cannot serve the request — try another one".
+ * Tavily documents 432 (plan limit exceeded) and 433 (pay-as-you-go limit
+ * exceeded) as conditions of the key's account, so in a multi-key pool they
+ * are failover cases exactly like 401/429. Treating them as fatal made a pool
+ * with one exhausted account fail calls that a healthy key could have served.
+ */
+const RETRYABLE_STATUS = new Set([401, 429, 432, 433])
+
+/** Why a retryable status refused the key; the pool's exhaustion report uses it. */
+const RETRYABLE_REASON = new Map([
+  [401, 'invalid or missing key'],
+  [429, 'rate limit'],
+  [432, 'Tavily plan usage limit exceeded'],
+  [433, 'Tavily pay-as-you-go limit exceeded'],
+])
+
+const MANAGER_STATE = 'tavily-manager.json'
+
+/**
+ * The concurrency policy the card writes: the strategy, the per-account cap,
+ * and the account grouping the usage route inferred (or the user pinned). Read
+ * per call, so a card change takes effect on the next search, and tolerant of a
+ * missing or half-written file.
+ */
+function readConcurrencyPolicy() {
+  const state = readJsonFile(join(resolveDshHome(), MANAGER_STATE), {})
+  const accounts = state.accounts !== null && state.accounts !== undefined &&
+    typeof state.accounts === 'object' && !Array.isArray(state.accounts)
+    ? state.accounts
+    : {}
+  return {
+    balanced: state.strategy === 'load-balance',
+    cap: clampAccountCap(state.accountCap),
+    accounts,
+  }
+}
+
+/**
+ * The account one key belongs to. A key the usage route could not group is
+ * treated as its own account: the cap still applies to it, and it stays
+ * parallel-safe with keys of other accounts.
+ */
+function accountOfKey(key, accounts) {
+  const masked = maskValue(key)
+  const entry = accounts[masked]
+  if (entry !== null && entry !== undefined && typeof entry === 'object' &&
+    typeof entry.group === 'string' && entry.group.length > 0) {
+    return entry.group
+  }
+  return 'key:' + masked
+}
+
+/**
+ * Per-account in-flight limiter. One instance per installed tool half, so every
+ * tool call in the process shares it — which is the point: parallel calls queue
+ * behind one free account instead of hitting it at the same time.
+ */
+function createAccountScheduler() {
+  const active = new Map()
+  const waiters = []
+  const activeIn = (group) => active.get(group) ?? 0
+
+  const detach = (waiter) => {
+    if (waiter.signal !== null && waiter.signal !== undefined && waiter.onAbort !== undefined) {
+      waiter.signal.removeEventListener('abort', waiter.onAbort)
+    }
+  }
+
+  function wake() {
+    for (let index = 0; index < waiters.length; index += 1) {
+      const waiter = waiters[index]
+      if (activeIn(waiter.group) >= waiter.cap) continue
+      waiters.splice(index, 1)
+      index -= 1
+      detach(waiter)
+      // Reserve the slot here: waking two waiters must never hand out one slot.
+      active.set(waiter.group, activeIn(waiter.group) + 1)
+      waiter.resolve()
+    }
+  }
+
+  return {
+    activeIn,
+    /** Resolves `true` when the caller had to wait for capacity. */
+    async acquire(group, cap, signal) {
+      if (activeIn(group) < cap) {
+        active.set(group, activeIn(group) + 1)
+        return false
+      }
+      const aborted = () => new Error('tavily: request aborted while waiting for account capacity')
+      if (signal !== null && signal !== undefined && signal.aborted === true) throw aborted()
+      await new Promise((resolve, reject) => {
+        const waiter = { group, cap, resolve, signal }
+        if (signal !== null && signal !== undefined) {
+          waiter.onAbort = () => {
+            const index = waiters.indexOf(waiter)
+            if (index >= 0) waiters.splice(index, 1)
+            detach(waiter)
+            reject(aborted())
+          }
+          signal.addEventListener('abort', waiter.onAbort, { once: true })
+        }
+        waiters.push(waiter)
+      })
+      // The waker already reserved this slot.
+      return true
+    },
+    release(group) {
+      const next = activeIn(group) - 1
+      if (next <= 0) active.delete(group)
+      else active.set(group, next)
+      wake()
+    },
+  }
+}
+
 /** Clamp a number into [min, max]; `fallback` when not finite. */
 export function clampInt(value, min, max, fallback) {
   const n = Number(value)
@@ -103,7 +224,7 @@ async function callOperation(operation, body, key, exec, timeoutMs) {
       body: JSON.stringify(body),
       signal: combined.signal,
     })
-    if (response.status === 401 || response.status === 429) {
+    if (RETRYABLE_STATUS.has(response.status)) {
       return { kind: 'retryable', statusCode: response.status }
     }
     if (!response.ok) {
@@ -113,6 +234,8 @@ async function callOperation(operation, body, key, exec, timeoutMs) {
         if (parsed !== null && typeof parsed === 'object') {
           if (typeof parsed.error === 'string') message = parsed.error
           else if (parsed.error !== null && typeof parsed.error === 'object' && typeof parsed.error.message === 'string') message = parsed.error.message
+          // 432/433 and other API failures arrive as {"detail":{"error":…}}.
+          else if (parsed.detail !== undefined && parsed.detail !== null && typeof parsed.detail === 'object' && typeof parsed.detail.error === 'string') message = parsed.detail.error
         }
       } catch {
         /* keep the status message */
@@ -430,6 +553,7 @@ const NAVIGATION_PARAMETERS = {
  */
 export function installTavilyTool(ctx) {
   let rotation = 0
+  const scheduler = createAccountScheduler()
 
   // Every registration is collected into one disposer so the fiber tears the
   // half down cleanly: on row unload, and on the Plugins page's runtime unload
@@ -455,19 +579,66 @@ export function installTavilyTool(ctx) {
     return { keys: [], source: 'unconfigured' }
   }
 
-  /** Run `body` through the rotation/failover loop; returns the response JSON. */
+  /**
+   * Run `body` through the rotation/failover loop; returns the response JSON.
+   *
+   * The first key is picked from the shared round-robin cursor, and each
+   * attempt advances by exactly one slot, so one call tries every configured
+   * key at most once. Folding the attempt offset into the cursor instead
+   * (the earlier shape) skipped keys and repeated others whenever the pool
+   * size was even.
+   *
+   * Under `load-balance` the slot is chosen by account capacity instead: the
+   * first key, from the cursor, whose account is below the cap — so parallel
+   * calls spread across accounts, and a call queues rather than opening a
+   * second connection to one free account. Every other strategy keeps the
+   * plain order above, untouched.
+   */
   async function requestWithRotation(operation, body, timeoutMs, exec, detail) {
     const { keys, source } = await resolveKeys()
     if (keys.length === 0) {
       throw new Error(`tavily_${operation}: TAVILY_API_KEYS is not configured — add it via the Tavily Search settings card or ~/.dsh/.credentials.yaml (comma-separated)`)
     }
     const attempts = keys.length
+    const first = rotation
+    const policy = readConcurrencyPolicy()
+    const groups = policy.balanced ? keys.map((key) => accountOfKey(key, policy.accounts)) : null
+    const signal = exec !== null && exec !== undefined ? exec.signal : undefined
+    const tried = new Set()
+    const pickBalancedIndex = (fallbackAttempt) => {
+      let leastBusy = null
+      let leastLoad = Infinity
+      for (let step = 0; step < keys.length; step += 1) {
+        const candidate = (first + step) % keys.length
+        if (tried.has(candidate)) continue
+        const load = scheduler.activeIn(groups[candidate])
+        if (load < policy.cap) return candidate
+        if (load < leastLoad) {
+          leastBusy = candidate
+          leastLoad = load
+        }
+      }
+      return leastBusy === null ? (first + fallbackAttempt) % keys.length : leastBusy
+    }
     let lastRetryable = 0
     for (let attempt = 0; attempt < attempts; attempt++) {
-      const index = (rotation + attempt) % keys.length
+      const index = groups === null ? (first + attempt) % keys.length : pickBalancedIndex(attempt)
+      tried.add(index)
       const key = keys[index]
+      const group = groups === null ? null : groups[index]
+      if (group !== null) {
+        const waited = await scheduler.acquire(group, policy.cap, signal)
+        if (waited) {
+          ctx.logger.info(`tavily_${operation}: waited for account capacity %j`, { group, cap: policy.cap, keys: keys.length })
+        }
+      }
       const started = Date.now()
-      const outcome = await callOperation(operation, body, key, exec, timeoutMs)
+      let outcome
+      try {
+        outcome = await callOperation(operation, body, key, exec, timeoutMs)
+      } finally {
+        if (group !== null) scheduler.release(group)
+      }
       if (outcome.kind === 'ok') {
         rotation = (index + 1) % keys.length
         const extra = typeof detail === 'function' ? detail(outcome.json) : {}
@@ -478,7 +649,8 @@ export function installTavilyTool(ctx) {
       rotation = (index + 1) % keys.length
       ctx.logger.info(`tavily_${operation}: key %d/%d returned HTTP %d, rotating`, index + 1, keys.length, outcome.statusCode)
     }
-    throw new Error(`tavily_${operation}: all ${attempts} configured key(s) failed with HTTP ${lastRetryable} (invalid key or rate limit)`)
+    const reason = RETRYABLE_REASON.get(lastRetryable) ?? 'invalid key or rate limit'
+    throw new Error(`tavily_${operation}: all ${attempts} configured key(s) failed with HTTP ${lastRetryable} (${reason})`)
   }
 
   tools.register(defineTool({