oira666_pi-limits-wait 0.6.3 → 0.6.5

package/README.md

@@ -27,9 +27,10 @@ Restart Pi or run `/reload` after installing.
 - For any other (unclassified) provider error — including an API `Internal server error` and HTTP `400`/`403` failures — waits and retries the same model. There is deliberately no permanent-HTTP-status allowlist. This retry-all behavior is enabled by default and uses a configurable retry limit and interval; it does not alter model-limit wait intervals.
 - Reports retry/fallback waits in chat notifications without adding persistent TUI status lines.
 - In non-interactive hosts (`rpc` and `json` modes) the extension never reports a planned retry as an error. While it is waiting to retry, it publishes `info`-level liveliness notifications (plus a `setStatus` entry) describing what it is doing, how long it will wait, and why. An `error`-level notification is emitted only when the extension stops retrying and surfaces the terminal failure. See [Non-interactive (RPC) reporting](#non-interactive-rpc-reporting).
-- Leaves terminal input to Pi, so built-in selectors remain usable during a wait. Choosing another model with `/model` immediately ends the pending wait and continues with the newly selected provider/model.
+- 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
 
@@ -235,10 +236,10 @@ 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.
 
-Fallback behavior — strict priority, highest first:
+Fallback priority and manual selection:
 
-1. When the default/user-selected model is listed, `fallback-models` is authoritative, top to bottom—even if that selected model is lower in the list. A selection **outside** the list stays first; only on its failure does the configured fallback list apply. Unrelated model calls and internal synthetic models are not rerouted.
-2. Every LLM call restarts from the top of that order. The highest-priority model that is not known to be unavailable is used, so the first model runs whenever it can, and a lower-priority model is used **only** while every model above it is unavailable — and only for as long as that stays true.
+1. Normally, `fallback-models` sets priority from top to bottom. An unlisted startup model is tried before the list. Unrelated model calls and internal synthetic models are not rerouted.
+2. A manual model change overrides that priority, whether the chosen model is listed or not. Calls stay on that model without probing higher-priority models while it works. If it fails, the override ends and normal retry/fallback rules apply. After failover, the configured list controls priority again; recovery of the manually chosen model does not restore its override. Abort and context-overflow recovery do not cancel the override. Without an override, each LLM call starts at the top of the order and uses the highest-priority available model.
 3. If the active model becomes rate-limited, the extension immediately moves down to the next usable model in the order.
 4. As soon as a higher-priority model works again, the next call goes back to it and Pi's active model is switched back automatically. That happens in three ways:
    - the remembered limit for that model expires;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-limits-wait",
-  "version": "0.6.3",
+  "version": "0.6.5",
   "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
   "type": "module",
   "exports": {
@@ -18,7 +18,7 @@
   },
   "repository": {
     "type": "git",
-    "url": "https://github.com/gee666/pi-limits-wait.git"
+    "url": "git+https://github.com/gee666/pi-limits-wait.git"
   },
   "bugs": {
     "url": "https://github.com/gee666/pi-limits-wait/issues"
@@ -29,7 +29,7 @@
   },
   "scripts": {
     "check": "tsc -p tsconfig.json --noEmit",
-    "test": "tsx src/__tests__/test-extension.ts",
+    "test": "tsx src/__tests__/test-extension.ts && tsx src/__tests__/test-wait-ui.ts",
     "pack:check": "npm pack --dry-run",
     "prepublishOnly": "npm run check && npm test && npm run pack:check"
   },

package/src/constants.ts

@@ -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.2";
+export const EXTENSION_VERSION = "0.6.5";
 
 export const SETTINGS_FILE_NAME = "limits-wait.json";
 export const FALLBACK_MODELS_KEY = "fallback-models";

package/src/index.ts

@@ -1,9 +1,9 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { isInternalSyntheticModel } from "./models.js";
+import { isInternalSyntheticModel, modelKey } from "./models.js";
 import { sanitiseAnthropicPayloadSystem } from "./prompt.js";
 import { RetrySummaryCoordinator } from "./retry-summary.js";
 import { loadFallbackSettings } from "./settings.js";
-import { allWaitingDisabled, loadUnknownErrorRetrySettings } from "./ui.js";
+import { allWaitingDisabled, installWaitRetryEditor, loadUnknownErrorRetrySettings } from "./ui.js";
 import { consumeExpectedModelSelection, state } from "./state.js";
 import { disableModelRuntimeInterception, installModelRuntimeInterception } from "./stream.js";
 
@@ -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,
@@ -125,11 +125,23 @@ export default function (pi: ExtensionAPI) {
     on(event: "cache_warming_decision", handler: () => { action: "stop" }): void;
   }).on("cache_warming_decision", () => ({ action: "stop" }));
 
