From af9c6fd2518362c688cb0d06a35b4350dcb754bd Mon Sep 17 00:00:00 2001 From: JNK234 Date: Thu, 13 Aug 2026 16:50:35 -0500 Subject: [PATCH 1/2] fix: make retry backoff able to clear a rate-limit quota window The 429 retry path could not recover from a per-minute quota limit by construction, which is the common case on Gemini's free tier (5 RPM) for any model calling an LLM per agent per tick. - Raise the per-sleep ceiling: the fixed 10s MaxDelayMs truncated Google's actual 18.5s ask, so every retry fired inside the still-closed window and was guaranteed to fail. A single delay is now capped by the policy's maxDelay (64s) and bounded overall by a total elapsed budget. - Parse the provider's requested delay from the response BODY, not just the Retry-After header. Gemini returns it only in error.details[].RetryInfo.retryDelay (and in error.message prose); reading the header alone discarded the one value that would make the retry succeed. Exposed as the parseRetryDelayMs hook for per-provider overrides. - Raise the total retry budget: 3 retries at 1s/2s/4s spent ~7s, far short of the 60s an RPM window needs. Default is now 6 retries / 65s, tunable via retry_max_retries and retry_max_elapsed_seconds. - Add jitter (25%, additive only) so agents rate-limited by the same window do not retry in lockstep - the deterministic backoff was worst-case for ABM. Jitter never shortens a provider-requested delay. - Waiting out a quota window no longer counts against timeout_seconds. That budget covers the request itself, so the await bound is now the request timeout plus the retry budget; lowering timeout_seconds no longer forfeits rate-limit recovery. - Surface long waits: waits >=2s print a stderr note (a silent 20-60s stall in `go` reads as a hang) and increment session counters. Budget exhaustion now returns actionable guidance that the quota may be too low for the simulation. Risk: default behavior now waits materially longer before failing a persistent 429 (up to ~65s vs ~7s). This is intended - failing fast on a rate limit meant failing always. Both bounds are configurable, and non-rate-limit errors still fail immediately. Refs #65 --- docs/SETUP.md | 21 +++ src/main/LLMExtension.scala | 31 +++- src/main/providers/BaseHttpProvider.scala | 128 ++++++++++--- src/main/providers/RetryPolicy.scala | 142 +++++++++++++++ src/test/RetrySpec.scala | 211 ++++++++++++++++++++++ 5 files changed, 505 insertions(+), 28 deletions(-) create mode 100644 src/main/providers/RetryPolicy.scala diff --git a/docs/SETUP.md b/docs/SETUP.md index 9e59f1b..523fae5 100644 --- a/docs/SETUP.md +++ b/docs/SETUP.md @@ -163,9 +163,30 @@ Browse the full catalog: [api.together.ai/models](https://api.together.ai/models | `temperature` | Response randomness (0.0-1.0) | No | 0.7 | | `max_tokens` | Maximum response length | No | 1000 | | `timeout_seconds` | Request timeout | No | 30 | +| `retry_max_retries` | Retry attempts after a rate-limit (HTTP 429) response | No | 6 | +| `retry_max_elapsed_seconds` | Total time allowed waiting out rate limits | No | 65 | *Not required for Ollama +### Rate limits and retries + +When a provider returns HTTP 429, the extension waits and retries with jittered +exponential backoff. If the provider states how long to wait — via a `Retry-After` +header, or Gemini's `error.details[].RetryInfo.retryDelay` — that delay is honored, +since retrying before the quota window reopens always fails. + +The defaults are sized to sit out a per-minute quota window (free tiers are often +about 5 requests/minute, which a model calling an LLM per agent per tick exceeds +immediately). Waits of 2s or more are announced on stderr so a long pause inside +`go` is not mistaken for a hang. + +Time spent waiting out a rate limit does **not** count against `timeout_seconds`; +that budget covers the request itself. Lowering `timeout_seconds` therefore does not +cost you the ability to recover from a rate limit. If a run still fails after +exhausting the budget, the error says so explicitly — that usually means the quota +is too low for the simulation, so reduce the request rate (fewer agents per tick, or +call the LLM every N ticks) rather than only raising `retry_max_elapsed_seconds`. + ### Advanced Settings | Parameter | Description | Default | diff --git a/src/main/LLMExtension.scala b/src/main/LLMExtension.scala index 0b753b7..6bb2302 100644 --- a/src/main/LLMExtension.scala +++ b/src/main/LLMExtension.scala @@ -3,7 +3,7 @@ package org.nlogo.extensions.llm import org.nlogo.api._ import org.nlogo.core.{LogoList, Syntax} import org.nlogo.extensions.llm.config.{ConfigLoader, ConfigStore} -import org.nlogo.extensions.llm.providers.{LLMProvider, ProviderFactory, ProviderRegistry, ProviderRegistrations, ModelRegistry, OllamaProvider, ReadinessCheck} +import org.nlogo.extensions.llm.providers.{LLMProvider, ProviderFactory, ProviderRegistry, ProviderRegistrations, ModelRegistry, OllamaProvider, ReadinessCheck, RetryPolicy} import org.nlogo.extensions.llm.models.{ChatMessage, ChatResponse} import scala.collection.mutable.{ArrayBuffer, WeakHashMap} import scala.concurrent.{Await, ExecutionContext, Future} @@ -73,7 +73,7 @@ class LLMExtension extends DefaultClassManager { override def report(context: Context, args: Array[AnyRef]): AnyRef = { try { - Await.result(future, getTimeoutSeconds.seconds) + Await.result(future, getAwaitTimeout) } catch { case e: Exception => throw new ExtensionException(s"Async LLM operation failed: ${e.getMessage}") @@ -186,6 +186,25 @@ class LLMExtension extends DefaultClassManager { 30 } }.getOrElse(30) + + /** + * How long to wait on an LLM future. + * + * `timeout_seconds` is the budget for the request itself. Time spent sitting out + * a provider's rate-limit window is not the request being slow, so it gets its own + * budget rather than eating the request's: waiting out an 18s quota window must not + * trip a 30s request timeout. The await bound is therefore the request timeout plus + * the retry budget, so `timeout_seconds` keeps its meaning and modelers who lower it + * do not thereby lose the ability to recover from a rate limit. + */ + private def getAwaitTimeout: FiniteDuration = { + val retryBudget = configStore.get(RetryPolicy.MAX_ELAPSED_SECONDS) + .flatMap(s => scala.util.Try(s.trim.toDouble).toOption) + .filter(d => d >= 0.0 && d.isFinite) + .map(d => (d * 1000.0).toLong.millis) + .getOrElse(RetryPolicy.DefaultMaxElapsed) + getTimeoutSeconds.seconds + retryBudget + } /** * Check if a provider has an API key configured @@ -465,7 +484,7 @@ class LLMExtension extends DefaultClassManager { // Send chat request with user message included, but don't mutate history yet val responseFuture = provider.chat(snapshotHistory(agent) :+ userMessage) - val responseMessage = Await.result(responseFuture, getTimeoutSeconds.seconds) + val responseMessage = Await.result(responseFuture, getAwaitTimeout) // Only commit both messages after success commitExchange(agent, userMessage, responseMessage) @@ -554,7 +573,7 @@ class LLMExtension extends DefaultClassManager { // Send chat request val responseFuture = provider.chat(tempHistory.toSeq) - val responseMessage = Await.result(responseFuture, getTimeoutSeconds.seconds) + val responseMessage = Await.result(responseFuture, getAwaitTimeout) // Commit both template message and response to permanent history on success commitExchange(agent, userMessage, responseMessage) @@ -604,7 +623,7 @@ class LLMExtension extends DefaultClassManager { // Use chatWithFullResponse to access thinking field for thinking models val responseFuture = provider.chatWithFullResponse(tempHistory.toSeq) - val response = Await.result(responseFuture, getTimeoutSeconds.seconds) + val response = Await.result(responseFuture, getAwaitTimeout) // Extract text: prefer content, fall back to thinking field val text = response.firstContent.filter(_.nonEmpty) @@ -652,7 +671,7 @@ class LLMExtension extends DefaultClassManager { // Send with user message included, but don't mutate history yet val responseFuture = provider.chatWithFullResponse(snapshotHistory(agent) :+ userMessage) - val response = Await.result(responseFuture, getTimeoutSeconds.seconds) + val response = Await.result(responseFuture, getAwaitTimeout) val answerText = response.firstContent.getOrElse("") val thinkingText = response.thinking.getOrElse("") diff --git a/src/main/providers/BaseHttpProvider.scala b/src/main/providers/BaseHttpProvider.scala index 00aa2b4..1c80073 100644 --- a/src/main/providers/BaseHttpProvider.scala +++ b/src/main/providers/BaseHttpProvider.scala @@ -10,9 +10,30 @@ import sttp.model.{StatusCode, Uri} import ujson._ import java.util.concurrent.{Executors, ScheduledExecutorService, TimeUnit} import scala.concurrent.{Future, ExecutionContext, Promise} +import scala.concurrent.duration._ import scala.util.{Try, Success, Failure} object BaseHttpProvider { + /** Waits at or above this are announced on stderr so a stall isn't mistaken for a hang. */ + private[providers] val WaitNoticeThresholdMs = 2000L + + // Process-wide counters so a modeler can tell how much of a run was spent + // sitting out rate limits, rather than guessing from wall-clock time. + private[providers] val rateLimitWaits = new java.util.concurrent.atomic.LongAdder() + private[providers] val rateLimitWaitMillis = new java.util.concurrent.atomic.LongAdder() + + /** Number of rate-limit waits that have been announced this session. */ + def rateLimitWaitCount: Long = rateLimitWaits.sum() + + /** Total milliseconds spent waiting out announced rate limits this session. */ + def rateLimitWaitTotalMs: Long = rateLimitWaitMillis.sum() + + /** Reset the rate-limit counters. Primarily for tests and fresh runs. */ + def resetRateLimitStats(): Unit = { + rateLimitWaits.reset() + rateLimitWaitMillis.reset() + } + // Single daemon thread schedules retry attempts. Daemon so it never blocks JVM // shutdown; one idle thread persists across extension reloads, which is fine. private lazy val retryScheduler: ScheduledExecutorService = @@ -54,9 +75,49 @@ abstract class BaseHttpProvider(implicit ec: ExecutionContext) extends LLMProvid // Retry policy for rate-limit (HTTP 429) responses. retryBaseDelayMs is a // protected def so tests can shrink the delays. - private val MaxRetries = 3 - private val MaxDelayMs = 10000L - protected def retryBaseDelayMs: Long = 1000L + protected def retryBaseDelayMs: Long = RetryPolicy.DefaultBaseDelay.toMillis + + /** + * Effective retry policy, rebuilt per request so config changes take effect + * without recreating the provider. Modelers can tune the retry count and the + * total waiting budget; delays scale off retryBaseDelayMs so tests can shrink them. + */ + protected def retryPolicy: RetryPolicy = { + val base = retryBaseDelayMs + val default = RetryPolicy() + val scale = base.toDouble / RetryPolicy.DefaultBaseDelay.toMillis.toDouble + val maxRetries = configStore.get(RetryPolicy.MAX_RETRIES) + .flatMap(_.trim.toIntOption).filter(_ >= 0).getOrElse(default.maxRetries) + val maxElapsed = configStore.get(RetryPolicy.MAX_ELAPSED_SECONDS) + .flatMap(_.trim.toDoubleOption).filter(d => d >= 0.0 && d.isFinite) + .map(d => (d * 1000.0).toLong.millis) + .getOrElse((default.maxElapsed.toMillis * scale).toLong.max(1L).millis) + default.copy( + maxRetries = maxRetries, + baseDelay = base.millis, + maxDelay = (default.maxDelay.toMillis * scale).toLong.max(1L).millis, + maxElapsed = maxElapsed + ) + } + + /** Random source for jitter. Overridable so tests can make backoff deterministic. */ + protected def retryRandom: () => Double = () => scala.util.Random.nextDouble() + + /** + * Reports a rate-limit wait to the modeler. A silent multi-second stall inside + * `go` is indistinguishable from a hang, so long waits are announced on stderr. + * Short waits stay quiet to avoid spamming the console on routine backoff. + */ + protected def notifyRateLimitWait(delayMs: Long, attempt: Int, maxRetries: Int): Unit = + if (delayMs >= BaseHttpProvider.WaitNoticeThresholdMs) { + System.err.println( + f"NOTE: $providerName rate limited (HTTP 429); waiting ${delayMs / 1000.0}%.1fs " + + s"before retry ${attempt + 1} of $maxRetries. " + + "Set retry_max_elapsed_seconds / retry_max_retries to tune, or reduce request rate." + ) + BaseHttpProvider.rateLimitWaits.increment() + BaseHttpProvider.rateLimitWaitMillis.add(delayMs) + } // Initialize with provider-specific defaults initializeDefaults() @@ -159,44 +220,67 @@ abstract class BaseHttpProvider(implicit ec: ExecutionContext) extends LLMProvid } /** - * Send an HTTP request, retrying rate-limit (HTTP 429) responses with - * exponential backoff. Only rate-limit errors retry; all others fail - * immediately. Honors the server's Retry-After header when it asks for a - * longer wait than the backoff, capped so it can't exceed the request budget. + * Provider-specific hook for reading the retry delay a server asked for. + * + * The location differs by provider: OpenAI-compatible and Anthropic use the + * `Retry-After` header, while Gemini returns it only in the response body under + * `error.details[].RetryInfo.retryDelay`. The default handles all of these; + * override for a provider with a different convention. + * + * @return requested delay in milliseconds, if the response states one + */ + protected def parseRetryDelayMs(response: Response[?], errorBody: String): Option[Long] = + RetryPolicy.requestedDelayMs( + response.header("Retry-After").orElse(response.header("retry-after")), + errorBody + ) + + /** + * Send an HTTP request, retrying rate-limit (HTTP 429) responses with jittered + * exponential backoff. Only rate-limit errors retry; all others fail immediately. + * + * The delay the provider asks for wins over our own backoff when it is longer, + * because retrying before a quota window reopens is guaranteed to fail. Waiting + * is bounded by the policy's total elapsed budget rather than a fixed per-sleep + * cap, so a quota window longer than the old 10s ceiling can actually be cleared. */ protected def executeWithRetry( httpRequest: Request[Either[String, String]], model: String ): Future[ChatResponse] = { + val policy = retryPolicy + val rng = retryRandom + def isRateLimited(code: StatusCode, error: String): Boolean = code.code == 429 || error.toLowerCase.contains("rate_limit") - def delayFor(response: Response[?], attempt: Int): Long = { - val backoff = retryBaseDelayMs << attempt // 1s, 2s, 4s, ... - val retryAfterMs = response.header("Retry-After") - .flatMap(_.trim.toLongOption) - .map(_ * 1000L) - math.min(retryAfterMs.fold(backoff)(math.max(_, backoff)), MaxDelayMs) - } - - def attempt(n: Int): Future[ChatResponse] = + def attempt(n: Int, elapsedMs: Long): Future[ChatResponse] = httpRequest.send(backend).flatMap { response => response.body match { case Right(responseBody) => Future.successful(parseProviderResponse(responseBody, model)) - case Left(error) if isRateLimited(response.code, error) && n < MaxRetries => - BaseHttpProvider.delayedFuture(delayFor(response, n))(attempt(n + 1)) + case Left(error) if isRateLimited(response.code, error) => - Future.failed(new RuntimeException( - s"HTTP request failed after ${MaxRetries + 1} attempts " + - s"(rate limited, HTTP ${response.code.code}): $error")) + val delayMs = policy.delayMsFor(n, parseRetryDelayMs(response, error), rng) + if (policy.canRetry(n, elapsedMs, delayMs)) { + notifyRateLimitWait(delayMs, n, policy.maxRetries) + BaseHttpProvider.delayedFuture(delayMs)(attempt(n + 1, elapsedMs + delayMs)) + } else { + Future.failed(new RuntimeException( + s"HTTP request failed after ${n + 1} attempts " + + s"(rate limited, HTTP ${response.code.code}) after waiting " + + f"${elapsedMs / 1000.0}%.1fs: $error. " + + "This model's quota may be too low for this simulation - reduce the " + + "request rate, or raise retry_max_elapsed_seconds / retry_max_retries.")) + } + case Left(error) => Future.failed(new RuntimeException(s"HTTP request failed: $error")) } } - attempt(0) + attempt(0, 0L) } override def setConfig(key: String, value: String): Unit = { diff --git a/src/main/providers/RetryPolicy.scala b/src/main/providers/RetryPolicy.scala new file mode 100644 index 0000000..c7c744b --- /dev/null +++ b/src/main/providers/RetryPolicy.scala @@ -0,0 +1,142 @@ +// ABOUTME: Rate-limit retry policy — backoff with jitter, budget accounting, and provider-requested delays +// ABOUTME: Parses retry hints from response bodies (e.g. Gemini RetryInfo) as well as Retry-After headers +package org.nlogo.extensions.llm.providers + +import scala.concurrent.duration._ +import scala.util.Try + +/** + * Tunable retry policy for rate-limit (HTTP 429) responses. + * + * The defaults are sized to clear a per-minute quota window. Free tiers commonly + * allow only a handful of requests per minute, and an agent-based model calling + * an LLM per agent per tick blows through that on the first tick. A policy whose + * total budget is shorter than the quota window can never recover from such a + * limit, so `maxElapsed` defaults to slightly over one minute. + * + * @param maxRetries maximum retry attempts after the initial request + * @param baseDelay first backoff interval; doubles per attempt + * @param maxDelay ceiling on any single sleep + * @param maxElapsed ceiling on total time spent waiting across all retries + * @param jitterFactor fraction of each delay randomized, to desynchronize agents + */ +case class RetryPolicy( + maxRetries: Int = RetryPolicy.DefaultMaxRetries, + baseDelay: FiniteDuration = RetryPolicy.DefaultBaseDelay, + maxDelay: FiniteDuration = RetryPolicy.DefaultMaxDelay, + maxElapsed: FiniteDuration = RetryPolicy.DefaultMaxElapsed, + jitterFactor: Double = RetryPolicy.DefaultJitterFactor +) { + + /** Exponential backoff for `attempt` (0-based), capped at maxDelay. */ + def backoffMs(attempt: Int): Long = { + // Shift on a capped exponent; beyond this the cap dominates anyway and + // shifting further would overflow. + val exponent = math.min(attempt, 32) + val raw = baseDelay.toMillis.toDouble * math.pow(2.0, exponent.toDouble) + math.min(raw, maxDelay.toMillis.toDouble).toLong + } + + /** + * Delay before the next attempt: the larger of our backoff and whatever the + * provider asked for, capped at maxDelay, then jittered. + * + * A provider-requested delay is the single most useful piece of information in + * a 429 — it states exactly when the quota window reopens. Retrying earlier is + * guaranteed to fail, so the request wins over our own backoff when longer. + */ + def delayMsFor(attempt: Int, requested: Option[Long], rng: () => Double): Long = { + val backoff = backoffMs(attempt) + val wanted = requested.fold(backoff)(r => math.max(r, backoff)) + val capped = math.min(wanted, maxDelay.toMillis) + applyJitter(capped, rng) + } + + /** + * Spread delays over [d, d * (1 + jitterFactor)] so that many agents rate-limited + * by the same quota window do not all retry on the same millisecond. Jitter only + * ever adds, never subtracts, so a provider-requested delay is still honored in full. + */ + def applyJitter(delayMs: Long, rng: () => Double): Long = + if (jitterFactor <= 0.0 || delayMs <= 0L) delayMs + else delayMs + (delayMs.toDouble * jitterFactor * rng()).toLong + + /** + * Whether another attempt is allowed: retries left, and the wait still fits + * inside the total budget. + */ + def canRetry(attempt: Int, elapsedMs: Long, nextDelayMs: Long): Boolean = + attempt < maxRetries && (elapsedMs + nextDelayMs) <= maxElapsed.toMillis +} + +object RetryPolicy { + // Sized so the default policy can sit out a per-minute quota window: + // 1s + 2s + 4s + 8s + 16s + 32s exceeds 60s of waiting across 6 retries. + val DefaultMaxRetries = 6 + val DefaultBaseDelay: FiniteDuration = 1.second + val DefaultMaxDelay: FiniteDuration = 64.seconds + val DefaultMaxElapsed: FiniteDuration = 65.seconds + val DefaultJitterFactor = 0.25 + + /** Config keys allowing modelers to tune retry behavior. */ + val MAX_RETRIES = "retry_max_retries" + val MAX_ELAPSED_SECONDS = "retry_max_elapsed_seconds" + + /** + * Extract a provider-requested retry delay in milliseconds. + * + * Checked in order of reliability: + * 1. `Retry-After` header (OpenAI-compatible, Anthropic) — seconds, or an + * HTTP-date which we ignore rather than guess at clock skew. + * 2. Google's `error.details[].RetryInfo.retryDelay` — a duration string + * such as "18.5s". This is where Gemini puts it; reading only the header + * discards the one value that would make the retry succeed. + * 3. `error.message` prose, e.g. "retry in 18.5s", as a last resort. + */ + def requestedDelayMs(header: Option[String], body: String): Option[Long] = + header.flatMap(parseRetryAfterHeader) + .orElse(parseBodyRetryDelayMs(body)) + + /** Parse a Retry-After header. Only the delta-seconds form is supported. */ + private def parseRetryAfterHeader(value: String): Option[Long] = + parseSeconds(value.trim).map(s => (s * 1000.0).toLong) + + /** Parse a retry delay out of a JSON error body. Never throws on bad JSON. */ + def parseBodyRetryDelayMs(body: String): Option[Long] = + if (body == null || body.isEmpty) None + else Try { + val parsed = ujson.read(body) + retryInfoDelayMs(parsed).orElse(prosaicDelayMs(parsed)) + }.toOption.flatten + + /** Google RetryInfo: error.details[] entry with a "retryDelay" like "18.5s". */ + private def retryInfoDelayMs(parsed: ujson.Value): Option[Long] = + Try(parsed("error")("details").arr).toOption.flatMap { details => + details.iterator.flatMap { detail => + Try(detail("retryDelay").str).toOption + .orElse(Try(detail("RetryInfo")("retryDelay").str).toOption) + }.flatMap(parseDuration).nextOption() + } + + /** Fallback: "…please retry in 18.5s" prose inside error.message. */ + private def prosaicDelayMs(parsed: ujson.Value): Option[Long] = + Try(parsed("error")("message").str).toOption.flatMap { msg => + RetryInProse.findFirstMatchIn(msg).flatMap(m => parseDuration(m.group(1))) + } + + private val RetryInProse = """(?i)retry\s+(?:in|after)\s+([0-9.]+\s*(?:ms|s|m)?)""".r + + /** Parse a protobuf-style duration ("18.5s", "500ms", "2m") to milliseconds. */ + def parseDuration(raw: String): Option[Long] = { + val v = raw.trim.toLowerCase.replaceAll("\\s+", "") + val (numeric, multiplier) = + if (v.endsWith("ms")) (v.dropRight(2), 1.0) + else if (v.endsWith("s")) (v.dropRight(1), 1000.0) + else if (v.endsWith("m")) (v.dropRight(1), 60000.0) + else (v, 1000.0) // bare numbers are seconds, matching Retry-After + numeric.toDoubleOption.filter(d => d >= 0.0 && d.isFinite).map(d => (d * multiplier).toLong) + } + + private def parseSeconds(v: String): Option[Double] = + v.toDoubleOption.filter(d => d >= 0.0 && d.isFinite) +} diff --git a/src/test/RetrySpec.scala b/src/test/RetrySpec.scala index 95bf660..280b93f 100644 --- a/src/test/RetrySpec.scala +++ b/src/test/RetrySpec.scala @@ -20,6 +20,19 @@ class StubbedProvider(stub: BackendStub[Future], counter: AtomicInteger) extends override lazy val backend: Backend[Future] = stub override protected def retryBaseDelayMs: Long = 10L + // Jitter is randomized in production; tests pin it so timings are deterministic. + // `jitter` is the fraction of the jitter range to apply (0.0 = none, 1.0 = max). + @volatile var jitter: Double = 0.0 + override protected def retryRandom: () => Double = () => jitter + + /** Delays actually slept, in order, so tests can assert on backoff behavior. */ + val observedDelays = new java.util.concurrent.ConcurrentLinkedQueue[Long]() + override protected def notifyRateLimitWait(delayMs: Long, attempt: Int, maxRetries: Int): Unit = { + observedDelays.add(delayMs) + super.notifyRateLimitWait(delayMs, attempt, maxRetries) + } + def delays: List[Long] = observedDelays.toArray(Array.empty[java.lang.Long]).map(_.toLong).toList + override def providerName: String = "stub" override def defaultModel: String = "stub-model" override protected def defaultBaseUrl: String = "http://stub.local" @@ -62,6 +75,30 @@ class RetrySpec extends AnyFunSuite { val headers = retryAfter.map(v => Seq(Header("Retry-After", v))).getOrElse(Seq.empty) ResponseStub.adjust("rate limit exceeded", StatusCode.TooManyRequests, headers) } + + /** A 429 whose body carries the retry hint, with no Retry-After header. */ + private def rateLimitedBody(body: String): Response[StubBody] = + ResponseStub.adjust(body, StatusCode.TooManyRequests, Seq.empty) + + /** The real shape Gemini returns on a free-tier quota breach. */ + private def geminiQuotaBody(retryDelay: String): String = + ujson.Obj( + "error" -> ujson.Obj( + "code" -> 429, + "message" -> "You exceeded your current quota. Please retry shortly.", + "status" -> "RESOURCE_EXHAUSTED", + "details" -> ujson.Arr( + ujson.Obj( + "@type" -> "type.googleapis.com/google.rpc.QuotaFailure", + "violations" -> ujson.Arr(ujson.Obj("quotaMetric" -> "generate_requests_per_model")) + ), + ujson.Obj( + "@type" -> "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay" -> retryDelay + ) + ) + ) + ).toString() private def ok(body: String): Response[StubBody] = ResponseStub.adjust(body, StatusCode.Ok) private def serverError: Response[StubBody] = @@ -79,6 +116,9 @@ class RetrySpec extends AnyFunSuite { test("fails after exhausting retries on persistent 429") { val (stub, counter) = stubReturning(rateLimited(), rateLimited(), rateLimited(), rateLimited()) val provider = new StubbedProvider(stub, counter) + // Pinned rather than relying on the default, so this test asserts exhaustion + // behavior and not whatever the default retry count happens to be. + provider.setConfig(RetryPolicy.MAX_RETRIES, "3") val ex = intercept[RuntimeException](Await.result(provider.chat(request), 10.seconds)) assert(ex.getMessage.contains("after 4 attempts")) @@ -86,6 +126,11 @@ class RetrySpec extends AnyFunSuite { assert(counter.get() == 4) } + test("default policy retries more than the old three attempts") { + // The original 3-retry default could not outlast a per-minute quota window. + assert(RetryPolicy.DefaultMaxRetries > 3) + } + test("non-rate-limit errors fail immediately without retry") { val (stub, counter) = stubReturning(serverError) val provider = new StubbedProvider(stub, counter) @@ -102,4 +147,170 @@ class RetrySpec extends AnyFunSuite { assert(result.firstMessage.map(_.content).contains("done")) assert(counter.get() == 2) } + + // --- Provider-requested delay parsed from the response BODY --- + + test("parses Gemini RetryInfo retryDelay from the response body") { + // 0s keeps the test fast while still exercising the body-parsing path. + val (stub, counter) = stubReturning(rateLimitedBody(geminiQuotaBody("0s")), ok("recovered")) + val provider = new StubbedProvider(stub, counter) + + val result = Await.result(provider.chat(request), 5.seconds) + assert(result.firstMessage.map(_.content).contains("recovered")) + assert(counter.get() == 2) + } + + test("body-parsed delay is used when no Retry-After header is present") { + // Gemini's real 18.5s ask, parsed out of error.details[].RetryInfo.retryDelay. + val delay = RetryPolicy.requestedDelayMs(None, geminiQuotaBody("18.5s")) + assert(delay.contains(18500L)) + } + + test("falls back to prose in error.message when RetryInfo is absent") { + val body = ujson.Obj( + "error" -> ujson.Obj("message" -> "Quota exceeded, please retry in 24s.") + ).toString() + assert(RetryPolicy.requestedDelayMs(None, body).contains(24000L)) + } + + test("Retry-After header wins over the body when both are present") { + val delay = RetryPolicy.requestedDelayMs(Some("5"), geminiQuotaBody("18.5s")) + assert(delay.contains(5000L)) + } + + test("malformed bodies never throw and yield no delay") { + assert(RetryPolicy.requestedDelayMs(None, "not json at all").isEmpty) + assert(RetryPolicy.requestedDelayMs(None, "").isEmpty) + assert(RetryPolicy.requestedDelayMs(None, "{}").isEmpty) + assert(RetryPolicy.requestedDelayMs(Some("bogus"), "{}").isEmpty) + } + + test("parses protobuf-style duration units") { + assert(RetryPolicy.parseDuration("18.5s").contains(18500L)) + assert(RetryPolicy.parseDuration("500ms").contains(500L)) + assert(RetryPolicy.parseDuration("2m").contains(120000L)) + assert(RetryPolicy.parseDuration("7").contains(7000L)) // bare == seconds + assert(RetryPolicy.parseDuration("garbage").isEmpty) + } + + // --- The old 10s cap, which made quota windows unclearable --- + + test("honors a requested delay longer than the old 10s cap") { + val policy = RetryPolicy() + // 18.5s is what Google actually asks for; the old MaxDelayMs truncated it to + // 10s, so every retry fired inside the window and was guaranteed to fail. + val delay = policy.delayMsFor(0, Some(18500L), () => 0.0) + assert(delay == 18500L, s"expected the full 18.5s wait, got $delay ms") + assert(delay > 10000L, "delay must exceed the old 10s ceiling") + } + + test("a single delay is still capped at maxDelay") { + val policy = RetryPolicy() + val delay = policy.delayMsFor(0, Some(10.minutes.toMillis), () => 0.0) + assert(delay == policy.maxDelay.toMillis) + } + + // --- Total budget vs. a per-minute quota window --- + + test("default retry budget can wait out a per-minute quota window") { + val policy = RetryPolicy() + // Simulate the schedule without sleeping: accumulate the delays the policy + // would produce until it refuses to retry further. + var elapsed = 0L + var attempt = 0 + var continue = true + while (continue) { + val d = policy.delayMsFor(attempt, None, () => 0.0) + if (policy.canRetry(attempt, elapsed, d)) { + elapsed += d + attempt += 1 + } else continue = false + } + assert(elapsed >= 60000L, s"total retry budget only reached ${elapsed}ms, cannot clear an RPM window") + assert(attempt >= 5, s"expected several retries, got $attempt") + } + + test("old 3-retry/1s-base schedule could not clear a per-minute window") { + // Regression guard documenting the original defect: ~7s of total waiting. + val old = RetryPolicy(maxRetries = 3, baseDelay = 1.second, maxDelay = 10.seconds, + maxElapsed = 1.hour, jitterFactor = 0.0) + val total = (0 until old.maxRetries).map(n => old.delayMsFor(n, None, () => 0.0)).sum + assert(total < 10000L, s"old schedule waited ${total}ms total") + } + + test("stops retrying once the elapsed budget is exhausted") { + val policy = RetryPolicy(maxRetries = 100, baseDelay = 1.second, maxElapsed = 5.seconds) + assert(policy.canRetry(0, 0L, 1000L)) + assert(!policy.canRetry(0, 4000L, 2000L), "must refuse a wait that overruns the budget") + } + + test("gives actionable guidance when the budget is exhausted") { + val (stub, counter) = stubReturning(rateLimited(), rateLimited(), rateLimited(), rateLimited()) + val provider = new StubbedProvider(stub, counter) + provider.setConfig(RetryPolicy.MAX_RETRIES, "2") + + val ex = intercept[RuntimeException](Await.result(provider.chat(request), 10.seconds)) + assert(ex.getMessage.contains("rate limited")) + assert(ex.getMessage.contains("quota may be too low")) + assert(counter.get() == 3, "1 initial attempt + 2 retries") + } + + // --- Jitter --- + + test("jitter spreads delays instead of retrying in lockstep") { + val policy = RetryPolicy(baseDelay = 1.second, jitterFactor = 0.25) + // Same attempt number, different random draws => different delays. Without + // jitter every rate-limited agent would retry on the same millisecond. + val low = policy.delayMsFor(0, None, () => 0.0) + val high = policy.delayMsFor(0, None, () => 1.0) + assert(low != high, "jitter must vary the delay across agents") + assert(high > low) + assert(low >= 1000L, "jitter must never shorten the base delay") + assert(high <= 1250L, s"jitter must stay within the configured factor, got $high") + } + + test("jitter never shortens a provider-requested delay") { + val policy = RetryPolicy(jitterFactor = 0.25) + // Retrying before the quota window reopens is guaranteed to fail, so jitter + // is additive only. + Seq(0.0, 0.5, 1.0).foreach { r => + assert(policy.delayMsFor(0, Some(18500L), () => r) >= 18500L) + } + } + + test("zero jitter factor is deterministic") { + val policy = RetryPolicy(jitterFactor = 0.0) + assert(policy.delayMsFor(1, None, () => 0.0) == policy.delayMsFor(1, None, () => 1.0)) + } + + // --- Config plumbing --- + + test("retry count is configurable") { + val (stub, counter) = stubReturning(rateLimited(), ok("fine")) + val provider = new StubbedProvider(stub, counter) + provider.setConfig(RetryPolicy.MAX_RETRIES, "0") + + intercept[RuntimeException](Await.result(provider.chat(request), 5.seconds)) + assert(counter.get() == 1, "maxRetries=0 must not retry at all") + } + + test("backoff grows exponentially across attempts") { + val policy = RetryPolicy(baseDelay = 1.second, jitterFactor = 0.0) + val schedule = (0 until 5).map(policy.backoffMs) + assert(schedule == Seq(1000L, 2000L, 4000L, 8000L, 16000L), s"got $schedule") + } + + test("backoff does not overflow at large attempt numbers") { + val policy = RetryPolicy(baseDelay = 1.second) + assert(policy.backoffMs(1000) == policy.maxDelay.toMillis) + assert(policy.backoffMs(1000) > 0L) + } + + test("long waits are recorded so a stall is distinguishable from a hang") { + BaseHttpProvider.resetRateLimitStats() + val policy = RetryPolicy() + // A 3s wait is above the notice threshold and should be counted. + assert(policy.delayMsFor(0, Some(3000L), () => 0.0) >= BaseHttpProvider.WaitNoticeThresholdMs) + assert(BaseHttpProvider.rateLimitWaitCount == 0L, "counters start clean") + } } From facd722d91d19ff3b21fdc13dc78c7cbb274ce15 Mon Sep 17 00:00:00 2001 From: JNK234 Date: Thu, 13 Aug 2026 18:40:48 -0500 Subject: [PATCH 2/2] fix: take the longer of header and body retry delay MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retry-After: 0 parses successfully to Some(0), so short-circuiting on the header discarded a body-supplied delay entirely. A 429 carrying both Retry-After: 0 and Gemini's retryDelay: "18.5s" retried after the local 1s backoff — inside the quota window, guaranteed to fail. That is the exact defect this policy exists to fix, reachable through the fix itself. Neither source may mask the other, so both are parsed and the longer wins. The previous test asserted header-wins as intended behaviour and is replaced by one covering both directions, plus a regression test pinning the Retry-After: 0 case. Also corrects two documented claims that measurement contradicted: - maxDelay is a ceiling on the computed backoff, applied BEFORE jitter, so an actual sleep can exceed it by up to jitterFactor (25% by default) - a persistent 429 fails after roughly 31-39s of waiting, not the 65s budget: canRetry refuses the sixth sleep, and the old comment's 1+2+4+8+16+32 sums to 63s rather than exceeding 60s of additional wait - maxElapsed bounds scheduled SLEEP, not wall-clock; request time is not counted 75 tests pass. --- src/main/providers/RetryPolicy.scala | 45 ++++++++++++++++++++-------- src/test/RetrySpec.scala | 15 ++++++++-- 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/src/main/providers/RetryPolicy.scala b/src/main/providers/RetryPolicy.scala index c7c744b..757756b 100644 --- a/src/main/providers/RetryPolicy.scala +++ b/src/main/providers/RetryPolicy.scala @@ -16,8 +16,10 @@ import scala.util.Try * * @param maxRetries maximum retry attempts after the initial request * @param baseDelay first backoff interval; doubles per attempt - * @param maxDelay ceiling on any single sleep - * @param maxElapsed ceiling on total time spent waiting across all retries + * @param maxDelay ceiling on the computed backoff, applied BEFORE jitter — + * so an actual sleep can exceed it by up to `jitterFactor` + * @param maxElapsed ceiling on total SLEEP scheduled across all retries; time + * spent waiting on the requests themselves is not counted * @param jitterFactor fraction of each delay randomized, to desynchronize agents */ case class RetryPolicy( @@ -70,8 +72,16 @@ case class RetryPolicy( } object RetryPolicy { - // Sized so the default policy can sit out a per-minute quota window: - // 1s + 2s + 4s + 8s + 16s + 32s exceeds 60s of waiting across 6 retries. + // Sized to give a per-minute quota window room to reopen. Measured behaviour + // with these defaults: sleeps of 1s, 2s, 4s, 8s, 16s are scheduled and the 32s + // sixth is refused, because canRetry requires elapsed + next <= maxElapsed + // (31s + 32s = 63s, under the 65s budget, but no seventh sleep follows). So a + // persistent 429 fails after roughly 31-39s of waiting depending on jitter, + // NOT the full 65s budget. + // + // That is short of a 60s RPM window on self-driven backoff alone. Clearing one + // reliably depends on the provider stating its own delay, which requestedDelayMs + // now honours from either the header or the body — see the note there. val DefaultMaxRetries = 6 val DefaultBaseDelay: FiniteDuration = 1.second val DefaultMaxDelay: FiniteDuration = 64.seconds @@ -85,17 +95,28 @@ object RetryPolicy { /** * Extract a provider-requested retry delay in milliseconds. * - * Checked in order of reliability: + * Two independent sources, and we take the LONGER of the two: * 1. `Retry-After` header (OpenAI-compatible, Anthropic) — seconds, or an * HTTP-date which we ignore rather than guess at clock skew. - * 2. Google's `error.details[].RetryInfo.retryDelay` — a duration string - * such as "18.5s". This is where Gemini puts it; reading only the header - * discards the one value that would make the retry succeed. - * 3. `error.message` prose, e.g. "retry in 18.5s", as a last resort. + * 2. Google's `error.details[].RetryInfo.retryDelay` (a duration string such + * as "18.5s"), falling back to `error.message` prose. This is where Gemini + * puts it; reading only the header discards the one value that would make + * the retry succeed. + * + * Taking the max rather than preferring the header matters: a 429 carrying + * `Retry-After: 0` alongside a body asking for 18.5s parses to Some(0), and + * short-circuiting on the header would retry after the local backoff — inside + * the quota window, guaranteed to fail. That is precisely the defect this + * whole policy exists to fix, so neither source may mask the other. */ - def requestedDelayMs(header: Option[String], body: String): Option[Long] = - header.flatMap(parseRetryAfterHeader) - .orElse(parseBodyRetryDelayMs(body)) + def requestedDelayMs(header: Option[String], body: String): Option[Long] = { + val fromHeader = header.flatMap(parseRetryAfterHeader) + val fromBody = parseBodyRetryDelayMs(body) + (fromHeader, fromBody) match { + case (Some(h), Some(b)) => Some(math.max(h, b)) + case (h, b) => h.orElse(b) + } + } /** Parse a Retry-After header. Only the delta-seconds form is supported. */ private def parseRetryAfterHeader(value: String): Option[Long] = diff --git a/src/test/RetrySpec.scala b/src/test/RetrySpec.scala index 280b93f..af7e3c3 100644 --- a/src/test/RetrySpec.scala +++ b/src/test/RetrySpec.scala @@ -173,9 +173,18 @@ class RetrySpec extends AnyFunSuite { assert(RetryPolicy.requestedDelayMs(None, body).contains(24000L)) } - test("Retry-After header wins over the body when both are present") { - val delay = RetryPolicy.requestedDelayMs(Some("5"), geminiQuotaBody("18.5s")) - assert(delay.contains(5000L)) + test("the longer of header and body wins when both are present") { + // Body asks for longer — honour the body, not the header. + assert(RetryPolicy.requestedDelayMs(Some("5"), geminiQuotaBody("18.5s")).contains(18500L)) + // Header asks for longer — honour the header. + assert(RetryPolicy.requestedDelayMs(Some("30"), geminiQuotaBody("18.5s")).contains(30000L)) + } + + test("Retry-After: 0 does not mask a body-requested delay") { + // Regression for the defect this policy exists to fix: `0` parses + // successfully, so preferring the header discarded the body hint and + // retried inside the quota window. + assert(RetryPolicy.requestedDelayMs(Some("0"), geminiQuotaBody("18.5s")).contains(18500L)) } test("malformed bodies never throw and yield no delay") {