summaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorAdam Malczewski <[email protected]>2026-06-28 14:50:27 +0900
committerAdam Malczewski <[email protected]>2026-06-28 14:50:27 +0900
commitfb4a9217b55dd3ba11670104ac23536416d36940 (patch)
tree4ae4aed5ca1bbc09620e0ab53e54cf54b5cbd2f2
parentdf7d45b58076b6ab47bea3148872cac11401d510 (diff)
downloaddispatch-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.
-rw-r--r--packages/provider-concurrency/src/concurrency-manager.test.ts168
-rw-r--r--packages/provider-concurrency/src/concurrency-manager.ts169
-rw-r--r--packages/provider-concurrency/src/extension.ts88
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);