fix(agents): bound probe throttle cache

altaywtf · altaywtf · commit bbc4254b9455 · 2026-03-10T00:50:55.000+03:00
diff --git a/src/agents/model-fallback.probe.test.ts b/src/agents/model-fallback.probe.test.ts
@@ -251,6 +251,36 @@ describe("runWithModelFallback – probe logic", () => {
     expectPrimaryProbeSuccess(result, run, "probed-ok");
   });
 
+  it("prunes stale probe throttle entries before checking eligibility", () => {
+    _probeThrottleInternals.lastProbeAttempt.set(
+      "stale",
+      NOW - _probeThrottleInternals.PROBE_STATE_TTL_MS - 1,
+    );
+    _probeThrottleInternals.lastProbeAttempt.set("fresh", NOW - 5_000);
+
+    expect(_probeThrottleInternals.lastProbeAttempt.has("stale")).toBe(true);
+
+    expect(_probeThrottleInternals.isProbeThrottleOpen(NOW, "fresh")).toBe(false);
+
+    expect(_probeThrottleInternals.lastProbeAttempt.has("stale")).toBe(false);
+    expect(_probeThrottleInternals.lastProbeAttempt.has("fresh")).toBe(true);
+  });
+
+  it("caps probe throttle state by evicting the oldest entries", () => {
+    for (let i = 0; i < _probeThrottleInternals.MAX_PROBE_KEYS; i += 1) {
+      _probeThrottleInternals.lastProbeAttempt.set(`key-${i}`, NOW - (i + 1));
+    }
+
+    _probeThrottleInternals.markProbeAttempt(NOW, "freshest");
+
+    expect(_probeThrottleInternals.lastProbeAttempt.size).toBe(
+      _probeThrottleInternals.MAX_PROBE_KEYS,
+    );
+    expect(_probeThrottleInternals.lastProbeAttempt.has("freshest")).toBe(true);
+    expect(_probeThrottleInternals.lastProbeAttempt.has("key-255")).toBe(false);
+    expect(_probeThrottleInternals.lastProbeAttempt.has("key-0")).toBe(true);
+  });
+
   it("handles non-finite soonest safely (treats as probe-worthy)", async () => {
     const cfg = makeCfg();
 
diff --git a/src/agents/model-fallback.ts b/src/agents/model-fallback.ts
@@ -342,17 +342,51 @@ const lastProbeAttempt = new Map<string, number>();
 const MIN_PROBE_INTERVAL_MS = 30_000; // 30 seconds between probes per key
 const PROBE_MARGIN_MS = 2 * 60 * 1000;
 const PROBE_SCOPE_DELIMITER = "::";
+const PROBE_STATE_TTL_MS = 24 * 60 * 60 * 1000;
+const MAX_PROBE_KEYS = 256;
 
 function resolveProbeThrottleKey(provider: string, agentDir?: string): string {
   const scope = String(agentDir ?? "").trim();
   return scope ? `${scope}${PROBE_SCOPE_DELIMITER}${provider}` : provider;
 }
 
+function pruneProbeState(now: number): void {
+  for (const [key, ts] of lastProbeAttempt) {
+    if (!Number.isFinite(ts) || ts <= 0 || now - ts > PROBE_STATE_TTL_MS) {
+      lastProbeAttempt.delete(key);
+    }
+  }
+}
+
+function enforceProbeStateCap(): void {
+  while (lastProbeAttempt.size > MAX_PROBE_KEYS) {
+    let oldestKey: string | null = null;
+    let oldestTs = Number.POSITIVE_INFINITY;
+    for (const [key, ts] of lastProbeAttempt) {
+      if (ts < oldestTs) {
+        oldestKey = key;
+        oldestTs = ts;
+      }
+    }
+    if (!oldestKey) {
+      break;
+    }
+    lastProbeAttempt.delete(oldestKey);
+  }
+}
+
 function isProbeThrottleOpen(now: number, throttleKey: string): boolean {
+  pruneProbeState(now);
   const lastProbe = lastProbeAttempt.get(throttleKey) ?? 0;
   return now - lastProbe >= MIN_PROBE_INTERVAL_MS;
 }
 
+function markProbeAttempt(now: number, throttleKey: string): void {
+  pruneProbeState(now);
+  lastProbeAttempt.set(throttleKey, now);
+  enforceProbeStateCap();
+}
+
 function shouldProbePrimaryDuringCooldown(params: {
   isPrimary: boolean;
   hasFallbackCandidates: boolean;
@@ -383,7 +417,12 @@ export const _probeThrottleInternals = {
   lastProbeAttempt,
   MIN_PROBE_INTERVAL_MS,
   PROBE_MARGIN_MS,
+  PROBE_STATE_TTL_MS,
+  MAX_PROBE_KEYS,
   resolveProbeThrottleKey,
+  isProbeThrottleOpen,
+  pruneProbeState,
+  markProbeAttempt,
 } as const;
 
 type CooldownDecision =
@@ -536,7 +575,7 @@ export async function runWithModelFallback<T>(params: {
         }
 
         if (decision.markProbe) {
-          lastProbeAttempt.set(probeThrottleKey, now);
+          markProbeAttempt(now, probeThrottleKey);
         }
         if (
           decision.reason === "rate_limit" ||