+  pi.on("session_start", (event, ctx) => {
+    state.sharedCtx = ctx;
+    installWaitRetryEditor(ctx, event.reason);
+  });
+
   pi.on("model_select", (event, ctx) => {
     state.sharedCtx = ctx;
     if (consumeExpectedModelSelection(event.model)) return;
     if (isInternalSyntheticModel(event.model)) return;
     state.primaryModel = event.model;
+    state.manualModelOverride = event.source === "restore" ? undefined : event.model;
+    state.manualModelOverrideReleased = false;
+    // An explicit choice deserves a fresh attempt, even if previously blocked.
+    if (state.manualModelOverride) {
+      state.rateLimitMemory.delete(modelKey(event.model));
+      state.nonRetryableFailureMemory.delete(modelKey(event.model));
+    }
     state.primaryThinkingLevel = pi.getThinkingLevel();
     state.userModelSelectionGeneration++;
     // A user-selected model takes precedence over a pending retry. Wake every

package/src/models.ts

@@ -38,14 +38,16 @@ export function configuredAttempt(model: Model<Api>): FallbackModel {
 }
 
 /**
- * Listed models are authoritative, top to bottom. An explicitly selected
- * model outside the list stays first; the configured list is used only if
- * that selection fails. Unrelated calls are excluded by isFallbackEligibleModel.
+ * A manual choice stays first until it fails. Otherwise use configured priority.
+ * An unlisted startup model stays first until an explicit selection has failed.
+ * Unrelated calls are excluded by isFallbackEligibleModel.
  */
 export function candidateOrder(current: Model<Api>): FallbackModel[] {
   const primary = getPrimaryModel(current);
   const listed = state.fallbackModels.some((entry) => modelKey(entry.model) === modelKey(primary));
-  const order: FallbackModel[] = listed ? [] : [configuredAttempt(primary)];
+  const order: FallbackModel[] = state.manualModelOverride
+    ? [configuredAttempt(state.manualModelOverride)]
+    : listed || state.manualModelOverrideReleased ? [] : [configuredAttempt(primary)];
   const seen = new Set(order.map((entry) => modelKey(entry.model)));
   for (const entry of state.fallbackModels) {
     const key = modelKey(entry.model);

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/state.ts

@@ -21,6 +21,8 @@ export const state = {
   probeIntervalMs: DEFAULT_PROBE_INTERVAL_MS,
   primaryModel: undefined as Model<Api> | undefined,
   primaryThinkingLevel: undefined as ModelThinkingLevel | undefined,
+  manualModelOverride: undefined as Model<Api> | undefined,
+  manualModelOverrideReleased: false,
   expectedModelSelections: new Map<string, Set<symbol>>(),
   userModelSelectionGeneration: 0,
   activeWaitSkips: new Set<() => void>(),
@@ -62,6 +64,8 @@ export function resetRuntimeStateForTests(models: FallbackModel[], ctx?: Extensi
   state.sharedCtx = ctx;
   state.primaryModel = undefined;
   state.primaryThinkingLevel = undefined;
+  state.manualModelOverride = undefined;
+  state.manualModelOverrideReleased = false;
   state.rateLimitMemory.clear();
   state.nonRetryableFailureMemory.clear();
   state.expectedModelSelections.clear();

package/src/stream.ts

@@ -64,6 +64,8 @@ export function __configureFallbackModelsForTests(models: FallbackModel[], ctx?:
   state.sharedCtx = ctx;
   state.primaryModel = undefined;
   state.primaryThinkingLevel = undefined;
+  state.manualModelOverride = undefined;
+  state.manualModelOverrideReleased = false;
   state.rateLimitMemory.clear();
   state.nonRetryableFailureMemory.clear();
   state.ambientStatusCleanup?.();
@@ -217,6 +219,15 @@ export function streamWithLimitsRetry(
     const earliestCandidateWakeup = (current: Model<Api>) => findCandidateWakeup(current, exclusions());
     let observedUserModelSelection = state.userModelSelectionGeneration;
     let forceAttemptAfterSkip = false;
+    const releaseManualOverride = () => {
+      // A stale failure must not revoke a newer choice, even of the same model.
+      if (observedUserModelSelection === state.userModelSelectionGeneration
+        && state.manualModelOverride
+        && modelKey(state.manualModelOverride) === modelKey(attempt.model)) {
+        state.manualModelOverride = undefined;
+        state.manualModelOverrideReleased = true;
+      }
+    };
 
     const prepareAttemptForPi = async (): Promise<boolean> => {
       const originalKey = modelKey(model);
@@ -279,9 +290,10 @@ export function streamWithLimitsRetry(
         observedUserModelSelection = state.userModelSelectionGeneration;
         const selected = state.primaryModel ?? state.sharedCtx?.model;
         if (selected) {
-          // New selections supersede the old attempt. Listed selections still
-          // obey configured priority; an out-of-list selection stays first.
-          attempt = allowFallback ? initialAttempt(selected) ?? configuredAttempt(selected) : configuredAttempt(selected);
+          // Try the user's exact choice, not the first configured candidate.
+          attempt = configuredAttempt(selected);
+          forceAttemptAfterSkip = true;
+          nonRetryableAttempts.clear();
           triedModels.clear();
           lastExhaustedFailure = undefined;
         }
@@ -503,7 +515,10 @@ export function streamWithLimitsRetry(
               continue;
             }
 
-            if (event.type === "error") retryPeriod?.recordReason(event.error.errorMessage ?? "Provider stream error");
+            if (event.type === "error") {
+              if (event.reason !== "aborted") releaseManualOverride();
+              retryPeriod?.recordReason(event.error.errorMessage ?? "Provider stream error");
+            }
             if (event.type === "done" || event.type === "error") {
               retryPeriod?.finalize();
               clearLivelinessStatus();
@@ -587,6 +602,10 @@ export function streamWithLimitsRetry(
         return;
       }
 
+      // A superseded attempt cannot freeze or limit the user's newer choice.
+      if (observedUserModelSelection !== state.userModelSelectionGeneration) continue;
+      releaseManualOverride();
+
       if (probe && !committed) {
         // A probe is exactly one request, regardless of HTTP/network/unknown
         // classification. Abort and context overflow have already terminated.

package/src/ui.ts

@@ -1,3 +1,4 @@
+import { CustomEditor, type ExtensionContext, type SessionStartEvent } from "@earendil-works/pi-coding-agent";
 import { randomUUID } from "node:crypto";
 import { readFile, rename, unlink } from "node:fs/promises";
 import {
@@ -35,6 +36,51 @@ import {
 } from "./telemetry.js";
 import type { RetryReason } from "./types.js";
 
+/**
+ * Install at session_start, never when a wait starts or ends: replacing the
+ * editor also moves focus and would otherwise dismiss an active picker.
+ * Intercept submissions, not terminal keys, so autocomplete and other focused
+ * components get to handle Enter first. Preserve an existing editor factory.
+ */
+export function installWaitRetryEditor(
+  ctx: ExtensionContext,
+  reason: SessionStartEvent["reason"] = "startup",
+): void {
+  if (ctx.mode !== "tui") return;
+  const previous = ctx.ui.getEditorComponent();
+  ctx.ui.setEditorComponent((tui, theme, keybindings) => {
+    const editor = previous?.(tui, theme, keybindings) ?? new CustomEditor(tui, theme, keybindings);
+    // Pi resets custom editors before reload and only transfers their text.
+    // Unlike startup/session switching, reload does not populate history after
+    // session_start. Seed our new default editor from the active branch only
+    // on reload; existing custom editors remain responsible for their history.
+    if (!previous && reason === "reload") {
+      for (const entry of ctx.sessionManager.getBranch()) {
+        if (entry.type !== "message" || entry.message.role !== "user") continue;
+        const content = entry.message.content;
+        const text = typeof content === "string"
+          ? content
+          : content.filter((block) => block.type === "text").map((block) => block.text).join("");
+        if (text) editor.addToHistory?.(text);
+      }
+    }
+    let onSubmit = editor.onSubmit;
+    // Pi assigns its submit callback after the factory returns.
+    Object.defineProperty(editor, "onSubmit", {
+      configurable: true,
+      get: () => (text: string) => {
+        if (!text.trim() && state.activeWaitSkips.size > 0) {
+          for (const skip of [...state.activeWaitSkips]) skip();
+          return;
+        }
+        onSubmit?.call(editor, text);
+      },
+      set: (handler: typeof editor.onSubmit) => { onSubmit = handler; },
+    });
+    return editor;
+  });
+}
+
 export function formatDuration(ms: number): string {
   const totalSecs = Math.max(0, Math.ceil(ms / 1_000));
   const hours = Math.floor(totalSecs / 3_600);
@@ -452,10 +498,9 @@ export function waitForRetry(
       return;
     }
 
-    // Raw terminal listeners run before whichever component currently owns
-    // focus. Consuming Enter here would steal confirmation from built-in
-    // selectors such as /model. Model changes and the out-of-band control
-    // channel call this skip function directly instead.
+    // The session editor's empty-submit handler, model changes, and the
+    // out-of-band control channel share this skip function. No raw input hook:
+    // it would steal Enter from whichever picker currently owns focus.
     state.activeWaitSkips.add(skipWait);
 
     signal?.addEventListener("abort", onAbort);