oira666_pi-limits-wait 0.4.1 → 0.4.3

package/package.json

@@ -1,49 +1,49 @@
-{
-  "name": "oira666_pi-limits-wait",
-  "version": "0.4.1",
-  "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
-  "type": "module",
-  "license": "MIT",
-  "files": [
-    "src",
-    "!src/__tests__",
-    "README.md",
-    "LICENSE"
-  ],
-  "engines": {
-    "node": ">=18"
-  },
-  "repository": {
-    "type": "git",
-    "url": "https://github.com/gee666/pi-limits-wait.git"
-  },
-  "bugs": {
-    "url": "https://github.com/gee666/pi-limits-wait/issues"
-  },
-  "homepage": "https://github.com/gee666/pi-limits-wait#readme",
-  "publishConfig": {
-    "access": "public"
-  },
-  "scripts": {
-    "check": "tsc -p tsconfig.json --noEmit",
-    "test": "tsx src/__tests__/test-extension.ts",
-    "pack:check": "npm pack --dry-run",
-    "prepublishOnly": "npm run check && npm test && npm run pack:check"
-  },
-  "pi": {
-    "extensions": [
-      "./src/index.ts"
-    ]
-  },
-  "dependencies": {},
-  "peerDependencies": {
-    "@mariozechner/pi-ai": "*",
-    "@mariozechner/pi-coding-agent": "*"
-  },
-  "devDependencies": {
-    "@mariozechner/pi-ai": "^0.66.1",
-    "@mariozechner/pi-coding-agent": "^0.66.1",
-    "tsx": "^4.20.6",
-    "typescript": "^5.9.3"
-  }
-}
+{
+  "name": "oira666_pi-limits-wait",
+  "version": "0.4.3",
+  "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
+  "type": "module",
+  "license": "MIT",
+  "files": [
+    "src",
+    "!src/__tests__",
+    "README.md",
+    "LICENSE"
+  ],
+  "engines": {
+    "node": ">=18"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/gee666/pi-limits-wait.git"
+  },
+  "bugs": {
+    "url": "https://github.com/gee666/pi-limits-wait/issues"
+  },
+  "homepage": "https://github.com/gee666/pi-limits-wait#readme",
+  "publishConfig": {
+    "access": "public"
+  },
+  "scripts": {
+    "check": "tsc -p tsconfig.json --noEmit",
+    "test": "tsx src/__tests__/test-extension.ts",
+    "pack:check": "npm pack --dry-run",
+    "prepublishOnly": "npm run check && npm test && npm run pack:check"
+  },
+  "pi": {
+    "extensions": [
+      "./src/index.ts"
+    ]
+  },
+  "dependencies": {},
+  "peerDependencies": {
+    "@mariozechner/pi-ai": "*",
+    "@mariozechner/pi-coding-agent": "*"
+  },
+  "devDependencies": {
+    "@mariozechner/pi-ai": "^0.66.1",
+    "@mariozechner/pi-coding-agent": "^0.66.1",
+    "tsx": "^4.20.6",
+    "typescript": "^5.9.3"
+  }
+}

package/src/index.ts

