oira666_pi-limits-wait 0.4.4 → 0.4.7

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-limits-wait",
-  "version": "0.4.4",
+  "version": "0.4.7",
   "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

@@ -59,6 +59,12 @@ type RateLimitMemory = {
   deadline: number;
 };
 
+type NonRetryableFailureMemory = {
+  failedAt: number;
+  deadline: number;
+  errorMessage: string;
+};
+
 // ─── Shared UI context ────────────────────────────────────────────────────────
 
 let sharedCtx: ExtensionContext | undefined;
@@ -67,6 +73,8 @@ let restoreFetch: (() => void) | undefined;
 let ambientStatusCleanup: (() => void) | undefined;
 let modelStatusCleanup: (() => void) | undefined;
 let activeProviderRequests = 0;
+let activeFallbackProviderRequests = 0;
+let lastObservedRateLimitError: { at: number; message: string } | undefined;
 const wrappedApis = new Set<Api>();
 const builtinStreamSimpleByApi = new Map<Api, StreamSimpleFn>();
 let fallbackModels: FallbackModel[] = [];
@@ -75,7 +83,7 @@ let primaryThinkingLevel: ThinkingLevel | undefined;
 let suppressNextModelSelect = false;
 let settingsSignature: string | undefined;
 const rateLimitMemory = new Map<string, RateLimitMemory>();
-const nonRetryableFailureMemory = new Map<string, { failedAt: number; deadline: number; errorMessage: string }>();
+const nonRetryableFailureMemory = new Map<string, NonRetryableFailureMemory>();
 
 // ─── Anthropic subscription prompt sanitisation ──────────────────────────────
 
