oira666_pi-limits-wait 0.6.4 → 0.6.6

package/README.md

@@ -19,7 +19,7 @@ Restart Pi or run `/reload` after installing.
 - On Pi 0.86+, stops background cache warming via `cache_warming_decision` while retry interception is enabled, preventing background refreshes from triggering foreground retries or fallbacks. Normal request caching is unchanged. With `PI_LIMITS_WAIT_DISABLE_ALL_WAITING=true`, Pi's cache-warming policy is left untouched. As with any handler for this event, a later extension can override the decision.
 - Reload-safe ownership ensures `/reload` transfers interception cleanly and aborts waits still owned by the previous extension instance. The interception has one process-global active extension owner, matching the interactive Pi CLI's single active session runtime lifecycle; SDK processes running multiple concurrent Pi sessions are not supported.
 - On rate-limit errors (`429`, `rate_limit`, `too many requests`, quota/reset messages), waits and retries in a loop.
-- Remembers provider retry timing exposed through response callbacks or provider errors (`retry-after`, `retry-after-ms`, `retry in ...`, reset messages). With fallback probing enabled, explicit Retry-After deadlines may intentionally be probed early; set the probe interval to `0` to honor those deadlines (unless a wait is manually skipped).
+- Remembers provider retry timing exposed through response callbacks or provider errors (`retry-after`, `retry-after-ms`, `retry in ...`, reset messages). Explicit provider retry/reset deadlines are always honored: no early probes occur during that period (unless a wait is manually skipped).
 - 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.
@@ -30,7 +30,7 @@ Restart Pi or run `/reload` after installing.
 - During a wait, pressing Enter in an empty editor retries immediately. Typed commands, autocomplete, and model-picker navigation still work normally. Choosing another model with `/model` immediately ends the pending wait and tries that exact model, even if it was previously marked unavailable.
 - Keeps a manually chosen model ahead of the configured priority list until it fails. After failure, normal retry/fallback rules apply and subsequent calls use the configured order again.
 - Optionally falls back to configured models when the current/default model is rate-limited. Every fallback to a different model—even on the same provider—deliberately discards request-owned credentials, headers, environment overrides, and provider-specific options so the target runtime resolves its own configuration.
-- Preserves Pi's `before_provider_request` / `after_provider_response` callback flow and sanitises only the exact Anthropic OAuth system-block shape generated by current Pi; Pi remains responsible for OAuth identity insertion.
+- Preserves Pi's `before_provider_request` / `after_provider_response` callback flow. When Pi's exact Anthropic OAuth identity block is present, sanitises both the initial prompt and later system-message updates, repeating the Claude Code identity in textual updates. User/assistant messages and tool blocks stay untouched. Pi remains responsible for initial OAuth identity insertion and authentication; this does not repair invalid credentials.
 
 ### Unsupported interception limitation
 
@@ -234,7 +234,7 @@ Each `fallback-models` entry supports:
 
 Top-level settings keys:
 
-- `probe-interval-seconds` — optional. How often a limited/frozen higher-priority model is re-probed. Default `60`, `0` disables probing. `PI_LIMITS_WAIT_PROBE_INTERVAL` overrides it. Both sources accept only safe non-negative integers (numeric strings are also accepted); values above 2147483 seconds are capped. Invalid file values use the default, and invalid environment values use the file value or default.
+- `probe-interval-seconds` — optional. How often a higher-priority model with an estimated limit/freeze is silently re-probed. Explicit provider retry/reset deadlines are never probed early. Default `60`, `0` disables probing. `PI_LIMITS_WAIT_PROBE_INTERVAL` overrides it. Both sources accept only safe non-negative integers (numeric strings are also accepted); values above 2147483 seconds are capped. Invalid file values use the default, and invalid environment values use the file value or default.
 
 Fallback priority and manual selection:
 
@@ -245,10 +245,10 @@ Fallback priority and manual selection:
    - the remembered limit for that model expires;
    - the model answers a probe (see below);
    - successful output clears the limit/freeze known when its attempt began, but never a newer concurrent failure. Committed streams are never interrupted to reclaim a model.
