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..757756b --- /dev/null +++ b/src/main/providers/RetryPolicy.scala @@ -0,0 +1,163 @@ +// 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 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( + 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 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 + 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. + * + * 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"), 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] = { + 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] = + 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..af7e3c3 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,179 @@ 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("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") { + 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") + } }