diff options
| author | Adam Malczewski <[email protected]> | 2026-06-28 14:50:27 +0900 |
|---|---|---|
| committer | Adam Malczewski <[email protected]> | 2026-06-28 14:50:27 +0900 |
| commit | fb4a9217b55dd3ba11670104ac23536416d36940 (patch) | |
| tree | 4ae4aed5ca1bbc09620e0ab53e54cf54b5cbd2f2 | |
| parent | df7d45b58076b6ab47bea3148872cac11401d510 (diff) | |
| download | dispatch-fb4a9217b55dd3ba11670104ac23536416d36940.tar.gz dispatch-fb4a9217b55dd3ba11670104ac23536416d36940.zip | |
fix(concurrency): usage-gate fast-path overshoot + fetchUsage rejection safety + persist auto-reduce notice
Bug 1 (overshoot): acquire() fast-path now queues while a recycle-poll is in
flight (gatePolling), and the recycle defers its inFlight decrement until the
usage poll resolves — holding inFlight inflated through the poll window so a
concurrent caller cannot sneak through before upstream confirms room. The gated
path admits exactly one waiter per fresh poll (grantOne), robust against stale
upstream counts. Common-case throughput preserved: the fast-path still grants
immediately when no poll is in flight.
Bug 2 (unhandled rejection): fetchUsage is wrapped in safeFetchUsage — a
throwing getUsage() is caught, treated as undefined (cooldown-only fallback),
and reported via the new onUsagePollError opt (warn-level). No unhandled
promise rejection can crash the process.
Bug 3 (lost notice): loadLimits now calls restoreLimit (a new lower-level
state-seeding method that does NOT clear autoReduced) instead of setLimit
(which is a manual user action that clears the notice). The auto-reduce marker
(autoReducedFrom) is persisted under auto-reduce:<providerId> and restored on
activate via loadAutoReduce, so the frontend banner survives a restart.
+6 tests covering each bug.
3 files changed, 393 insertions, 32 deletions
diff --git a/packages/provider-concurrency/src/concurrency-manager.test.ts b/packages/provider-concurrency/src/concurrency-manager.test.ts index 6db8964..185c6c2 100644 --- a/packages/provider-concurrency/src/concurrency-manager.test.ts +++ b/packages/provider-concurrency/src/concurrency-manager.test.ts @@ -64,6 +64,7 @@ function createManager(opts?: { releaseCooldownMs?: number; fetchUsage?: (providerId: string) => Promise<ProviderUsage | undefined>; onLimitReduced?: (providerId: string, newLimit: number, oldLimit: number) => void; + onUsagePollError?: (providerId: string, err: unknown) => void; }): { manager: ConcurrencyService; timers: ReturnType<typeof createFakeTimers>; @@ -77,6 +78,7 @@ function createManager(opts?: { ...(opts?.releaseCooldownMs !== undefined ? { releaseCooldownMs: opts.releaseCooldownMs } : {}), ...(opts?.fetchUsage !== undefined ? { fetchUsage: opts.fetchUsage } : {}), ...(opts?.onLimitReduced !== undefined ? { onLimitReduced: opts.onLimitReduced } : {}), + ...(opts?.onUsagePollError !== undefined ? { onUsagePollError: opts.onUsagePollError } : {}), setTimeout: timers.setTimeout, clearTimeout: timers.clearTimeout, setInterval: timers.setInterval, @@ -797,4 +799,170 @@ describe("createConcurrencyManager", () => { expect(() => timers.advance(1000)).not.toThrow(); }); }); + + // ─── Bug 1: usage-gate fast-path anti-overshoot ───────────────────────── + + it("Bug 1: a concurrent acquire during a recycle-poll queues instead of fast-pathing (no overshoot)", async () => { + // The poll resolves only on an explicit microtask flush (deferred), so a + // concurrent acquire arriving mid-poll must see gatePolling/inflated inFlight. + let resolvePoll: (snap: ProviderUsage) => void = () => {}; + const pollCalled: number[] = []; + const { manager } = createManager({ + fetchUsage: () => + new Promise<ProviderUsage>((resolve) => { + pollCalled.push(1); + resolvePoll = resolve; + }), + }); + manager.setLimit("umans", 1); + manager.setCooldown("umans", 0); + + // Hold the single slot. + const release1 = await manager.acquire("umans", "c1", 0); + expect(manager.getStatus("umans")?.inFlight).toBe(1); + + // Queue a waiter (c2). Cooldown is 0, but the gate defers admission until a poll. + let c2Granted = false; + const p2 = manager.acquire("umans", "c2", 10).then((r) => { + c2Granted = true; + return r; + }); + await Promise.resolve(); + await Promise.resolve(); + + // Release c1 → recycle → poll started (inFlight held inflated during poll). + release1(); + await Promise.resolve(); // let recycle schedule the poll + await Promise.resolve(); + expect(pollCalled.length).toBeGreaterThanOrEqual(1); + // inFlight is still 1 (the recycle's decrement is deferred until the poll). + expect(manager.getStatus("umans")?.inFlight).toBe(1); + + // A NEW acquire arriving mid-poll: inFlight is 1 (== limit) → must QUEUE, + // not fast-path. Even if it saw inFlight < limit, gatePolling would route it + // through the queue. Either way it must NOT be granted yet. + let c3Granted = false; + const p3 = manager.acquire("umans", "c3", 20).then((r) => { + c3Granted = true; + return r; + }); + await Promise.resolve(); + await Promise.resolve(); + expect(c3Granted).toBe(false); + expect(manager.getStatus("umans")?.queued).toBeGreaterThanOrEqual(1); + + // Resolve the poll with room (0 < 1) → c2 admitted (inFlight: decrement then + // re-increment for the grant). c3 stays queued (one admission per poll). + resolvePoll({ concurrentSessions: 0 }); + const release2 = await p2; + expect(c2Granted).toBe(true); + // c3 NOT admitted by this poll (one per poll). + expect(c3Granted).toBe(false); + + release2(); + void p3; + }); + + it("Bug 1: when no poll is in flight, the fast-path still grants immediately (common-case throughput preserved)", async () => { + const { manager } = createManager({ + fetchUsage: async () => ({ concurrentSessions: 0 }), + }); + manager.setLimit("umans", 4); + + // Nowhere near the limit, no recycle in progress → fast-path, no poll. + const release = await manager.acquire("umans", "c1", 0); + expect(manager.getStatus("umans")?.inFlight).toBe(1); + release(); + }); + + // ─── Bug 2: fetchUsage exceptions don't become unhandled rejections ────── + + it("Bug 2: a throwing fetchUsage is treated as undefined (cooldown-only fallback) and fires onUsagePollError", async () => { + let pollError: { providerId: string; err: unknown } | undefined; + const { manager } = createManager({ + fetchUsage: async () => { + throw new Error("usage endpoint exploded"); + }, + onUsagePollError: (providerId, err) => { + pollError = { providerId, err }; + }, + }); + manager.setLimit("umans", 1); + manager.setCooldown("umans", 0); + + const release1 = await manager.acquire("umans", "c1", 0); + // Queue a waiter; release → recycle → poll THROWS. + const p2 = manager.acquire("umans", "c2", 10); + await Promise.resolve(); + await Promise.resolve(); + + // Must NOT reject / throw unhandled — swallow + fall back to granting. + release1(); + const release2 = await p2; // resolves (cooldown-only fallback grants). + expect(release2).toBeTypeOf("function"); + expect(pollError?.providerId).toBe("umans"); + expect(pollError?.err).toBeInstanceOf(Error); + release2(); + }); + + it("Bug 2: no unhandled promise rejection is left when fetchUsage throws (process stays clean)", async () => { + const rejections: unknown[] = []; + const handler = (reason: unknown) => rejections.push(reason); + process.on("unhandledRejection", handler); + try { + const { manager } = createManager({ + fetchUsage: async () => { + throw new Error("boom"); + }, + }); + manager.setLimit("umans", 1); + manager.setCooldown("umans", 0); + + const release1 = await manager.acquire("umans", "c1", 0); + manager.acquire("umans", "c2", 10).then((r) => r()); // queue + auto-release + await Promise.resolve(); + await Promise.resolve(); + release1(); + // Let the swallowed poll + grant settle fully. + await new Promise((r) => setTimeout(r, 5)); + await new Promise((r) => setTimeout(r, 5)); + expect(rejections).toEqual([]); + } finally { + process.off("unhandledRejection", handler); + } + }); + + // ─── Bug 3: persisted auto-reduced limit keeps its notice across restart ─ + + it("Bug 3: restoreLimit (startup) preserves the auto-reduce notice that setLimit (manual) clears", () => { + const { manager } = createManager(); + manager.setLimit("umans", 4); + manager.reportRateLimit("umans"); // 4 -> 3, autoReduced + expect(manager.getStatus("umans")?.autoReduced).toBe(true); + expect(manager.getStatus("umans")?.autoReducedFrom).toBe(4); + + // Simulate a restart: a fresh manager restores the persisted limit (3) + + // the auto-reduce marker (autoReducedFrom=4) via restoreLimit. + const { manager: restarted } = createManager(); + restarted.restoreLimit("umans", 3, 4); + const status = restarted.getStatus("umans"); + expect(status?.limit).toBe(3); + expect(status?.autoReduced).toBe(true); + expect(status?.autoReducedFrom).toBe(4); + expect(status?.notice).toContain("auto-reduced to 3"); + + // Contrast: a MANUAL setLimit clears the notice (user took control). + restarted.setLimit("umans", 4); + expect(restarted.getStatus("umans")?.autoReduced).toBe(false); + expect(restarted.getStatus("umans")?.autoReducedFrom).toBeUndefined(); + }); + + it("Bug 3: restoreLimit without autoReducedFrom does not synthesize a notice", () => { + const { manager } = createManager(); + manager.restoreLimit("umans", 4); + const status = manager.getStatus("umans"); + expect(status?.limit).toBe(4); + expect(status?.autoReduced).toBe(false); + expect(status?.notice).toBeUndefined(); + }); }); diff --git a/packages/provider-concurrency/src/concurrency-manager.ts b/packages/provider-concurrency/src/concurrency-manager.ts index 2d3d857..1d05bb0 100644 --- a/packages/provider-concurrency/src/concurrency-manager.ts +++ b/packages/provider-concurrency/src/concurrency-manager.ts @@ -117,6 +117,15 @@ export interface ConcurrencyLimiter { export interface ConcurrencyService extends ConcurrencyLimiter { /** Set the concurrency limit for a provider (MANUAL — clears the auto-reduce notice). Creates the state if new. */ setLimit(providerId: string, limit: number): void; + /** + * Restore a persisted limit on startup WITHOUT clearing the auto-reduce + * notice (Bug 3). Unlike {@link setLimit} (a manual user action that signals + * "the user took control"), this seeds state from disk: it applies the limit + * and, when `autoReducedFrom` is provided, re-marks the state as auto-reduced + * so the frontend banner survives a restart. Used by the extension's + * `loadLimits`/`loadAutoReduce` on activate. + */ + restoreLimit(providerId: string, limit: number, autoReducedFrom?: number): void; /** Get the configured limit, or `undefined` when none. */ getLimit(providerId: string): number | undefined; /** Remove the limit for a provider (makes it unlimited). */ @@ -216,6 +225,13 @@ export interface ConcurrencyManagerOpts { readonly onPause?: (providerId: string, durationMs: number) => void; /** Fired when a 429 adaptively reduces a provider's limit (for persistence + logging). */ readonly onLimitReduced?: (providerId: string, newLimit: number, oldLimit: number) => void; + /** + * Fired when the injected `fetchUsage` throws (network/parse failure beyond the + * graceful-undefined path). The manager treats a thrown poll as "no usage info" + * (cooldown-only fallback) — this callback is for WARN-level logging only. The + * poll never becomes an unhandled rejection. + */ + readonly onUsagePollError?: (providerId: string, err: unknown) => void; } /** Min interval between usage-gate fallback repolls (ms). The release trigger is immediate. */ @@ -277,14 +293,29 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre released = true; state.slots.delete(id); - // Recycle the slot: decrement inFlight + attempt to grant the next waiter. - // With a release cooldown > 0, defer this by the cooldown duration so - // the upstream provider has time to decrement its concurrent_sessions + // Recycle the slot: free its inFlight count + attempt to grant the next + // waiter. With a release cooldown > 0, defer this by the cooldown duration + // so the upstream provider has time to decrement its concurrent_sessions // counter — preventing an N+1 overshoot from accounting lag. During the // cooldown, inFlight stays incremented, so new acquires queue. const recycle = () => { - state.inFlight--; - void tryGrantNext(providerId); + if (fetchUsage === undefined || state.queue.length === 0) { + // No usage gate, OR no one waiting (the lag window is irrelevant when + // there is no waiter to admit) → free the slot immediately. With no + // gate, also drain the queue (grant all that fit). + state.inFlight--; + if (fetchUsage === undefined) grantLoop(state, providerId); + return; + } + // Usage gate configured + a waiter exists → hold inFlight inflated + // DURING the poll window (gatePolling is set synchronously inside + // pollAndGrant, the inFlight decrement is deferred until the poll + // resolves). This closes the overshoot gap: a concurrent acquire arriving + // between the cooldown firing and the poll resolving sees the slot as + // still occupied (inFlight >= limit) and queues instead of fast-pathing. + // pollAndGrant(decrementOnPoll=true) decrements inFlight after observing + // the post-release upstream state, then admits one waiter if there is room. + void pollAndGrant(providerId, state, true); }; if (state.cooldownMs > 0) { const timer = setTimeout(() => { @@ -307,8 +338,9 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre /** * Grant queued waiters WITHOUT the usage gate (the fast path used when no - * `fetchUsage` is configured, or after a poll confirmed upstream has room). - * Grants while there is internal room (`inFlight < limit`). Synchronous. + * `fetchUsage` is configured, or as the cooldown-only fallback when a poll + * returns no usage info). Grants while there is internal room + * (`inFlight < limit`). Synchronous. */ function grantLoop(state: ProviderState, providerId: string): void { while (state.queue.length > 0 && state.inFlight < state.limit) { @@ -326,15 +358,56 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre } /** + * Admit exactly ONE queued waiter (the front of the queue), if there is + * internal room. Used by the usage-gated path so each admission is confirmed + * by a FRESH upstream poll — admitting multiple from a single (possibly stale) + * poll risks an N+1 overshoot when the upstream count lags. Additional waiters + * are admitted on subsequent repolls. + */ + function grantOne(state: ProviderState, providerId: string): void { + if (state.queue.length === 0) return; + if (state.inFlight >= state.limit) return; + const waiter = state.queue[0]; + if (waiter === undefined) return; + state.queue.shift(); + const releaseFn = grantSlot(state, providerId, waiter.conversationId); + waiter.resolve(releaseFn); + // If the queue drained, disarm the fallback timer. + if (state.queue.length === 0 && state.gateRepollTimer !== undefined) { + clearTimeout(state.gateRepollTimer); + state.gateRepollTimer = undefined; + } + } + + /** + * Invoke the injected `fetchUsage`, treating ANY thrown error as "no usage + * info available" (cooldown-only fallback) — so a throwing `getUsage()` never + * becomes an unhandled rejection. The `onUsagePollError` opt is fired for + * WARN-level logging. Returns `undefined` on throw (Bug 2 fix). + */ + async function safeFetchUsage(providerId: string): Promise<ProviderUsage | undefined> { + if (fetchUsage === undefined) return undefined; + try { + return await fetchUsage(providerId); + } catch (err) { + opts.onUsagePollError?.(providerId, err); + return undefined; + } + } + + /** * Drain the queue, gated on the upstream usage poll when `fetchUsage` is - * configured. Called from release (post-cooldown), setLimit, pause-expiry, - * and the repoll timer. Async because the usage poll is an injected I/O + * configured. Called from setLimit, pause-expiry, and the repoll timer (NOT + * from release — that goes through {@link recycleGated}, which holds inFlight + * inflated during the poll). Async because the usage poll is an injected I/O * effect; callers fire-and-forget the returned promise. * - * Only QUEUED agents are gated. The fast-path immediate grant in `acquire` - * (when `inFlight < limit`) is NOT gated — it fires only when there is - * clearly internal room, and the cooldown keeps `inFlight` inflated during - * the accounting-lag window so the fast-path does not fire then. + * The fast-path immediate grant in `acquire` (when `inFlight < limit`) is + * disabled while `gatePolling` is true — `acquire` queues instead, so a + * concurrent caller cannot sneak through the accounting-lag / poll window + * (anti-overshoot). When no poll is in flight the fast-path is safe: the + * cooldown keeps `inFlight` inflated during the lag window, and a recycle + * sets `gatePolling` synchronously before decrementing. * * Each successful poll admits at most ONE queued waiter (each admission pushes * the upstream count back toward the limit); additional waiters are admitted @@ -364,28 +437,41 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre await pollAndGrant(providerId, state); } - /** Poll upstream usage, then grant if there is headroom. */ - async function pollAndGrant(providerId: string, state: ProviderState): Promise<void> { + /** + * Shared poll-then-admit. `decrementOnPoll` is true for the recycle path + * (the released slot's inFlight decrement is deferred until the poll resolves, + * holding inFlight inflated so concurrent acquires queue — anti-overshoot) and + * false for the drain path (setLimit/pause-expiry/repoll — no slot to account). + * Admits at most ONE waiter on a successful poll. + */ + async function pollAndGrant( + providerId: string, + state: ProviderState, + decrementOnPoll = false, + ): Promise<void> { state.gatePolling = true; try { - const snapshot = await fetchUsage?.(providerId); + const snapshot = await safeFetchUsage(providerId); + + // For the recycle path, the released slot is now truly freed (the poll + // has observed the post-release upstream state). + if (decrementOnPoll) { + state.inFlight--; + } // Conditions may have changed during the async poll — re-check. if (state.paused) return; if (state.queue.length === 0) return; - if (state.inFlight >= state.limit) return; if (snapshot === undefined) { - // No usage info available → fall back to cooldown-only (grant). - grantLoop(state, providerId); + // No usage info available → fall back to cooldown-only (grant one). + grantOne(state, providerId); return; } if (snapshot.concurrentSessions < state.limit) { - // Upstream has room — admit exactly ONE queued waiter (grantLoop will - // stop after one because granting increments inFlight toward limit, and - // each admission pushes upstream back toward the limit). - grantLoop(state, providerId); + // Upstream has room — admit exactly ONE queued waiter. + grantOne(state, providerId); } // else: upstream at/over limit → keep queued; repoll timer handles retry. } finally { @@ -455,7 +541,19 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre } if (!state.paused && state.inFlight < state.limit) { - return Promise.resolve(grantSlot(state, providerId, conversationId)); + // Usage-gate anti-overshoot: while a recycle-poll is in flight, the + // inFlight count is momentarily unreliable (a released slot's decrement + // is deferred until the poll resolves — see pollAndGrant). A concurrent + // caller that fast-pathed now could overshoot the upstream limit before + // the poll confirms room. So route it through the queue instead; the + // in-flight poll will re-check (gateRepollRequested) and admit it once + // upstream confirms room. When no poll is in flight the fast-path is + // safe (the cooldown keeps inFlight inflated through the lag window). + if (fetchUsage !== undefined && state.gatePolling) { + // falls through to the queue path below + } else { + return Promise.resolve(grantSlot(state, providerId, conversationId)); + } } // Cannot grant immediately — the request will be queued. @@ -528,6 +626,29 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre void tryGrantNext(providerId); }, + restoreLimit(providerId, limit, autoReducedFrom) { + // Startup restoration (Bug 3): seed state from disk WITHOUT the manual + // "user took control" semantics, so a persisted auto-reduced limit keeps + // its notice/banner across a restart. When autoReducedFrom is provided, + // re-mark the state as auto-reduced (rebuild the notice). + let state = states.get(providerId); + if (state === undefined) { + state = makeState(limit, seedCooldown(providerId)); + states.set(providerId, state); + } else { + state.limit = limit; + } + if (autoReducedFrom !== undefined && autoReducedFrom > limit) { + state.autoReduced = true; + state.autoReducedFrom = autoReducedFrom; + state.notice = + `Concurrency limit auto-reduced to ${limit} after a 429 — ` + + "restore manually when ready."; + } + // A higher limit may let queued requests through. + void tryGrantNext(providerId); + }, + getLimit(providerId) { return states.get(providerId)?.limit; }, diff --git a/packages/provider-concurrency/src/extension.ts b/packages/provider-concurrency/src/extension.ts index 261d531..378655b 100644 --- a/packages/provider-concurrency/src/extension.ts +++ b/packages/provider-concurrency/src/extension.ts @@ -41,14 +41,23 @@ const WATCHDOG_INTERVAL_MS = 30 * 1000; const DEFAULT_PAUSE_MS = 30 * 1000; const RELEASE_COOLDOWN_MS = 350; -/** Storage key prefix for persisted cooldowns (limits are stored under the bare providerId). */ +/** + * Storage key prefixes. Limits are stored under the bare `<providerId>` key + * (unchanged for backward compatibility). Cooldowns + the adaptive-headroom + * auto-reduce marker are stored under their own prefixed keys so they persist + * independently without loadLimits misreading them as limits. + */ const COOLDOWN_KEY_PREFIX = "cooldown:"; +const AUTOREDUCE_KEY_PREFIX = "auto-reduce:"; /** * Wrap a `ConcurrencyService` so `setLimit`/`removeLimit`/`setCooldown` persist * to the given `StorageNamespace`. All other methods delegate directly to the * inner service. Persistence is fire-and-forget — a storage write failure logs * a warning but does NOT fail the API call (the in-memory value is already set). + * + * `restoreLimit` is NOT persisted here — it is a startup restore FROM disk, so + * it delegates straight through (the value is already on disk). */ function createPersistedService( inner: ConcurrencyService, @@ -66,7 +75,13 @@ function createPersistedService( err: err instanceof Error ? err.message : String(err), }), ); + // A MANUAL limit set clears the auto-reduce notice (the user took + // control) → drop the persisted auto-reduce marker too. + storage.delete(`${AUTOREDUCE_KEY_PREFIX}${providerId}`).catch(() => { + /* absent marker is fine */ + }); }, + restoreLimit: inner.restoreLimit.bind(inner), removeLimit(providerId) { inner.removeLimit(providerId); storage.delete(providerId).catch((err) => @@ -75,6 +90,9 @@ function createPersistedService( err: err instanceof Error ? err.message : String(err), }), ); + storage.delete(`${AUTOREDUCE_KEY_PREFIX}${providerId}`).catch(() => { + /* absent marker is fine */ + }); }, setCooldown(providerId, cooldownMs) { inner.setCooldown(providerId, cooldownMs); @@ -96,9 +114,14 @@ function createPersistedService( } /** - * Load saved limits from storage and apply them to the manager. - * Called during activate, before the service is registered. Skips cooldown - * keys (prefixed `cooldown:`) — those are loaded by {@link loadCooldowns}. + * Load saved limits from storage and apply them to the manager via + * `restoreLimit` (NOT `setLimit` — Bug 3). `setLimit` is a MANUAL user action + * that clears the auto-reduce notice; using it at startup would wipe the + * persisted auto-reduce banner. `restoreLimit` seeds the limit WITHOUT clearing + * the notice, and `loadAutoReduce` re-applies the notice afterward. + * + * Skips prefixed keys (cooldown:/auto-reduce:) — those are loaded by their + * own loaders. */ async function loadLimits( storage: StorageNamespace, @@ -108,18 +131,51 @@ async function loadLimits( const keys = await storage.keys(); for (const key of keys) { if (key.startsWith(COOLDOWN_KEY_PREFIX)) continue; // cooldown settings + if (key.startsWith(AUTOREDUCE_KEY_PREFIX)) continue; // auto-reduce markers const providerId = key; const raw = await storage.get(providerId); if (raw === null) continue; const limit = Number.parseInt(raw, 10); if (!Number.isNaN(limit) && limit > 0) { - manager.setLimit(providerId, limit); + manager.restoreLimit(providerId, limit); logger.info(`provider-concurrency: restored limit ${limit} for "${providerId}"`); } } } /** + * Load saved auto-reduce markers and re-apply them via `restoreLimit` so the + * frontend banner survives a restart (Bug 3). A marker is stored under + * `auto-reduce:<providerId>` with the value = the ORIGINAL limit before + * reduction (autoReducedFrom). The current (reduced) limit was already restored + * by {@link loadLimits}; this call re-marks it as auto-reduced. + */ +async function loadAutoReduce( + storage: StorageNamespace, + manager: ConcurrencyService, + logger: Logger, +): Promise<void> { + const keys = await storage.keys(AUTOREDUCE_KEY_PREFIX); + for (const key of keys) { + const providerId = key.slice(AUTOREDUCE_KEY_PREFIX.length); + if (providerId.length === 0) continue; + const raw = await storage.get(key); + if (raw === null) continue; + const autoReducedFrom = Number.parseInt(raw, 10); + if (!Number.isNaN(autoReducedFrom) && autoReducedFrom > 0) { + const currentLimit = manager.getLimit(providerId); + if (currentLimit !== undefined && currentLimit < autoReducedFrom) { + manager.restoreLimit(providerId, currentLimit, autoReducedFrom); + logger.info( + `provider-concurrency: restored auto-reduce notice for "${providerId}" ` + + `(${autoReducedFrom} -> ${currentLimit})`, + ); + } + } + } +} + +/** * Load saved cooldowns from storage and apply them to the manager. * Cooldowns are stored under `cooldown:<providerId>` keys (distinct from the * bare-`<providerId>` limit keys) so the two settings persist independently. @@ -185,21 +241,37 @@ export async function activate(host: HostAPI): Promise<void> { oldLimit, newLimit, }); - // Persist the reduced (one-way) limit so it survives a restart. + // Persist the reduced (one-way) limit so it survives a restart, AND the + // auto-reduce marker (autoReducedFrom) so the banner survives too (Bug 3). storage.set(providerId, String(newLimit)).catch((err) => logger.warn("provider-concurrency: failed to persist auto-reduced limit", { providerId, err: err instanceof Error ? err.message : String(err), }), ); + storage.set(`${AUTOREDUCE_KEY_PREFIX}${providerId}`, String(oldLimit)).catch((err) => + logger.warn("provider-concurrency: failed to persist auto-reduce marker", { + providerId, + err: err instanceof Error ? err.message : String(err), + }), + ); + }, + onUsagePollError: (providerId, err) => { + // A throwing getUsage() is treated as "no usage info" (cooldown-only + // fallback) by the manager — this is WARN-level observability only (Bug 2). + logger.warn("provider-concurrency: usage poll failed — falling back to cooldown-only", { + providerId, + err: err instanceof Error ? err.message : String(err), + }); }, }; const inner = createConcurrencyManager(managerOpts); - // Restore persisted limits + cooldowns before registering the service so the - // first request sees the correct configuration. + // Restore persisted limits + auto-reduce notices + cooldowns before registering + // the service so the first request sees the correct configuration. await loadLimits(storage, inner, logger); + await loadAutoReduce(storage, inner, logger); await loadCooldowns(storage, inner, logger); const service = createPersistedService(inner, storage, logger); |
