oira666_pi-limits-wait 0.4.8 → 0.4.10

package/README.md

@@ -19,6 +19,8 @@ Restart Pi or run `/reload` after installing.
 - Uses provider retry timing when available (`retry-after`, `retry-after-ms`, `retry in ...`, reset messages).
 - If no retry timing is available for a rate limit, waits 30 minutes before retrying.
 - On `server_is_overloaded`, waits 5 minutes, then retries. If the provider is still overloaded after Pi's normal retries, it waits another 5 minutes and repeats.
+- On transient network/transport failures — including undici idle-timeout aborts (`UND_ERR_HEADERS_TIMEOUT` / `UND_ERR_BODY_TIMEOUT`), `fetch failed`, `terminated`, `ECONNRESET`, `ETIMEDOUT`, etc. — treats the error as retryable with a short backoff (default 15s, or the provider's `retry-after`) instead of giving up. This prevents a stalled streaming request from turning into a silent hang.
+- For any other (unclassified) non-retryable error, retries the same model a few times (default 3) before falling back / freezing it, so a one-off hiccup does not immediately sideline a model.
 - Shows a countdown in the Pi status/working line.
 - Press Enter during the countdown to skip the wait and retry immediately.
 - Optionally falls back to configured models when the current/default model is rate-limited.
@@ -74,6 +76,12 @@ Fallback behavior:
 
 When the settings are loaded, the extension shows the full usable fallback model list. When models become rate-limited, it shows a live countdown for each limited model.
 
+## Environment variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `PI_LIMITS_WAIT_FREEZING_ENABLED` | `true` | When a model keeps failing with a non-retryable error, it is normally "frozen" for 1 hour and skipped in favour of other configured models. Set this to `false` (also accepts `0`, `no`, `off`) to disable freezing entirely: the extension will instead try each configured candidate once (after the bounded retries) and then surface the error, never blocking on a long "model-frozen" wait. Useful for non-interactive / subagent runs where no one can press Enter to skip. |
+
 ## License
 
 MIT

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-limits-wait",
-  "version": "0.4.8",
+  "version": "0.4.10",
   "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
   "type": "module",
   "license": "MIT",

package/src/index.ts

@@ -31,15 +31,22 @@ const PI_IDENTITY_SENTENCE_PATTERN =
 
 const DEFAULT_RATE_LIMIT_WAIT_MS = 30 * 60 * 1_000; // 30 minutes
 const DEFAULT_OVERLOADED_WAIT_MS = 5 * 60 * 1_000; // 5 minutes
-const DEFAULT_NON_RETRYABLE_FREEZE_MS = 60 * 60 * 1_000; // 1 hour
+const DEFAULT_NETWORK_WAIT_MS = 15 * 1_000; // 15 seconds
+const DEFAULT_NON_RETRYABLE_FREEZE_MS = 10 * 60 * 1_000; // 1 hour
+/** How many times a non-retryable error is retried before we freeze/fallback. */
+const NON_RETRYABLE_MAX_ATTEMPTS = 3;
+/** Short backoff between non-retryable retries. */
+const NON_RETRYABLE_RETRY_DELAY_MS = 5 * 1_000; // 5 seconds
 const SETTINGS_FILE_NAME = "limits-wait.json";
 const FALLBACK_MODELS_KEY = "fallback-models";
+/** Env var to toggle the 1-hour model-freeze behaviour. Default: enabled. */
+const FREEZING_ENV_VAR = "PI_LIMITS_WAIT_FREEZING_ENABLED";
 
 /** Key used with ctx.ui.setStatus() for the countdown line. */
 const STATUS_KEY = "limits-wait";
 const MODELS_STATUS_KEY = "limits-wait-models";
 
-type RetryReason = "rate-limit" | "overloaded" | "authentication" | "model-frozen";
+type RetryReason = "rate-limit" | "overloaded" | "authentication" | "model-frozen" | "network" | "retry";
 
 type RetryableError = {
   reason: RetryReason;
@@ -78,10 +85,12 @@ let ambientStatusCleanup: (() => void) | undefined;
 let modelStatusCleanup: (() => void) | undefined;
 let activeProviderRequests = 0;
 let activeFallbackProviderRequests = 0;
-let lastObservedRateLimitError: { at: number; message: string } | undefined;
+let lastObservedHttpError: { at: number; message: string } | undefined;
 const wrappedApis = new Set<Api>();
 const builtinStreamSimpleByApi = new Map<Api, StreamSimpleFn>();
 let fallbackModels: FallbackModel[] = [];
+let nonRetryableMaxAttempts = NON_RETRYABLE_MAX_ATTEMPTS;
+let nonRetryableRetryDelayMs = NON_RETRYABLE_RETRY_DELAY_MS;
 let primaryModel: Model<Api> | undefined;
 let primaryThinkingLevel: ThinkingLevel | undefined;
 let suppressNextModelSelect = false;
@@ -346,11 +355,15 @@ function ensureRateLimitedModelsStatus(): void {
   };
 }
 
+function notifyRetryableError(model: Model<Api>, retryable: RetryableError, errorMessage?: string): void {
+  const detail = errorMessage ? ` Error: ${formatErrorDetail(errorMessage)}` : "";
+  sharedCtx?.ui.notify(`${formatModel(model)} ${reasonLabel(retryable.reason).toLowerCase()} for ${formatDuration(retryable.waitMs)}.${detail}`, "warning");
+}
+
 function rememberRateLimit(model: Model<Api>, retryable: RetryableError, errorMessage?: string): void {
   const deadline = Date.now() + retryable.waitMs;
   rateLimitMemory.set(modelKey(model), { reason: retryable.reason, limitedAt: Date.now(), deadline });
-  const detail = errorMessage ? ` Error: ${formatErrorDetail(errorMessage)}` : "";
-  sharedCtx?.ui.notify(`${formatModel(model)} ${reasonLabel(retryable.reason).toLowerCase()} for ${formatDuration(retryable.waitMs)}.${detail}`, "warning");
+  notifyRetryableError(model, retryable, errorMessage);
   ensureRateLimitedModelsStatus();
 }
 
@@ -423,10 +436,21 @@ function isAbortErrorEvent(event: AssistantMessageEvent, signal?: AbortSignal):
   return reason === "aborted" || stopReason === "aborted" || isAbortMessage(message);
 }
 
