@kylebrodeur/pi-model-router 0.1.2 → 0.1.3

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/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.3",
   "type": "module",
   "description": "Intelligent per-turn model router extension for the pi coding agent (Enhanced Fork)",
   "keywords": [