-5. A still-limited higher-priority model is eligible for a probe every **`probe-interval-seconds`** (default 60s), on the next LLM call or while waiting with all candidates blocked. This intentionally permits early probes even before an explicit provider `Retry-After` deadline; `0` disables early probes and honors remembered deadlines unless manually skipped. A probe is one normal request sent to that model instead of the current fallback:
+5. A higher-priority model with an estimated limit/freeze (no explicit provider timing) is eligible for a probe every **`probe-interval-seconds`** (default 60s), on the next LLM call or while waiting with all candidates blocked. Explicit provider `Retry-After` and reset deadlines always block early probes, regardless of this interval; `0` disables early probes altogether. Manual wait skipping can still retry immediately. Probes and their recoverable failures are silent: no probe notices or retry/freeze warnings are shown. A probe is one normal request sent to that model instead of the current fallback:
    - if it succeeds, the model is reclaimed immediately and the request continues on it;
    - if it fails before output is committed, it costs just one attempt—even for HTTP 403, network or unknown errors—then falls back down the order and postpones the next probe by another interval. Abort and context-overflow errors remain terminal rather than triggering fallback. With freezing disabled, unknown/HTTP 403 probe failures only postpone probes of the existing limit (without extending its deadline or creating a freeze); that candidate is exhausted for the current request. Exhausted and not-yet-due blocked candidates are skipped, and if no eligible candidate remains the unknown error is surfaced rather than looping or waiting on a freeze.
-6. When *every* configured model is unavailable, the extension waits — but only until the earliest deadline **or the next due probe**, whichever comes first, so recovery is never delayed by a long remembered deadline.
+6. When *every* configured model is unavailable, the extension waits — but only until the earliest deadline **or the next eligible probe**, whichever comes first (never probing before an explicit provider deadline), so recovery is never delayed by a long remembered deadline.
 7. Rate-limit/freeze memory lives only in memory and is cleared when Pi restarts.
 
 The extension rereads these files before every LLM call, so changes made during a long-running task take effect on the next LLM call without restarting Pi or using `/reload`. Tool calls do not trigger a settings read. When the settings are loaded, the extension shows the full usable fallback model list. When models become rate-limited, it reports the wait in chat notifications.
