@kylebrodeur/pi-model-router 0.1.2 → 0.1.4

package/CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- Transparent wait and retry interception for string-based rate limit errors (e.g., "quota will reset after X seconds")
 - Ollama auto-sync feature
 - Rate-limit fallback with transparent HTTP error handling (402, 429, 503, 529)
 - Feature toggles in config (`features` object)

package/LEARNINGS.md

@@ -73,10 +73,11 @@ The fallback mechanism uses a user-configurable sequence of models: `fallbackSeq
 *   **Key benefit**: Prevents catastrophic failures when a primary model is unavailable.
 
 ### 3. Graceful Error Handling
-The extension transparently handles errors. For "out of credits" (`402`) or "rate limit" (`429`), it automatically switches to a fallback model and emits a custom session entry (`router-fallback`) for headless tooling to detect.
+The extension transparently handles errors. For "out of credits" (`402`) or "rate limit" (`429`), it automatically switches to a fallback model and emits a custom session entry (`router-fallback`) for headless tooling to detect. 
+Additionally, for string-based 429 errors specifying a cooldown (e.g., "quota will reset after 58s"), the router can intercept the stream, pause for the required duration (if under `shortDelayThreshold`), and automatically retry the original request without failing the turn.
 
 *   **When to use**: For any extension exposed to external API services.
-*   **Key insight**: Never mask API errors; provide enough detail (status codes) in UI notifications for users to diagnose.
+*   **Key insight**: Never mask API errors; provide enough detail (status codes) in UI notifications for users to diagnose, but handle transient issues (like short rate limits) invisibly where possible.
 
 ## 🔌 Pi Integration Patterns
 

package/README.md

@@ -121,6 +121,30 @@ Copy the example config to one of:
 
 **Priority:** Project config `.pi/model-router.json` overrides user config `~/.pi/agent/model-router.json`. Both override defaults.
 
+### Rate Limit Interception & Fallback
+
+The router can gracefully handle 429 Rate Limit and Quota errors. If the error specifies a wait time (e.g., "reset after 58s"), the router will pause and automatically retry the prompt if the wait time is under your threshold. If it exceeds the threshold or is unparseable, it fails over to the next available model in your fallback sequence.
+
+```json
+{
+  "rateLimitFallback": {
+    "enabled": true,
+    "shortDelayThreshold": 60,
+    "autoFallback": true,
+    "autoRestore": true,
+    "restoreCheckInterval": 300,
+    "fallbackSequence": ["anthropic/claude-3-haiku-20240307", "ollama/*"]
+  }
+}
+```
+
+| Field | Description |
+|-------|-------------|
+| `shortDelayThreshold` | Maximum time (in seconds) the router will pause and wait to retry when encountering a rate limit. If the cooldown is longer than this, it triggers a fallback. |
+| `fallbackSequence` | Array of model IDs (or wildcards like `ollama/*`) to try if the primary model fails or the wait time is too long. |
+| `autoFallback` | (Optional) Automatically switch session to the fallback model globally after a hard failure. |
+| `autoRestore` | (Optional) If fallback was triggered, automatically try to restore the original cloud model after `restoreCheckInterval` seconds. |
+
 ### Progressive Enhancement Configs
 
 After installing optional extensions, copy one of these to `.pi/model-router.json`:

package/extensions/index.ts

@@ -25,7 +25,7 @@ import { registerCommands } from './commands';
 import { registerRouterProvider } from './provider';
 // ─── Feature modules (added by fork) ────────────────────────────────────────
 import { initializeOllamaSync } from './ollama-sync';
-import { initializeRateLimitFallback } from './rate-limit';
+import { initializeRateLimitFallback, checkAndRestore } from './rate-limit';
 
 // ─── Plugin Detection & Progressive Integration ──────────────────────────
 interface PluginStatus {
@@ -35,7 +35,7 @@ interface PluginStatus {
 
 const detectPlugins = (pi: ExtensionAPI): PluginStatus => {
   const tools = (pi as any).tools ?? {};
-  const log = (pi as any).log;
+  const log = (pi as any).log || console;
   return {
     ledger: typeof tools.append_ledger === 'function',
     agentBus: typeof tools.link_send === 'function',
@@ -48,7 +48,7 @@ const detectAndIntegratePlugins = (
   debugEnabled: boolean,
 ) => {
   const plugins = detectPlugins(pi);
-  const log = (pi as any).log;
+  const log = (pi as any).log || console;
 
   // Ledger integration: log routing decisions to qmd-ledger
   const shouldIntegrateLedger = features?.ledgerIntegration === true;
@@ -575,6 +575,28 @@ const routerExtension = (pi: ExtensionAPI) => {
         await setModelInternally(routerModel);
       }
     }
+
+    // Auto-restore from rate-limit fallback
+    const rateLimitCfg = (currentConfig.rateLimitFallback ?? {}) as Record<
+      string,
+      unknown
+    >;
+    if (rateLimitCfg.autoRestore === true) {
+      const result = await checkAndRestore(
+        pi,
+        ctx,
+        currentConfig.features?.contextCompression === true,
+        (rateLimitCfg.restoreCheckInterval as number) ?? 300,
+      );
+      if (result.attempted && result.success) {
+        ctx.ui.notify(`[Router] Auto-restored: ${result.message}`, 'info');
+        pi.appendEntry('router-auto-restore', {
+          restored: true,
+          message: result.message,
+        });
+      }
+    }
+
     persistState();
     actions.updateStatus(ctx);
   });

package/extensions/provider.ts

@@ -30,6 +30,20 @@ import {
   hasImageAttachment,
 } from './routing';
 
+const rateLimitRegex = /(?:429|rate limit|quota).*?(?:reset after|try again in|wait)\s*(\d+)\s*([smh])/i;
+
+function extractWaitTimeMs(errorText: string): number | null {
+  const match = errorText.match(rateLimitRegex);
+  if (!match) return null;
+  const value = parseInt(match[1], 10);
+  const unit = match[2].toLowerCase();
+  
+  if (unit === 's') return value * 1000;
+  if (unit === 'm') return value * 60000;
+  if (unit === 'h') return value * 3600000;
+  return null;
+}
+
 export const createErrorMessage = (
   model: Model<Api>,
   message: string,
@@ -457,74 +471,109 @@ export const registerRouterProvider = (
             const apiKey = auth.apiKey;
             const headers = auth.headers;
 
-            try {
-              // HONESTY CHECK & AUTO-TRUNCATION
-              // If the picked model has a smaller context than what we reported, truncate now.
-              let effectiveContext = context;
-              const targetLimit = targetModel.contextWindow || 128_000;
-              if (targetLimit < model.contextWindow!) {
-                effectiveContext = truncateContext(context, targetLimit);
-              }
+            let retryCount = 0;
+            let modelSuccess = false;
 
-              const thinkingOverride = actions.getThinkingOverride(
-                model.id,
-                decision.tier,
-              );
-              const delegatedReasoning =
-                targetModel.reasoning &&
-                (thinkingOverride ?? decision.thinking) !== 'off'
-                  ? (thinkingOverride ?? decision.thinking)
-                  : undefined;
-
-              if (state.lastExtensionContext) {
-                if (delegatedReasoning) {
-                  state.lastExtensionContext.ui.setHiddenThinkingLabel?.(
-                    `Thinking (${targetProvider}/${targetModelId})...`,
-                  );
-                } else {
-                  state.lastExtensionContext.ui.setHiddenThinkingLabel?.();
+            while (retryCount < 2) {
+              let contentReceived = false;
+              try {
+                // HONESTY CHECK & AUTO-TRUNCATION
+                // If the picked model has a smaller context than what we reported, truncate now.
+                let effectiveContext = context;
+                const targetLimit = targetModel.contextWindow || 128_000;
+                if (targetLimit < model.contextWindow!) {
+                  effectiveContext = truncateContext(context, targetLimit);
                 }
-              }
 
-              const delegatedStream = streamSimple(
-                targetModel,
-                effectiveContext,
-                {
-                  ...options,
-                  apiKey,
-                  headers,
-                  ...(delegatedReasoning
-                    ? { reasoning: delegatedReasoning }
-                    : {}),
-                },
-              );
+                const thinkingOverride = actions.getThinkingOverride(
+                  model.id,
+                  decision.tier,
+                );
+                const delegatedReasoning =
+                  targetModel.reasoning &&
+                  (thinkingOverride ?? decision.thinking) !== 'off'
+                    ? (thinkingOverride ?? decision.thinking)
+                    : undefined;
+
+                if (state.lastExtensionContext) {
+                  if (delegatedReasoning) {
+                    state.lastExtensionContext.ui.setHiddenThinkingLabel?.(
+                      `Thinking (${targetProvider}/${targetModelId})...`,
+                    );
+                  } else {
+                    state.lastExtensionContext.ui.setHiddenThinkingLabel?.();
+                  }
+                }
 
-              let contentReceived = false;
-              for await (const event of delegatedStream) {
-                if (event.type === 'done') {
-                  const cost = event.message.usage?.cost?.total ?? 0;
-                  state.accumulatedCost += cost;
+                const delegatedStream = streamSimple(
+                  targetModel,
+                  effectiveContext,
+                  {
+                    ...options,
+                    apiKey,
+                    headers,
+                    ...(delegatedReasoning
+                      ? { reasoning: delegatedReasoning }
+                      : {}),
+                  },
+                );
+
+                for await (const event of delegatedStream) {
+                  if (event.type === 'done') {
+                    const cost = event.message.usage?.cost?.total ?? 0;
+                    state.accumulatedCost += cost;
+                  }
+                  if (event.type === 'error' && !contentReceived) {
+                    throw new Error(
+                      (event as any).error?.errorMessage ||
+                        'Model failed before sending content.',
+                    );
+                  }
+                  const isContent =
+                    event.type === 'text_delta' ||
+                    event.type === 'thinking_delta' ||
+                    event.type === 'toolcall_delta' ||
+                    event.type === 'toolcall_end';
+                  if (isContent) contentReceived = true;
+                  stream.push(event);
                 }
-                if (event.type === 'error' && !contentReceived) {
-                  throw new Error(
-                    (event as any).error?.errorMessage ||
-                      'Model failed before sending content.',
-                  );
+                modelSuccess = true;
+                success = true;
+                if (i > 0) decision.isFallback = true;
+                break; // break the retry loop
+              } catch (err) {
+                const errMsg = err instanceof Error ? err.message : String(err);
+                const waitMs = extractWaitTimeMs(errMsg);
+                const maxWaitMs = (state.currentConfig.rateLimitFallback?.shortDelayThreshold ?? 60) * 1000;
+
+                if (waitMs && waitMs <= maxWaitMs && retryCount === 0 && !contentReceived) {
+                  const partialMsg = {
+                    role: 'assistant',
+                    content: [],
+                    api: model.api,
+                    provider: targetProvider,
+                    model: targetModelId,
+                    usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+                    timestamp: Date.now(),
+                  } as unknown as AssistantMessage;
+
+                  stream.push({
+                    type: 'text_delta',
+                    contentIndex: 0,
+                    delta: `\n_⏳ [Router] Rate limit reached on ${targetProvider}/${targetModelId}. Waiting ${Math.ceil(waitMs/1000)}s before retrying..._\n`,
+                    partial: partialMsg
+                  });
+                  await new Promise(resolve => setTimeout(resolve, waitMs + 1000)); // buffer 1s
+                  retryCount++;
+                  continue; // try the same model again
                 }
-                const isContent =
-                  event.type === 'text_delta' ||
-                  event.type === 'thinking_delta' ||
-                  event.type === 'toolcall_delta' ||
-                  event.type === 'toolcall_end';
-                if (isContent) contentReceived = true;
-                stream.push(event);
+
+                lastError = err;
+                break; // model failed completely, break retry loop to go to next fallback model
               }
-              success = true;
-              if (i > 0) decision.isFallback = true;
-              break;
-            } catch (err) {
-              lastError = err;
             }
+
+            if (modelSuccess) break; // break fallback loop
           }
 
           if (!success) {

package/extensions/rate-limit.ts

@@ -28,12 +28,21 @@ export interface RateLimitEventEntry {
   httpStatus: number;
 }
 
+export interface ModelCapabilities {
+  vision: boolean;
+  reasoning: boolean;
+  contextWindow: number;
+  maxTokens: number;
+}
+
 export interface FallbackState {
   preferredModel?: string;
   fallbackActive: boolean;
   autoRestore: boolean;
   triggeredAt?: number;
   triggerReason?: 'rate_limit' | 'budget_exceeded' | 'manual';
+  lastRestoreAttempt?: number;
+  requiredCapabilities?: ModelCapabilities;
 }
 
 // ─── Config ─────────────────────────────────────────────────────────────────
@@ -58,26 +67,72 @@ let history: RateLimitEventEntry[] = [];
 
 // ─── Helpers ────────────────────────────────────────────────────────────────
 
+const getModelCapabilities = (model: {
+  input: string[];
+  reasoning: boolean;
+  contextWindow: number;
+  maxTokens: number;
+}): ModelCapabilities => ({
+  vision: model.input.includes('image'),
+  reasoning: model.reasoning,
+  contextWindow: model.contextWindow,
+  maxTokens: model.maxTokens,
+});
+
+const capabilitiesMatch = (
+  required: ModelCapabilities,
+  candidate: ModelCapabilities,
+): { match: boolean; missing: string[] } => {
+  const missing: string[] = [];
+  if (required.vision && !candidate.vision) missing.push('vision');
+  if (required.reasoning && !candidate.reasoning) missing.push('reasoning');
+  if (candidate.contextWindow < required.contextWindow)
+    missing.push(
+      `contextWindow ${candidate.contextWindow} < ${required.contextWindow}`,
+    );
+  if (candidate.maxTokens < required.maxTokens)
+    missing.push(`maxTokens ${candidate.maxTokens} < ${required.maxTokens}`);
+  return { match: missing.length === 0, missing };
+};
+
 const findBestFallbackModel = (
   ctx: ExtensionContext,
   sequence: string[],
-): { provider: string; id: string } | undefined => {
+  required?: ModelCapabilities,
+): { provider: string; id: string; missing?: string[] } | undefined => {
   const availableModels = ctx.modelRegistry.getAvailable();
 
+  let bestPartialMatch:
+    | { provider: string; id: string; missing: string[] }
+    | undefined;
+
   for (const pattern of sequence) {
     for (const model of availableModels) {
       const targetId = `${model.provider}/${model.id}`;
-      if (pattern === targetId)
-        return { provider: model.provider, id: model.id };
-      if (pattern.endsWith('*')) {
-        const prefix = pattern.slice(0, -1);
-        if (targetId.startsWith(prefix))
+      if (
+        pattern === targetId ||
+        (pattern.endsWith('*') && targetId.startsWith(pattern.slice(0, -1)))
+      ) {
+        if (required) {
+          const caps = getModelCapabilities(model);
+          const { match, missing } = capabilitiesMatch(required, caps);
+          if (match) {
+            return { provider: model.provider, id: model.id };
+          } else if (!bestPartialMatch) {
+            bestPartialMatch = {
+              provider: model.provider,
+              id: model.id,
+              missing,
+            };
+          }
+        } else {
           return { provider: model.provider, id: model.id };
+        }
       }
     }
   }
 
-  return undefined;
+  return bestPartialMatch;
 };
 
 // ─── Public API ─────────────────────────────────────────────────────────────
@@ -97,17 +152,26 @@ export const tryFallback = async (
 
   if (currentModel.provider !== 'ollama' && !state.fallbackActive) {
     state.preferredModel = `${currentModel.provider}/${currentModel.id}`;
+    state.requiredCapabilities = getModelCapabilities(currentModel);
   }
 
   const target = findBestFallbackModel(
     ctx,
     config.fallbackSequence.length > 0 ? config.fallbackSequence : ['ollama/*'],
+    state.requiredCapabilities,
   );
 
   if (!target) {
     return { success: false, message: 'No fallback models available' };
   }
 
+  if (target.missing) {
+    return {
+      success: false,
+      message: `Fallback model ${target.provider}/${target.id} lacks required capabilities: ${target.missing.join(', ')}`,
+    };
+  }
+
   const targetModel = ctx.modelRegistry.find(target.provider, target.id);
   if (!targetModel) {
     return {
@@ -153,9 +217,13 @@ export const tryRestore = async (
   pi: ExtensionAPI,
   ctx: ExtensionContext,
   contextCompressionEnabled: boolean = false,
-): Promise<{ success: boolean; message: string }> => {
+): Promise<{ success: boolean; message: string; restored: boolean }> => {
   if (!state.fallbackActive || !state.preferredModel) {
-    return { success: false, message: 'No preferred model stored' };
+    return {
+      success: false,
+      message: 'No preferred model stored',
+      restored: false,
+    };
   }
 
   const [provider, id] = state.preferredModel.split('/');
@@ -165,6 +233,7 @@ export const tryRestore = async (
     return {
       success: false,
       message: `Model ${state.preferredModel} not available`,
+      restored: false,
     };
   }
 
@@ -172,8 +241,9 @@ export const tryRestore = async (
   if (success) {
     state.fallbackActive = false;
     state.autoRestore = false;
+    state.lastRestoreAttempt = undefined;
+    state.requiredCapabilities = undefined;
 
-    // Context Compression Bridge: Instruct the model to summarize the fallback period
     if (contextCompressionEnabled) {
       pi.sendMessage(
         {
@@ -189,12 +259,47 @@ export const tryRestore = async (
 
   return {
     success,
+    restored: success,
     message: success
       ? `Restored ${state.preferredModel}`
       : 'Failed to restore model',
   };
 };
 
+/**
+ * Periodically check if the preferred cloud model is healthy and auto-restore.
+ * Call this from turn_end or another periodic hook.
+ */
+export const checkAndRestore = async (
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  contextCompressionEnabled: boolean = false,
+  restoreCheckIntervalSec: number = 300,
+): Promise<{ attempted: boolean; success: boolean; message: string }> => {
+  if (!state.autoRestore || !state.fallbackActive || !state.preferredModel) {
+    return {
+      attempted: false,
+      success: false,
+      message: 'Auto-restore not active',
+    };
+  }
+
+  const now = Date.now();
+  const intervalMs = restoreCheckIntervalSec * 1000;
+
+  if (state.lastRestoreAttempt && now - state.lastRestoreAttempt < intervalMs) {
+    return {
+      attempted: false,
+      success: false,
+      message: 'Restore throttled',
+    };
+  }
+
+  state.lastRestoreAttempt = now;
+  const result = await tryRestore(pi, ctx, contextCompressionEnabled);
+  return { attempted: true, ...result };
+};
+
 export const getFallbackState = (): FallbackState => {
   return { ...state };
 };
@@ -224,6 +329,8 @@ export const resetRateLimitState = (): void => {
   state = {
     fallbackActive: false,
     autoRestore: false,
+    lastRestoreAttempt: undefined,
+    requiredCapabilities: undefined,
   };
   history = [];
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@kylebrodeur/pi-model-router",
-  "version": "0.1.2",
+  "version": "0.1.4",
   "type": "module",
   "description": "Intelligent per-turn model router extension for the pi coding agent (Enhanced Fork)",
   "keywords": [