@@ -250,27 +258,31 @@ function formatDuration(ms: number): string {
 function updateRateLimitedModelsStatus(): void {
   const ctx = sharedCtx;
   if (!ctx) return;
-  const now = Date.now();
-  const lines: string[] = [];
-  for (const [key, entry] of rateLimitMemory) {
-    if (now >= entry.deadline) {
-      rateLimitMemory.delete(key);
-      continue;
+  try {
+    const now = Date.now();
+    const lines: string[] = [];
+    for (const [key, entry] of rateLimitMemory) {
+      if (now >= entry.deadline) {
+        rateLimitMemory.delete(key);
+        continue;
+      }
+      lines.push(`${key}: ${formatDuration(entry.deadline - now)}`);
     }
-    lines.push(`${key}: ${formatDuration(entry.deadline - now)}`);
-  }
-  for (const [key, entry] of nonRetryableFailureMemory) {
-    if (now >= entry.deadline) {
-      nonRetryableFailureMemory.delete(key);
-      continue;
+    for (const [key, entry] of nonRetryableFailureMemory) {
+      if (now >= entry.deadline) {
+        nonRetryableFailureMemory.delete(key);
+        continue;
+      }
+      lines.push(`${key}: frozen ${formatDuration(entry.deadline - now)}`);
     }
-    lines.push(`${key}: frozen ${formatDuration(entry.deadline - now)}`);
-  }
-  if (lines.length === 0) {
+    if (lines.length === 0) {
+      modelStatusCleanup?.();
+      return;
+    }
+    ctx.ui.setStatus(MODELS_STATUS_KEY, `🚦 Limited: ${lines.join(" | ")}`);
+  } catch {
     modelStatusCleanup?.();
-    return;
   }
-  ctx.ui.setStatus(MODELS_STATUS_KEY, `🚦 Limited: ${lines.join(" | ")}`);
 }
 
 function ensureRateLimitedModelsStatus(): void {
@@ -280,18 +292,23 @@ function ensureRateLimitedModelsStatus(): void {
   modelStatusCleanup = () => {
     clearInterval(ticker);
     modelStatusCleanup = undefined;
-    sharedCtx?.ui.setStatus(MODELS_STATUS_KEY, undefined);
+    try {
+      sharedCtx?.ui.setStatus(MODELS_STATUS_KEY, undefined);
+    } catch {
+      // The extension context can become stale after print-mode session teardown.
+    }
   };
 }
 
-function rememberRateLimit(model: Model<Api>, retryable: RetryableError): void {
+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 });
-  sharedCtx?.ui.notify(`${formatModel(model)} ${reasonLabel(retryable.reason).toLowerCase()} for ${formatDuration(retryable.waitMs)}.`, "warning");
+  const detail = errorMessage ? ` Error: ${formatErrorDetail(errorMessage)}` : "";
+  sharedCtx?.ui.notify(`${formatModel(model)} ${reasonLabel(retryable.reason).toLowerCase()} for ${formatDuration(retryable.waitMs)}.${detail}`, "warning");
   ensureRateLimitedModelsStatus();
 }
 
-function activeNonRetryableFailure(model: Model<Api>): { failedAt: number; deadline: number; errorMessage: string } | undefined {
+function activeNonRetryableFailure(model: Model<Api>): NonRetryableFailureMemory | undefined {
   const entry = nonRetryableFailureMemory.get(modelKey(model));
   if (!entry) return undefined;
   if (Date.now() >= entry.deadline) {
@@ -309,13 +326,63 @@ function nextAvailableCandidate(current: Model<Api>): FallbackModel | undefined
   return candidateOrder(current).find((entry) => !activeLimit(entry.model) && !hasNonRetryableFailure(entry.model));
 }
 
+function formatErrorDetail(errorMessage: string, maxLength = 500): string {
+  const oneLine = errorMessage.replace(/\s+/g, " ").trim();
+  if (!oneLine) return "unknown error";
+  return oneLine.length > maxLength ? `${oneLine.slice(0, maxLength - 1)}…` : oneLine;
+}
+
 function rememberNonRetryableFailure(model: Model<Api>, errorMessage: string): void {
   const deadline = Date.now() + DEFAULT_NON_RETRYABLE_FREEZE_MS;
   nonRetryableFailureMemory.set(modelKey(model), { failedAt: Date.now(), deadline, errorMessage });
-  sharedCtx?.ui.notify(`${formatModel(model)} failed; freezing it for ${formatDuration(DEFAULT_NON_RETRYABLE_FREEZE_MS)} and trying another configured model if available.`, "warning");
+  sharedCtx?.ui.notify(`${formatModel(model)} failed (${formatErrorDetail(errorMessage)}); freezing it for ${formatDuration(DEFAULT_NON_RETRYABLE_FREEZE_MS)} and trying another configured model if available.`, "warning");
   ensureRateLimitedModelsStatus();
 }
 
+function errorMessageWithCauses(err: unknown): string {
+  const messages: string[] = [];
+  let current: unknown = err;
+  const seen = new Set<unknown>();
+
+  while (current && !seen.has(current)) {
+    seen.add(current);
+    if (current instanceof Error) {
+      if (current.message) messages.push(current.message);
+      current = (current as Error & { cause?: unknown }).cause;
+      continue;
+    }
+    messages.push(String(current));
+    break;
+  }
+
+  return messages.filter(Boolean).join(" Cause: ") || String(err);
+}
+
+function isAbortMessage(msg: string): boolean {
+  const lower = msg.toLowerCase();
+  return (
+    lower.includes("operation aborted") ||
+    lower.includes("request aborted") ||
+    lower.includes("aborterror") ||
+    /\baborted\b/.test(lower)
+  );
+}
+
+function isAbortErrorEvent(event: AssistantMessageEvent, signal?: AbortSignal): boolean {
+  if (signal?.aborted) return true;
+  if (event.type !== "error") return false;
+  const reason = "reason" in event ? String(event.reason) : "";
+  const stopReason = event.error.stopReason ? String(event.error.stopReason) : "";
+  const message = event.error.errorMessage ?? "";
+  return reason === "aborted" || stopReason === "aborted" || isAbortMessage(message);
+}
+
+function recentObservedRateLimitError(maxAgeMs = 60_000): string | undefined {
+  if (!lastObservedRateLimitError) return undefined;
+  if (Date.now() - lastObservedRateLimitError.at > maxAgeMs) return undefined;
+  return lastObservedRateLimitError.message;
+}
+
 function earliestCandidateDeadline(current: Model<Api>): number | undefined {
   const deadlines = candidateOrder(current)
     .map((entry) => activeLimit(entry.model)?.deadline ?? activeNonRetryableFailure(entry.model)?.deadline)
@@ -605,6 +672,19 @@ export function waitForRateLimit(
 
 // ─── Early 429 observer ──────────────────────────────────────────────────────
 
+async function responseErrorMessage(response: Response): Promise<string> {
+  const retryAfterMs = retryAfterHeaderMs(response.headers);
+  const retryAfter = retryAfterMs !== undefined ? ` retry-after-ms ${retryAfterMs}` : "";
+  let body = "";
+  try {
+    body = await response.clone().text();
+  } catch {
+    // Ignore body read failures; status and retry headers are enough.
+  }
+  const detail = body.trim() ? ` ${formatErrorDetail(body)}` : "";
+  return `HTTP ${response.status}${response.statusText ? ` ${response.statusText}` : ""}${retryAfter}${detail}`;
+}
+
 function installFetchRateLimitObserver(): void {
   if (restoreFetch || typeof globalThis.fetch !== "function") return;
 
@@ -615,7 +695,12 @@ function installFetchRateLimitObserver(): void {
     if (activeProviderRequests > 0) {
       if (response.status === 429) {
         const waitMs = retryAfterHeaderMs(response.headers) ?? DEFAULT_RATE_LIMIT_WAIT_MS;
+        const message = await responseErrorMessage(response);
+        lastObservedRateLimitError = { at: Date.now(), message };
         showAmbientRetryStatus("rate-limit", waitMs);
+        if (activeFallbackProviderRequests > 0) {
+          throw new Error(message);
+        }
       } else if (ambientStatusCleanup && response.ok) {
         clearAmbientRetryStatus();
       }
@@ -648,6 +733,7 @@ export function __configureFallbackModelsForTests(
   primaryThinkingLevel = undefined;
   rateLimitMemory.clear();
   nonRetryableFailureMemory.clear();
+  lastObservedRateLimitError = undefined;
   ambientStatusCleanup?.();
   modelStatusCleanup?.();
 }
@@ -697,9 +783,11 @@ export function streamWithLimitsRetry(
 
     while (true) {
       activeProviderRequests++;
+      if (allowFallback) activeFallbackProviderRequests++;
       installFetchRateLimitObserver();
       const buffer: AssistantMessageEvent[] = [];
       let retryable: RetryableError | undefined;
+      let retryableErrorMessage = "";
       let nonRetryableError: { message: string; event?: AssistantMessageEvent } | undefined;
 
       try {
@@ -717,15 +805,39 @@ export function streamWithLimitsRetry(
             if (!committed) {
               if (event.type === "error") {
                 const errMsg = event.error.errorMessage ?? "";
-                retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
-                if (retryable) break;
+
+                if (isAbortErrorEvent(event, options?.signal)) {
+                  if (buffer.length > 0) {
+                    flush(buffer);
+                  } else {
+                    output.push({ type: "start", partial: freshMessage(attempt.model) });
+                    committed = true;
+                  }
+                  const error = freshMessage(attempt.model);
+                  error.stopReason = "aborted";
+                  error.errorMessage = errMsg || "Operation aborted";
+                  output.push({ type: "error", reason: "aborted", error });
+                  output.end();
+                  return;
+                }
+
+                retryable = getRetryableError(errMsg);
+                if (retryable) {
+                  retryableErrorMessage = errMsg;
+                  break;
+                }
+                const observedRateLimit = allowFallback ? recentObservedRateLimitError() : undefined;
+                retryable = observedRateLimit ? getRetryableError(observedRateLimit) : undefined;
+                if (retryable && observedRateLimit) {
+                  retryableErrorMessage = observedRateLimit;
+                  break;
+                }
 
                 if (allowFallback) {
                   nonRetryableError = { message: errMsg, event };
                   break;
                 }
 
-                // Non-retryable error without fallbacks: ensure start came first, then forward error.
                 if (buffer.length > 0) {
                   flush(buffer);
                 } else {
@@ -766,32 +878,40 @@ export function streamWithLimitsRetry(
             }
           }
         } catch (err) {
-          const errMsg = err instanceof Error ? err.message : String(err);
-          retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
+          const errMsg = errorMessageWithCauses(err);
+          const aborted = Boolean(options?.signal?.aborted) || isAbortMessage(errMsg);
+          retryable = !aborted ? getRetryableError(errMsg) : undefined;
+          if (retryable) retryableErrorMessage = errMsg;
+          if (!retryable && allowFallback) {
+            const observedRateLimit = recentObservedRateLimitError();
+            retryable = observedRateLimit ? getRetryableError(observedRateLimit) : undefined;
+            if (retryable && observedRateLimit) retryableErrorMessage = observedRateLimit;
+          }
           if (!retryable) {
-            if (allowFallback && !options?.signal?.aborted) {
+            if (allowFallback && !aborted) {
               nonRetryableError = { message: errMsg };
             } else {
-            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 = options?.signal?.aborted ? "aborted" : "error";
-            error.errorMessage = errMsg;
-            output.push({
-              type: "error",
-              reason: error.stopReason as "error" | "aborted",
-              error,
-            });
-            output.end();
-            return;
+              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.errorMessage = errMsg;
+              output.push({
+                type: "error",
+                reason: error.stopReason as "error" | "aborted",
+                error,
+              });
+              output.end();
+              return;
             }
           }
         }
       } finally {
         activeProviderRequests--;
+        if (allowFallback) activeFallbackProviderRequests--;
         if (activeProviderRequests === 0) {
           clearAmbientRetryStatus();
           restoreFetch?.();
@@ -838,6 +958,15 @@ export function streamWithLimitsRetry(
       }
 
       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;
+          }
+        }
         if (!committed) {
           if (buffer.length > 0) {
             flush(buffer);
@@ -855,7 +984,7 @@ export function streamWithLimitsRetry(
       }
 
       if (allowFallback) {
-        rememberRateLimit(attempt.model, retryable);
+        rememberRateLimit(attempt.model, retryable, retryableErrorMessage);
         const next = nextAvailableCandidate(attempt.model);
         if (next) {
           attempt = next;