@@ -28,13 +28,14 @@ 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 SETTINGS_KEY = "oira666_pi-limits-wait";
 
 /** 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";
+type RetryReason = "rate-limit" | "overloaded" | "authentication" | "model-frozen";
 
 type RetryableError = {
   reason: RetryReason;
@@ -47,7 +48,7 @@ type ConfiguredModel = {
   reasoningEffort?: ThinkingLevel;
 };
 
-type FallbackModel = {
+export type FallbackModel = {
   model: Model<Api>;
   reasoningEffort?: ThinkingLevel;
 };
@@ -74,6 +75,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 }>();
 
 // ─── Anthropic subscription prompt sanitisation ──────────────────────────────
 
@@ -113,6 +115,16 @@ function formatModel(model: Pick<Model<Api>, "provider" | "id">): string {
   return `${model.provider}/${model.id}`;
 }
 
+function isInternalSyntheticModel(model: Pick<Model<Api>, "provider" | "id">): boolean {
+  return model.provider.startsWith("pi-") || model.id.startsWith("synthetic-");
+}
+
+function isFallbackEligibleModel(model: Model<Api>): boolean {
+  if (isInternalSyntheticModel(model)) return false;
+  const key = modelKey(model);
+  return primaryModel ? key === modelKey(primaryModel) || fallbackModels.some((entry) => modelKey(entry.model) === key) : true;
+}
+
 function isThinkingLevel(value: unknown): value is ThinkingLevel {
   return ["off", "minimal", "low", "medium", "high", "xhigh"].includes(String(value));
 }
@@ -247,6 +259,13 @@ function updateRateLimitedModelsStatus(): void {
     }
     lines.push(`${key}: ${formatDuration(entry.deadline - now)}`);
   }
+  for (const [key, entry] of nonRetryableFailureMemory) {
+    if (now >= entry.deadline) {
+      nonRetryableFailureMemory.delete(key);
+      continue;
+    }
+    lines.push(`${key}: frozen ${formatDuration(entry.deadline - now)}`);
+  }
   if (lines.length === 0) {
     modelStatusCleanup?.();
     return;
@@ -272,13 +291,34 @@ function rememberRateLimit(model: Model<Api>, retryable: RetryableError): void {
   ensureRateLimitedModelsStatus();
 }
 
+function activeNonRetryableFailure(model: Model<Api>): { failedAt: number; deadline: number; errorMessage: string } | undefined {
+  const entry = nonRetryableFailureMemory.get(modelKey(model));
+  if (!entry) return undefined;
+  if (Date.now() >= entry.deadline) {
+    nonRetryableFailureMemory.delete(modelKey(model));
+    return undefined;
+  }
+  return entry;
+}
+
+function hasNonRetryableFailure(model: Model<Api>): boolean {
+  return Boolean(activeNonRetryableFailure(model));
+}
+
 function nextAvailableCandidate(current: Model<Api>): FallbackModel | undefined {
-  return candidateOrder(current).find((entry) => !activeLimit(entry.model));
+  return candidateOrder(current).find((entry) => !activeLimit(entry.model) && !hasNonRetryableFailure(entry.model));
+}
+
+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");
+  ensureRateLimitedModelsStatus();
 }
 
 function earliestCandidateDeadline(current: Model<Api>): number | undefined {
   const deadlines = candidateOrder(current)
-    .map((entry) => activeLimit(entry.model)?.deadline)
+    .map((entry) => activeLimit(entry.model)?.deadline ?? activeNonRetryableFailure(entry.model)?.deadline)
     .filter((deadline): deadline is number => typeof deadline === "number");
   return deadlines.length > 0 ? Math.min(...deadlines) : undefined;
 }
@@ -286,7 +326,7 @@ function earliestCandidateDeadline(current: Model<Api>): number | undefined {
 function initialAttempt(model: Model<Api>): FallbackModel {
   const configured = fallbackModels.find((entry) => modelKey(entry.model) === modelKey(model));
   const current = configured ?? { model, reasoningEffort: primaryThinkingLevel };
-  if (!fallbackEnabled() || !activeLimit(model)) return current;
+  if (!fallbackEnabled() || (!activeLimit(model) && !hasNonRetryableFailure(model))) return current;
   return nextAvailableCandidate(model) ?? current;
 }
 
@@ -351,6 +391,18 @@ export function isServerOverloadedError(msg: string): boolean {
   );
 }
 
+export function isAuthenticationRefreshError(msg: string): boolean {
+  const lower = msg.toLowerCase();
+  return (
+    /(?:^|\D)401(?:\D|$)/.test(msg) &&
+    (
+      lower.includes("authentication_error") ||
+      lower.includes("invalid authentication credentials") ||
+      lower.includes("invalid authentication")
+    )
+  );
+}
+
 export function parseRetryDelayMs(msg: string): number | undefined {
   const retryAfter = msg.match(/retry-after(?:-ms)?[^0-9]*(\d+(?:\.\d+)?)/i);
   if (retryAfter?.[1]) {
@@ -416,7 +468,10 @@ export function retryAfterHeaderMs(headers: Headers): number | undefined {
 
 export function getRetryableError(msg: string): RetryableError | undefined {
   if (isServerOverloadedError(msg)) {
-    return { reason: "overloaded", waitMs: DEFAULT_OVERLOADED_WAIT_MS };
+    return { reason: "overloaded", waitMs: parseRetryDelayMs(msg) ?? DEFAULT_OVERLOADED_WAIT_MS };
+  }
+  if (isAuthenticationRefreshError(msg)) {
+    return { reason: "authentication", waitMs: parseRetryDelayMs(msg) ?? DEFAULT_RATE_LIMIT_WAIT_MS };
   }
   if (isRateLimitError(msg)) {
     return { reason: "rate-limit", waitMs: rateLimitWaitMs(msg) };
@@ -427,7 +482,10 @@ export function getRetryableError(msg: string): RetryableError | undefined {
 // ─── Countdown UI ─────────────────────────────────────────────────────────────
 
 function reasonLabel(reason: RetryReason): string {
-  return reason === "overloaded" ? "Server overloaded" : "Rate limited";
+  if (reason === "overloaded") return "Server overloaded";
+  if (reason === "authentication") return "Authentication refresh pending";
+  if (reason === "model-frozen") return "Model frozen after error";
+  return "Rate limited";
 }
 
 function statusText(reason: RetryReason, deadline: number, allowSkip: boolean): string {
@@ -580,6 +638,20 @@ type StreamSimpleFn = (
   options?: SimpleStreamOptions,
 ) => AssistantMessageEventStream;
 
+export function __configureFallbackModelsForTests(
+  models: FallbackModel[],
+  ctx?: ExtensionContext,
+): void {
+  fallbackModels = models;
+  sharedCtx = ctx;
+  primaryModel = undefined;
+  primaryThinkingLevel = undefined;
+  rateLimitMemory.clear();
+  nonRetryableFailureMemory.clear();
+  ambientStatusCleanup?.();
+  modelStatusCleanup?.();
+}
+
 function freshMessage(model: Model<Api>): AssistantMessage {
   return {
     role: "assistant",
@@ -615,7 +687,8 @@ export function streamWithLimitsRetry(
 
   void (async () => {
     let committed = false;
-    let attempt: FallbackModel = initialAttempt(model);
+    const allowFallback = fallbackEnabled() && isFallbackEligibleModel(model);
+    let attempt: FallbackModel = allowFallback ? initialAttempt(model) : { model, reasoningEffort: primaryThinkingLevel };
 
     const flush = (buffer: AssistantMessageEvent[]) => {
       for (const event of buffer) output.push(event);
@@ -627,10 +700,11 @@ export function streamWithLimitsRetry(
       installFetchRateLimitObserver();
       const buffer: AssistantMessageEvent[] = [];
       let retryable: RetryableError | undefined;
+      let nonRetryableError: { message: string; event?: AssistantMessageEvent } | undefined;
 
       try {
         try {
-          const attemptOptions = fallbackEnabled()
+          const attemptOptions = allowFallback
             ? await optionsForModel(model, attempt, options)
             : options;
           const attemptDelegate = attempt.model.api === model.api
@@ -646,7 +720,12 @@ export function streamWithLimitsRetry(
                 retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
                 if (retryable) break;
 
-                // Non-retryable error: ensure start came first, then forward error.
+                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 {
@@ -664,7 +743,7 @@ export function streamWithLimitsRetry(
               }
 
               // First non-start, non-error event means this attempt is real.
-              if (fallbackEnabled()) await switchPiModel(attempt);
+              if (allowFallback) await switchPiModel(attempt);
               if (buffer.length > 0) {
                 flush(buffer);
               } else {
@@ -690,6 +769,9 @@ export function streamWithLimitsRetry(
           const errMsg = err instanceof Error ? err.message : String(err);
           retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
           if (!retryable) {
+            if (allowFallback && !options?.signal?.aborted) {
+              nonRetryableError = { message: errMsg };
+            } else {
             if (!committed) {
               // Synthetic start+error so the stream protocol stays valid.
               output.push({ type: "start", partial: freshMessage(attempt.model) });
@@ -705,6 +787,7 @@ export function streamWithLimitsRetry(
             });
             output.end();
             return;
+            }
           }
         }
       } finally {
@@ -715,6 +798,45 @@ export function streamWithLimitsRetry(
         }
       }
 
+      if (nonRetryableError) {
+        if (allowFallback && !committed) {
+          rememberNonRetryableFailure(attempt.model, nonRetryableError.message);
+          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;
+            }
+          }
+        }
+
+        if (!committed) {
+          if (buffer.length > 0) {
+            flush(buffer);
+          } else {
+            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;
+      }
+
       if (!retryable) {
         if (!committed) {
           if (buffer.length > 0) {
@@ -732,7 +854,7 @@ export function streamWithLimitsRetry(
         return;
       }
 
-      if (fallbackEnabled() && retryable.reason === "rate-limit") {
+      if (allowFallback) {
         rememberRateLimit(attempt.model, retryable);
         const next = nextAvailableCandidate(attempt.model);
         if (next) {
@@ -807,22 +929,23 @@ export default function (pi: ExtensionAPI) {
       suppressNextModelSelect = false;
       return;
     }
+    if (isInternalSyntheticModel(event.model)) return;
     primaryModel = event.model;
     primaryThinkingLevel = pi.getThinkingLevel();
   });
 
   pi.on("before_agent_start", (event, ctx) => {
     sharedCtx = ctx;
-    primaryModel ??= ctx.model;
+    if (!primaryModel && ctx.model && !isInternalSyntheticModel(ctx.model)) primaryModel = ctx.model;
     primaryThinkingLevel ??= pi.getThinkingLevel();
     loadFallbackSettings(ctx);
 
     // Some extensions may register providers after this extension loads. Wrap
     // any APIs that exist by the time an agent starts too.
     for (const model of ctx.modelRegistry.getAll()) {
-      registerWrappedApi(pi, model.api);
+      if (!isInternalSyntheticModel(model)) registerWrappedApi(pi, model.api);
     }
-    if (ctx.model) registerWrappedApi(pi, ctx.model.api);
+    if (ctx.model && !isInternalSyntheticModel(ctx.model)) registerWrappedApi(pi, ctx.model.api);
 
     // Anthropic subscription/OAuth requests identify as Claude Code, not Pi.
     // Check the session-wide Anthropic OAuth configuration rather than only the