+function recentObservedHttpError(maxAgeMs = 60_000): string | undefined {
+  if (!lastObservedHttpError) return undefined;
+  if (Date.now() - lastObservedHttpError.at > maxAgeMs) return undefined;
+  return lastObservedHttpError.message;
+}
+
 function recentObservedRateLimitError(maxAgeMs = 60_000): string | undefined {
-  if (!lastObservedRateLimitError) return undefined;
-  if (Date.now() - lastObservedRateLimitError.at > maxAgeMs) return undefined;
-  return lastObservedRateLimitError.message;
+  const message = recentObservedHttpError(maxAgeMs);
+  return message && isRateLimitError(message) ? message : undefined;
+}
+
+function errorMessageWithRecentHttpStatus(errorMessage: string): string {
+  const observed = recentObservedHttpError();
+  if (!observed || /(?:^|\D)[1-5]\d\d(?:\D|$)/.test(errorMessage) || errorMessage.includes(observed)) return errorMessage;
+  return `${observed}; Error: ${errorMessage}`;
 }
 
 function earliestCandidateDeadline(current: Model<Api>): number | undefined {
@@ -481,7 +505,7 @@ async function switchPiModel(entry: FallbackModel): Promise<void> {
 
 // ─── Retryable error detection ────────────────────────────────────────────────
 
-export function isRateLimitError(msg: string): boolean {
+function isExplicitRateLimitError(msg: string): boolean {
   const lower = msg.toLowerCase();
   return (
     /(?:^|\D)429(?:\D|$)/.test(msg) ||
@@ -489,7 +513,14 @@ export function isRateLimitError(msg: string): boolean {
     /rate\s*limit/.test(lower) ||
     lower.includes("too many requests") ||
     lower.includes("quota exceeded") ||
-    lower.includes("quota will reset") ||
+    lower.includes("quota will reset")
+  );
+}
+
+export function isRateLimitError(msg: string): boolean {
+  const lower = msg.toLowerCase();
+  return (
+    isExplicitRateLimitError(msg) ||
     lower.includes("retry delay") ||
     lower.includes("retry-after")
   );
@@ -504,6 +535,40 @@ export function isServerOverloadedError(msg: string): boolean {
   );
 }
 
+/**
+ * Transient network / HTTP transport failures. These are NOT API rate-limit or
+ * overload responses, but they are still worth retrying with a short backoff
+ * instead of treating them as a hard non-retryable failure (which would freeze
+ * the model for an hour). Most importantly this covers undici's idle-timeout
+ * aborts (`UND_ERR_HEADERS_TIMEOUT` / `UND_ERR_BODY_TIMEOUT`) that fire when a
+ * streaming request stalls without a response.
+ */
+export function isTransientNetworkError(msg: string): boolean {
+  const lower = msg.toLowerCase();
+  return (
+    lower.includes("und_err_headers_timeout") ||
+    lower.includes("und_err_body_timeout") ||
+    lower.includes("und_err_connect_timeout") ||
+    lower.includes("und_err_socket") ||
+    lower.includes("headers timeout") ||
+    lower.includes("body timeout") ||
+    lower.includes("connect timeout") ||
+    lower.includes("fetch failed") ||
+    lower.includes("terminated") ||
+    lower.includes("socket hang up") ||
+    lower.includes("other side closed") ||
+    lower.includes("network_error") ||
+    lower.includes("network error") ||
+    lower.includes("econnreset") ||
+    lower.includes("econnrefused") ||
+    lower.includes("etimedout") ||
+    lower.includes("enetunreach") ||
+    lower.includes("ehostunreach") ||
+    lower.includes("eai_again") ||
+    lower.includes("enotfound")
+  );
+}
+
 export function isAuthenticationRefreshError(msg: string): boolean {
   const lower = msg.toLowerCase();
   return (
@@ -586,9 +651,15 @@ export function getRetryableError(msg: string): RetryableError | undefined {
   if (isAuthenticationRefreshError(msg)) {
     return { reason: "authentication", waitMs: parseRetryDelayMs(msg) ?? DEFAULT_RATE_LIMIT_WAIT_MS };
   }
+  if (isTransientNetworkError(msg) && !isExplicitRateLimitError(msg)) {
+    return { reason: "network", waitMs: parseRetryDelayMs(msg) ?? DEFAULT_NETWORK_WAIT_MS };
+  }
   if (isRateLimitError(msg)) {
     return { reason: "rate-limit", waitMs: rateLimitWaitMs(msg) };
   }
+  if (isTransientNetworkError(msg)) {
+    return { reason: "network", waitMs: parseRetryDelayMs(msg) ?? DEFAULT_NETWORK_WAIT_MS };
+  }
   return undefined;
 }
 
@@ -598,9 +669,23 @@ function reasonLabel(reason: RetryReason): string {
   if (reason === "overloaded") return "Server overloaded";
   if (reason === "authentication") return "Authentication refresh pending";
   if (reason === "model-frozen") return "Model frozen after error";
+  if (reason === "network") return "Network/timeout error";
+  if (reason === "retry") return "Retrying after error";
   return "Rate limited";
 }
 
+/**
+ * Whether the 1-hour model-freeze behaviour is enabled. Controlled by the
+ * `PI_LIMITS_WAIT_FREEZING_ENABLED` env var (default enabled). Set it to
+ * `false`/`0`/`no`/`off` to disable freezing entirely.
+ */
+export function freezingEnabled(): boolean {
+  const raw = process.env[FREEZING_ENV_VAR];
+  if (raw === undefined) return true;
+  const value = raw.trim().toLowerCase();
+  return !(value === "false" || value === "0" || value === "no" || value === "off");
+}
+
 function statusText(reason: RetryReason, deadline: number, allowSkip: boolean): string {
   const remaining = Math.max(0, deadline - Date.now());
   const totalSecs = Math.ceil(remaining / 1_000);
@@ -692,14 +777,19 @@ export function waitForRetry(
     };
 
     tick();
+    // The interval only drives the visible 1s countdown; the actual wait is
+    // resolved precisely at the deadline by `timer` below (so a 200ms backoff
+    // resolves in ~200ms instead of being rounded up to the next 1s tick).
     const ticker = setInterval(() => {
       if (Date.now() >= deadline) { cleanup(); resolve("waited"); return; }
       tick();
     }, 1_000);
+    const timer = setTimeout(() => { if (!done) { cleanup(); resolve("waited"); } }, Math.max(0, deadline - Date.now()));
 
     function cleanup() {
       done = true;
       clearInterval(ticker);
+      clearTimeout(timer);
       signal?.removeEventListener("abort", onAbort);
       unsubInput?.();
       ctx?.ui.setStatus(STATUS_KEY, undefined);
@@ -739,13 +829,15 @@ function installFetchRateLimitObserver(): void {
     const response = await originalFetch(...args);
 
     if (activeProviderRequests > 0) {
-      if (response.status === 429) {
-        const waitMs = retryAfterHeaderMs(response.headers) ?? DEFAULT_RATE_LIMIT_WAIT_MS;
+      if (!response.ok) {
         const message = await responseErrorMessage(response);
-        lastObservedRateLimitError = { at: Date.now(), message };
-        showAmbientRetryStatus("rate-limit", waitMs);
-        if (activeFallbackProviderRequests > 0) {
-          throw new Error(message);
+        lastObservedHttpError = { at: Date.now(), message };
+        if (response.status === 429) {
+          const waitMs = retryAfterHeaderMs(response.headers) ?? DEFAULT_RATE_LIMIT_WAIT_MS;
+          showAmbientRetryStatus("rate-limit", waitMs);
+          if (activeFallbackProviderRequests > 0) {
+            throw new Error(message);
+          }
         }
       } else if (ambientStatusCleanup && response.ok) {
         clearAmbientRetryStatus();
@@ -769,6 +861,11 @@ type StreamSimpleFn = (
   options?: SimpleStreamOptions,
 ) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;
 
+export function __setNonRetryableTuningForTests(maxAttempts: number, retryDelayMs: number): void {
+  nonRetryableMaxAttempts = maxAttempts;
+  nonRetryableRetryDelayMs = retryDelayMs;
+}
+
 export function __configureFallbackModelsForTests(
   models: FallbackModel[],
   ctx?: ExtensionContext,
@@ -779,7 +876,7 @@ export function __configureFallbackModelsForTests(
   primaryThinkingLevel = undefined;
   rateLimitMemory.clear();
   nonRetryableFailureMemory.clear();
-  lastObservedRateLimitError = undefined;
+  lastObservedHttpError = undefined;
   ambientStatusCleanup?.();
   modelStatusCleanup?.();
 }
@@ -821,6 +918,11 @@ export function streamWithLimitsRetry(
     let committed = false;
     const allowFallback = fallbackEnabled() && isFallbackEligibleModel(model);
     let attempt: FallbackModel = allowFallback ? initialAttempt(model) : { model, reasoningEffort: primaryThinkingLevel };
+    // Per-stream bookkeeping: how many times each model has hit a non-retryable
+    // error, and which models we've exhausted (used to advance fallbacks when
+    // freezing is disabled).
+    const nonRetryableAttempts = new Map<string, number>();
+    const triedModels = new Set<string>();
 
     const flush = (buffer: AssistantMessageEvent[]) => {
       for (const event of buffer) output.push(event);
@@ -850,7 +952,8 @@ export function streamWithLimitsRetry(
           for await (const event of inner) {
             if (!committed) {
               if (event.type === "error") {
-                const errMsg = event.error.errorMessage ?? "";
+                const rawErrMsg = event.error.errorMessage ?? "";
+                const errMsg = errorMessageWithRecentHttpStatus(rawErrMsg);
 
                 if (isAbortErrorEvent(event, options?.signal)) {
                   if (buffer.length > 0) {
@@ -879,20 +982,11 @@ export function streamWithLimitsRetry(
                   break;
                 }
 
-                if (allowFallback) {
-                  nonRetryableError = { message: errMsg, event };
-                  break;
-                }
-
-                if (buffer.length > 0) {
-                  flush(buffer);
-                } else {
-                  output.push({ type: "start", partial: freshMessage(attempt.model) });
-                  committed = true;
-                }
-                output.push(event);
-                output.end();
-                return;
+                // Defer handling to the post-loop block, which applies the
+                // bounded non-retryable retry / fallback / freeze logic. This
+                // runs whether or not fallbacks are configured.
+                nonRetryableError = { message: errMsg, event };
+                break;
               }
 
               if (event.type === "start") {
@@ -924,7 +1018,7 @@ export function streamWithLimitsRetry(
             }
           }
         } catch (err) {
-          const errMsg = errorMessageWithCauses(err);
+          const errMsg = errorMessageWithRecentHttpStatus(errorMessageWithCauses(err));
           const aborted = Boolean(options?.signal?.aborted) || isAbortMessage(errMsg);
           retryable = !aborted ? getRetryableError(errMsg) : undefined;
           if (retryable) retryableErrorMessage = errMsg;
@@ -934,25 +1028,21 @@ export function streamWithLimitsRetry(
             if (retryable && observedRateLimit) retryableErrorMessage = observedRateLimit;
           }
           if (!retryable) {
-            if (allowFallback && !aborted) {
-              nonRetryableError = { message: errMsg };
-            } else {
+            if (aborted) {
               if (!committed) {
                 // Synthetic start+error so the stream protocol stays valid.
                 output.push({ type: "start", partial: freshMessage(attempt.model) });
                 committed = true;
               }
               const error = freshMessage(attempt.model);
-              error.stopReason = aborted ? "aborted" : "error";
+              error.stopReason = "aborted";
               error.errorMessage = errMsg;
-              output.push({
-                type: "error",
-                reason: error.stopReason as "error" | "aborted",
-                error,
-              });
+              output.push({ type: "error", reason: "aborted", error });
               output.end();
               return;
             }
+            // Defer to the post-loop block for bounded retries / fallback.
+            nonRetryableError = { message: errMsg };
           }
         }
       } finally {
@@ -964,72 +1054,122 @@ export function streamWithLimitsRetry(
         }
       }
 
-      if (nonRetryableError) {
-        if (allowFallback && !committed) {
-          rememberNonRetryableFailure(attempt.model, nonRetryableError.message);
-          const next = nextAvailableCandidate(attempt.model);
-          if (next) {
-            attempt = next;
-            continue;
-          }
+      if (!retryable) {
+        // Both explicit non-retryable error events/exceptions and a provider
+        // stream that ended without a terminal event land here.
+        const failureEvent = nonRetryableError?.event;
+        const failureMessage = nonRetryableError?.message ?? "Provider stream ended without a terminal event.";
 
-          const deadline = earliestCandidateDeadline(attempt.model);
-          if (deadline) {
-            const waitResult = await waitForRetry("model-frozen", Math.max(0, deadline - Date.now()), options?.signal);
-            if (waitResult !== "aborted") {
-              attempt = nextAvailableCandidate(attempt.model) ?? { model: getPrimaryModel(attempt.model), reasoningEffort: primaryThinkingLevel };
-              continue;
+        const emitFailure = () => {
+          if (!committed) {
+            if (buffer.length > 0) {
+              flush(buffer);
+            } else {
+              output.push({ type: "start", partial: freshMessage(attempt.model) });
+              committed = true;
             }
           }
-        }
-
-        if (!committed) {
-          if (buffer.length > 0) {
-            flush(buffer);
+          if (failureEvent) {
+            output.push(failureEvent);
           } else {
+            const error = freshMessage(attempt.model);
+            error.stopReason = "error";
+            error.errorMessage = failureMessage;
+            output.push({ type: "error", reason: "error", error });
+          }
+          output.end();
+        };
+
+        const emitAborted = () => {
+          if (!committed) {
             output.push({ type: "start", partial: freshMessage(attempt.model) });
             committed = true;
           }
-        }
-        if (nonRetryableError.event) {
-          output.push(nonRetryableError.event);
-        } else {
           const error = freshMessage(attempt.model);
-          error.stopReason = "error";
-          error.errorMessage = nonRetryableError.message;
-          output.push({ type: "error", reason: "error", error });
-        }
-        output.end();
-        return;
-      }
+          error.stopReason = "aborted";
+          error.errorMessage = "Request aborted while retrying after error.";
+          output.push({ type: "error", reason: "aborted", error });
+          output.end();
+        };
 
-      if (!retryable) {
-        if (allowFallback && !committed) {
-          const message = "Provider stream ended without a terminal event.";
-          rememberNonRetryableFailure(attempt.model, message);
-          const next = nextAvailableCandidate(attempt.model);
-          if (next) {
-            attempt = next;
-            continue;
-          }
-        }
+        // Nothing has been streamed yet: give the same model a few quick retries
+        // before we freeze it or fall back. This rides out transient failures
+        // that aren't recognised as retryable (stalls, odd 5xx, malformed
+        // errors) instead of immediately freezing the model for an hour.
         if (!committed) {
-          if (buffer.length > 0) {
-            flush(buffer);
-          } else {
-            output.push({ type: "start", partial: freshMessage(attempt.model) });
-            committed = true;
+          const key = modelKey(attempt.model);
+          const attempts = (nonRetryableAttempts.get(key) ?? 0) + 1;
+          nonRetryableAttempts.set(key, attempts);
+
+          if (attempts < nonRetryableMaxAttempts) {
+            const waitResult = await waitForRetry("retry", nonRetryableRetryDelayMs, options?.signal);
+            if (waitResult === "aborted") {
+              emitAborted();
+              return;
+            }
+            continue; // retry the same model
+          }
+
+          // Retries for this model are exhausted.
+          triedModels.add(key);
+
+          if (allowFallback) {
+            if (freezingEnabled()) {
+              rememberNonRetryableFailure(attempt.model, failureMessage);
+              const next = nextAvailableCandidate(attempt.model);
+              if (next) {
+                attempt = next;
+                continue;
+              }
+
+              const deadline = earliestCandidateDeadline(attempt.model);
+              if (deadline) {
+                const waitResult = await waitForRetry("model-frozen", Math.max(0, deadline - Date.now()), options?.signal);
+                if (waitResult !== "aborted") {
+                  attempt = nextAvailableCandidate(attempt.model) ?? { model: getPrimaryModel(attempt.model), reasoningEffort: primaryThinkingLevel };
+                  continue;
+                }
+              }
+            } else {
+              // Freezing disabled: walk through the remaining configured
+              // candidates (without freezing) and only give up once they are
+              // all exhausted. Never block on a long "model-frozen" wait.
+              const next = candidateOrder(attempt.model).find((entry) => !triedModels.has(modelKey(entry.model)));
+              if (next) {
+                attempt = next;
+                continue;
+              }
+            }
           }
         }
-        const error = freshMessage(attempt.model);
-        error.stopReason = "error";
-        error.errorMessage = "Provider stream ended without a terminal event.";
-        output.push({ type: "error", reason: "error", error });
-        output.end();
+
+        emitFailure();
         return;
       }
 
       if (allowFallback) {
+        // Network/transport failures usually mean the user's connection dropped
+        // or a socket timed out. Do not mark the model as unavailable, switch
+        // to another model, or enter model-freeze handling; keep retrying the
+        // same model with the short network backoff until it works again.
+        if (retryable.reason === "network") {
+          notifyRetryableError(attempt.model, retryable, retryableErrorMessage);
+          const waitResult = await waitForRetry(retryable.reason, retryable.waitMs, options?.signal);
+          if (waitResult === "aborted") {
+            if (!committed) {
+              output.push({ type: "start", partial: freshMessage(attempt.model) });
+              committed = true;
+            }
+            const error = freshMessage(attempt.model);
+            error.stopReason = "aborted";
+            error.errorMessage = `Request aborted during ${reasonLabel(retryable.reason).toLowerCase()} wait.`;
+            output.push({ type: "error", reason: "aborted", error });
+            output.end();
+            return;
+          }
+          continue;
+        }
+
         rememberRateLimit(attempt.model, retryable, retryableErrorMessage);
         const next = nextAvailableCandidate(attempt.model);
         if (next) {