@@ -263,7 +263,7 @@ The extension rereads these files before every LLM call, so changes made during
 | `PI_LIMITS_WAIT_LIVELINESS_INTERVAL` | `15` | Seconds between liveliness notifications published while waiting in non-interactive (`rpc`/`json`) modes. Must be a non-negative integer; values below 1 second are clamped to 1 second. Ignored in TUI mode. Also drives the structured status channel and is reported inside it as `livelinessIntervalMs`. |
 | `PI_LIMITS_WAIT_STATUS_JSON` | `true` | Whether the machine-readable status channel (`statusKey` `oira666.pi-limits-wait.json`) is published in non-interactive modes. Set to `false` (also accepts `0`, `no`, `off`) to restore exact 0.5.6 wire behavior. The prose channel is never affected. |
 | `PI_LIMITS_WAIT_CONTROL_FILE` | *(unset)* | Absolute path of the out-of-band control file polled once per second **while waiting**, used by the host to skip the current wait and retry immediately. Unset disables the feature and all filesystem access. See [Skipping a wait from the host](#skipping-a-wait-from-the-host). |
-| `PI_LIMITS_WAIT_PROBE_INTERVAL` | `60` | Seconds between re-probes of a higher-priority model that is still remembered as rate-limited or frozen. Lower values return to the preferred model sooner at the cost of more rejected probe requests; `0` disables probing, so a model is only reconsidered when its remembered deadline expires. Overrides the `probe-interval-seconds` settings key. Must be a safe non-negative integer; fractions, Infinity and unsafe values are rejected. Both file and environment values above 2147483 seconds are capped to Node's reliable timeout range. Invalid environment values use the file setting (or default). |
+| `PI_LIMITS_WAIT_PROBE_INTERVAL` | `60` | Seconds between silent re-probes of a higher-priority model with an estimated rate limit or freeze. Explicit provider retry/reset deadlines always block early probes. Lower values return to the preferred model sooner when no provider timing is known, at the cost of more rejected probe requests; `0` disables probing, so a model is only reconsidered when its remembered deadline expires. Overrides the `probe-interval-seconds` settings key. Must be a safe non-negative integer; fractions, Infinity and unsafe values are rejected. Both file and environment values above 2147483 seconds are capped to Node's reliable timeout range. Invalid environment values use the file setting (or default). |
 | `PI_LIMITS_WAIT_RETRY_INTERVAL` | `5` | Seconds to wait between unclassified/unknown provider-error retries. Must be a non-negative integer (values above 2147483 are capped to Node's maximum reliable timer delay). This does not affect model-limit wait intervals. |
 | `PI_LIMITS_WAIT_FREEZING_ENABLED` | `true` | When a model keeps failing with a non-retryable error, it is normally "frozen" for 10 minutes 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. |
 

package/package.json

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

package/src/constants.ts

@@ -25,8 +25,8 @@ export const DEFAULT_OVERLOADED_WAIT_MS = 5 * 60 * 1_000; // 5 minutes
 export const DEFAULT_NETWORK_WAIT_MS = 15 * 1_000; // 15 seconds
 export const DEFAULT_NON_RETRYABLE_FREEZE_MS = 10 * 60 * 1_000; // 10 minutes
 
-// A model that is currently limited/frozen is re-probed this often, so a
-// higher-priority model can be reclaimed long before its remembered deadline.
+// Estimated limits/freezes may be silently re-probed this often. Explicit
+// provider retry/reset deadlines always prevent early probes.
 export const DEFAULT_PROBE_INTERVAL_MS = 60 * 1_000; // 1 minute
 
 // Unknown provider errors are retried separately from rate-limit, overload,
@@ -43,7 +43,7 @@ export const DEFAULT_LIVELINESS_INTERVAL_MS = 15 * 1_000; // 15 seconds
 export const LIVELINESS_JSON_STATUS_KEY = "oira666.pi-limits-wait.json";
 /** Payload schema version for LIVELINESS_JSON_STATUS_KEY. Bump only on a breaking change. */
 export const LIVELINESS_JSON_VERSION = 1;
-export const EXTENSION_VERSION = "0.6.4";
+export const EXTENSION_VERSION = "0.6.6";
 
 export const SETTINGS_FILE_NAME = "limits-wait.json";
 export const FALLBACK_MODELS_KEY = "fallback-models";

package/src/index.ts

@@ -100,9 +100,9 @@ export type { FallbackModel } from "./types.js";
 
 export default function (pi: ExtensionAPI) {
   // Current pi-ai inserts the Claude Code identity as the first system block
-  // for Anthropic OAuth. Sanitise only that exact payload shape; authentication,
-  // identity insertion, Anthropic headers, and all other request composition
-  // remain owned by ModelRuntime.
+  // for Anthropic OAuth. Use that exact shape to guard sanitisation of the
+  // initial prompt and later system updates, which repeat the identity.
+  // Authentication, initial identity insertion, and headers remain runtime-owned.
   pi.on("before_provider_request", (event) => sanitiseAnthropicPayloadSystem(event.payload));
 
   // Keep prompt sanitisation and the runtime's normal Anthropic header path,

package/src/models.ts

@@ -1,6 +1,7 @@
 import type { Api, Model } from "@earendil-works/pi-ai";
 import { DEFAULT_NON_RETRYABLE_FREEZE_MS } from "./constants.js";
 import { expectModelSelection, state } from "./state.js";
+import { parseRetryDelayMs } from "./retry-errors.js";
 import type { FallbackModel, RetryableError } from "./types.js";
 import { formatErrorDetail, formatDuration, notifyLiveliness, reasonLabel } from "./ui.js";
 
@@ -82,19 +83,13 @@ export function notifyRetryingAfterError(model: Model<Api>, waitMs: number, erro
   notifyLiveliness(`${formatModel(model)} retrying after error in ${formatDuration(waitMs)}. Error: ${formatErrorDetail(errorMessage, Number.MAX_SAFE_INTEGER)}`);
 }
 
-function notifyProbe(model: Model<Api>, remainingMs: number): void {
-  const ctx = state.sharedCtx;
-  if (!ctx) return;
-  const remaining = remainingMs > 0 ? ` (${formatDuration(remainingMs)} left on its remembered limit)` : "";
-  try {
-    ctx.ui.notify(`Probing higher-priority model ${formatModel(model)}${remaining}.`, "info");
-  } catch { /* UI unavailable */ }
-}
-
-export function rememberRateLimit(model: Model<Api>, retryable: RetryableError, errorMessage?: string): void {
+export function rememberRateLimit(model: Model<Api>, retryable: RetryableError, errorMessage?: string, silent = false): void {
   const deadline = Date.now() + retryable.waitMs;
-  state.rateLimitMemory.set(modelKey(model), { reason: retryable.reason, limitedAt: Date.now(), deadline });
-  notifyRetryableError(model, retryable, errorMessage);
+  state.rateLimitMemory.set(modelKey(model), {
+    reason: retryable.reason, limitedAt: Date.now(), deadline,
+    explicitDeadline: errorMessage !== undefined && parseRetryDelayMs(errorMessage) !== undefined,
+  });
+  if (!silent) notifyRetryableError(model, retryable, errorMessage);
   ensureRateLimitedModelsStatus();
 }
 
@@ -118,9 +113,8 @@ export function probeIntervalMs(): number {
 /**
  * Remembered unusability of a model, merged across the rate-limit and frozen
  * memories, plus the moment it becomes worth re-probing. Remembered deadlines
- * are estimates (`retry-after` is often pessimistic, or absent and defaulted),
- * so a blocked model is re-probed periodically instead of being trusted until
- * its deadline.
+ * without provider timing are estimates and may be re-probed periodically.
+ * Explicit provider retry/reset deadlines must expire before any early probe.
  */
 export function blockedState(model: Model<Api>): { deadline: number; probeAt: number } | undefined {
   const limit = activeLimit(model);
@@ -134,7 +128,8 @@ export function blockedState(model: Model<Api>): { deadline: number; probeAt: nu
     frozen ? frozen.lastProbeAt ?? frozen.failedAt : 0,
   );
   const interval = probeIntervalMs();
-  const probeAt = interval > 0 ? Math.min(lastEvidence + interval, deadline) : deadline;
+  const estimatedProbeAt = interval > 0 ? Math.min(lastEvidence + interval, deadline) : deadline;
+  const probeAt = Math.max(limit?.explicitDeadline ? limit.deadline : 0, estimatedProbeAt);
   return { deadline, probeAt };
 }
 
@@ -158,9 +153,7 @@ export function clearUnavailability(model: Model<Api>, snapshot: ReturnType<type
 
 /** Record that a blocked model was just re-probed, so probes stay rate-limited. */
 export function markProbeAttempt(model: Model<Api>, now = Date.now()): void {
-  const remainingMs = (blockedState(model)?.deadline ?? now) - now;
   postponeProbe(model, now);
-  notifyProbe(model, remainingMs);
 }
 
 /** Postpone probing existing evidence without creating or extending a freeze. */
@@ -198,10 +191,10 @@ export function nextAvailableCandidate(current: Model<Api>, excluded?: ReadonlyS
   return nextAttemptCandidate(current, excluded)?.entry;
 }
 
-export function rememberNonRetryableFailure(model: Model<Api>, errorMessage: string): void {
+export function rememberNonRetryableFailure(model: Model<Api>, errorMessage: string, silent = false): void {
   const deadline = Date.now() + DEFAULT_NON_RETRYABLE_FREEZE_MS;
   state.nonRetryableFailureMemory.set(modelKey(model), { failedAt: Date.now(), deadline, errorMessage });
-  notifyLiveliness(`${formatModel(model)} failed (${formatErrorDetail(errorMessage)}); freezing it for ${formatDuration(DEFAULT_NON_RETRYABLE_FREEZE_MS)} and trying another configured model if available.`);
+  if (!silent) notifyLiveliness(`${formatModel(model)} failed (${formatErrorDetail(errorMessage)}); freezing it for ${formatDuration(DEFAULT_NON_RETRYABLE_FREEZE_MS)} and trying another configured model if available.`);
   ensureRateLimitedModelsStatus();
 }
 

package/src/prompt.ts

@@ -34,7 +34,8 @@ function isPiClaudeCodeIdentityBlock(block: unknown): block is SystemTextBlock {
  * current pi-ai runtime: the Claude Code identity is the first system block.
  * Looking only at that exact first block avoids treating a later user-authored
  * identity mention as proof of OAuth. The runtime remains solely responsible
- * for inserting the identity and selecting OAuth headers.
+ * for inserting the initial identity and selecting OAuth headers. Later system
+ * updates are sanitised too and repeat the identity alongside their instructions.
  */
 export function sanitiseAnthropicPayloadSystem(payload: unknown): unknown {
   if (!payload || typeof payload !== "object") return payload;
@@ -60,5 +61,50 @@ export function sanitiseAnthropicPayloadSystem(payload: unknown): unknown {
     }
   }
 
-  return changed ? { ...payload, system: newSystem } : payload;
+  // Pi 0.86+ can send prompt updates as system messages in the transcript.
+  // Only transform their text; tool additions/removals and conversation content
+  // must remain unchanged. Keep the OAuth guard above for every transformation.
+  const messages = (payload as { messages?: unknown }).messages;
+  const newMessages = Array.isArray(messages) ? messages.map((message: unknown) => {
+    if (!message || typeof message !== "object") return message;
+    const candidate = message as { role?: unknown; content?: unknown };
+    if (candidate.role !== "system") return message;
+
+    if (typeof candidate.content === "string") {
+      const cleaned = sanitiseSystemPrompt(candidate.content);
+      const content = cleaned ? `${CLAUDE_CODE_IDENTITY}\n\n${cleaned}` : CLAUDE_CODE_IDENTITY;
+      if (content === candidate.content) return message;
+      changed = true;
+      return { ...message, content };
+    }
+    if (!Array.isArray(candidate.content)) return message;
+
+    let messageChanged = false;
+    let hasIdentity = false;
+    const content = candidate.content.map((block: unknown) => {
+      if (!block || typeof block !== "object") return block;
+      const textBlock = block as SystemTextBlock;
+      if (textBlock.type !== "text" || typeof textBlock.text !== "string") return block;
+      const cleaned = sanitiseSystemPrompt(textBlock.text);
+      const text = hasIdentity ? cleaned
+        : cleaned ? `${CLAUDE_CODE_IDENTITY}\n\n${cleaned}` : CLAUDE_CODE_IDENTITY;
+      hasIdentity = true;
+      if (text === textBlock.text) return block;
+      messageChanged = true;
+      return { ...textBlock, text };
+    }).filter((block: unknown) => {
+      if (!block || typeof block !== "object") return true;
+      const candidate = block as SystemTextBlock;
+      return candidate.type !== "text" || candidate.text !== "";
+    });
+    if (!messageChanged) return message;
+    changed = true;
+    return { ...message, content };
+  }) : undefined;
+
+  return changed ? {
+    ...payload,
+    system: newSystem,
+    ...(newMessages ? { messages: newMessages } : {}),
+  } : payload;
 }

package/src/stream.ts

@@ -611,8 +611,8 @@ export function streamWithLimitsRetry(
         // classification. Abort and context overflow have already terminated.
         const failure = retryableErrorMessage || nonRetryableError?.message || "Provider stream ended without a terminal event.";
         if (retryable || freezingEnabled()) {
-          if (retryable) rememberRateLimit(attempt.model, retryable, failure);
-          else rememberNonRetryableFailure(attempt.model, failure);
+          if (retryable) rememberRateLimit(attempt.model, retryable, failure, true);
+          else rememberNonRetryableFailure(attempt.model, failure, true);
           retryPeriod?.recordImmediateRetry(failure);
           attempt = nextAvailableCandidate(attempt.model) ?? attempt;
           continue;

package/src/types.ts

@@ -30,6 +30,8 @@ export type RateLimitMemory = {
   reason: RetryReason;
   limitedAt: number;
   deadline: number;
+  /** Provider supplied retry/reset timing; never probe before this deadline. */
+  explicitDeadline?: boolean;
   /** Last time this model was re-probed while still remembered as limited. */
   lastProbeAt?: number;
 };