@alexeiled/pi-model-router 0.5.1 → 0.5.2

package/extensions/index.ts

@@ -21,7 +21,6 @@ import {
   saveLastRouterProfile,
 } from './state';
 import type {
-  CustomSessionEntry,
   RouterConfig,
   RouterPinByProfile,
   RouterThinkingByProfile,
@@ -58,6 +57,74 @@ const routerExtension = (pi: ExtensionAPI) => {
   let isInternalThinkingChange = false;
   let ignoreStartupThinkingEvent = false;
 
+  const runtimeState = {
+    get lastRegisteredModels() {
+      return lastRegisteredModels;
+    },
+    set lastRegisteredModels(value: string) {
+      lastRegisteredModels = value;
+    },
+    get currentConfig() {
+      return currentConfig;
+    },
+    get currentModelRegistry() {
+      return currentModelRegistry;
+    },
+    get lastExtensionContext() {
+      return lastExtensionContext;
+    },
+    get selectedProfile() {
+      return selectedProfile;
+    },
+    set selectedProfile(value: string | undefined) {
+      selectedProfile = value;
+    },
+    get routerEnabled() {
+      return routerEnabled;
+    },
+    set routerEnabled(value: boolean) {
+      routerEnabled = value;
+    },
+    get lastDecision() {
+      return lastDecision;
+    },
+    set lastDecision(value: RoutingDecision | undefined) {
+      lastDecision = value;
+    },
+    thinkingByProfile,
+    pinnedTierByProfile,
+    get accumulatedCost() {
+      return accumulatedCost;
+    },
+    set accumulatedCost(value: number) {
+      accumulatedCost = value;
+    },
+    get debugEnabled() {
+      return debugEnabled;
+    },
+    set debugEnabled(value: boolean) {
+      debugEnabled = value;
+    },
+    get widgetEnabled() {
+      return widgetEnabled;
+    },
+    set widgetEnabled(value: boolean) {
+      widgetEnabled = value;
+    },
+    get debugHistory() {
+      return debugHistory;
+    },
+    get lastNonRouterModel() {
+      return lastNonRouterModel;
+    },
+    set lastNonRouterModel(value: string | undefined) {
+      lastNonRouterModel = value;
+    },
+    get lastConfigWarnings() {
+      return lastConfigWarnings;
+    },
+  };
+
   const setModelInternally = async (
     model: NonNullable<ExtensionContext['model']>,
   ) => {
@@ -92,7 +159,7 @@ const routerExtension = (pi: ExtensionAPI) => {
   };
 
   const persistState = () => {
-    const state = buildPersistedState(
+    const state = buildPersistedState({
       routerEnabled,
       selectedProfile,
       pinnedTierByProfile,
@@ -103,7 +170,7 @@ const routerExtension = (pi: ExtensionAPI) => {
       lastDecision,
       lastNonRouterModel,
       accumulatedCost,
-    );
+    });
     const snapshot = JSON.stringify({
       ...state,
       timestamp: 0,
@@ -133,18 +200,16 @@ const routerExtension = (pi: ExtensionAPI) => {
     persistState,
     syncPiThinkingLevel: setThinkingLevelInternally,
     updateStatus: (ctx: ExtensionContext) =>
-      updateStatus(
-        ctx,
+      updateStatus(ctx, {
         routerEnabled,
         selectedProfile,
         pinnedTierByProfile,
-        thinkingByProfile,
         lastDecision,
         lastNonRouterModel,
         accumulatedCost,
         widgetEnabled,
         currentConfig,
-      ),
+      }),
     reloadConfig: (
       ctx?: ExtensionContext,
       options?: { preserveDebug?: boolean },
@@ -221,59 +286,13 @@ const routerExtension = (pi: ExtensionAPI) => {
       return true;
     },
     registerRouterProvider: () => {
-      registerRouterProvider(
-        pi,
-        {
-          get lastRegisteredModels() {
-            return lastRegisteredModels;
-          },
-          set lastRegisteredModels(v) {
-            lastRegisteredModels = v;
-          },
-          get currentConfig() {
-            return currentConfig;
-          },
-          get currentModelRegistry() {
-            return currentModelRegistry;
-          },
-          get lastExtensionContext() {
-            return lastExtensionContext;
-          },
-          get selectedProfile() {
-            return selectedProfile;
-          },
-          set selectedProfile(v) {
-            selectedProfile = v;
-          },
-          get routerEnabled() {
-            return routerEnabled;
-          },
-          set routerEnabled(v) {
-            routerEnabled = v;
-          },
-          get lastDecision() {
-            return lastDecision;
-          },
-          set lastDecision(v) {
-            lastDecision = v;
-          },
-          thinkingByProfile,
-          pinnedTierByProfile,
-          get accumulatedCost() {
-            return accumulatedCost;
-          },
-          set accumulatedCost(v) {
-            accumulatedCost = v;
-          },
-        },
-        {
-          persistState,
-          recordDebugDecision,
-          getThinkingOverride,
-          updateStatus: actions.updateStatus,
-          syncPiThinkingLevel: setThinkingLevelInternally,
-        },
-      );
+      registerRouterProvider(pi, runtimeState, {
+        persistState,
+        recordDebugDecision,
+        getThinkingOverride,
+        updateStatus: actions.updateStatus,
+        syncPiThinkingLevel: setThinkingLevelInternally,
+      });
     },
   };
 
@@ -318,16 +337,16 @@ const routerExtension = (pi: ExtensionAPI) => {
 
     await actions.ensureValidActiveRouterProfile(ctx);
 
-    const entries = ctx.sessionManager.getBranch() as CustomSessionEntry[];
-    const savedState = entries
-      .filter(
-        (entry) =>
-          entry.type === 'custom' && entry.customType === 'router-state',
+    const savedState = ctx.sessionManager
+      .getBranch()
+      .map((entry) =>
+        entry.type === 'custom' && entry.customType === 'router-state'
+          ? entry.data
+          : undefined,
       )
-      .map((entry) => entry.data)
-      .findLast((data) => isRouterPersistedState(data));
+      .findLast(isRouterPersistedState);
 
-    if (isRouterPersistedState(savedState)) {
+    if (savedState) {
       if (!hasExplicitStartupModel) {
         selectedProfile = resolveProfileName(
           currentConfig,
@@ -405,59 +424,7 @@ const routerExtension = (pi: ExtensionAPI) => {
     actions.updateStatus(ctx);
   };
 
-  registerCommands(
-    pi,
-    {
-      get currentConfig() {
-        return currentConfig;
-      },
-      get routerEnabled() {
-        return routerEnabled;
-      },
-      set routerEnabled(v) {
-        routerEnabled = v;
-      },
-      get selectedProfile() {
-        return selectedProfile;
-      },
-      set selectedProfile(v) {
-        selectedProfile = v;
-      },
-      pinnedTierByProfile,
-      thinkingByProfile,
-      get lastDecision() {
-        return lastDecision;
-      },
-      get lastNonRouterModel() {
-        return lastNonRouterModel;
-      },
-      set lastNonRouterModel(v) {
-        lastNonRouterModel = v;
-      },
-      get accumulatedCost() {
-        return accumulatedCost;
-      },
-      get debugEnabled() {
-        return debugEnabled;
-      },
-      set debugEnabled(v) {
-        debugEnabled = v;
-      },
-      get widgetEnabled() {
-        return widgetEnabled;
-      },
-      set widgetEnabled(v) {
-        widgetEnabled = v;
-      },
-      get debugHistory() {
-        return debugHistory;
-      },
-      get lastConfigWarnings() {
-        return lastConfigWarnings;
-      },
-    },
-    actions,
-  );
+  registerCommands(pi, runtimeState, actions);
 
   pi.on('session_start', async (event, ctx) => {
     isInitialized = true;
@@ -549,18 +516,20 @@ const routerExtension = (pi: ExtensionAPI) => {
 
     // User changed pi's thinking level (e.g. via shift+tab).
     // Apply as an all-tier thinking override for the active router profile.
-    thinkingByProfile[selectedProfile] ??= {};
-    const overrides = thinkingByProfile[selectedProfile];
+    let overrides = thinkingByProfile[selectedProfile];
+    if (!overrides) {
+      overrides = {};
+      thinkingByProfile[selectedProfile] = overrides;
+    }
     for (const t of ROUTER_TIERS) {
       overrides[t] = event.level;
     }
     persistState();
     actions.updateStatus(ctx);
     if (event.level !== 'off') {
-      const unsupported = getUnsupportedTiers(
-        currentConfig.profiles[selectedProfile],
-        event.level,
-      );
+      const activeProfile = currentConfig.profiles[selectedProfile];
+      if (!activeProfile) return;
+      const unsupported = getUnsupportedTiers(activeProfile, event.level);
       if (unsupported.length > 0) {
         ctx.ui.notify(
           `Router thinking (all) set to ${event.level}. ` +

package/extensions/provider.ts

@@ -14,6 +14,7 @@ import type {
   ExtensionAPI,
   ExtensionContext,
 } from '@earendil-works/pi-coding-agent';
+import { runClassifier } from './classifier';
 import {
   clampThinkingLevel,
   collectProfileThinkingLevels,
@@ -25,6 +26,13 @@ import {
   resolveMaxTokens,
 } from './config';
 import { DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from './constants';
+import { extractTextFromContent, hasImageAttachment } from './context';
+import {
+  buildRoutingDecision,
+  decideRouting,
+  phaseForTier,
+  resolveAvailableTier,
+} from './routing';
 import type {
   RouterConfig,
   RouterPinByProfile,
@@ -68,16 +76,6 @@ export const waitForRegistry = async (
   return undefined;
 };
 
-import {
-  buildRoutingDecision,
-  decideRouting,
-  extractTextFromContent,
-  hasImageAttachment,
-  phaseForTier,
-  resolveAvailableTier,
-  runClassifier,
-} from './routing';
-
 export const createErrorMessage = (
   model: Model<Api>,
   message: string,
@@ -130,19 +128,28 @@ const truncateContext = (context: Context, limit: number): Context => {
 
   // Drop only complete turns. Splitting an assistant/tool-result pair corrupts transcripts.
   // This text estimate cannot guarantee a fit for images/tools or one oversized active turn.
+  const systemMessages = messages.filter(
+    (message) => message.role === 'system',
+  );
   let remaining = totalTokens;
-  let startIndex = 0;
-  for (let i = 1; i < messages.length && remaining > limit; i++) {
-    if (messages[i].role !== 'user') continue;
-    while (startIndex < i) remaining -= messageTokens[startIndex++];
+  let nextRemovableIndex = 0;
+  for (let i = 0; i < messages.length && remaining > limit; i += 1) {
+    if (messages[i]?.role !== 'user') continue;
+    while (nextRemovableIndex < i) {
+      const candidate = messages[nextRemovableIndex];
+      if (candidate?.role !== 'system') {
+        remaining -= messageTokens[nextRemovableIndex] ?? 0;
+      }
+      nextRemovableIndex += 1;
+    }
   }
   return {
     ...context,
     messages: [
+      ...systemMessages,
       ...messages
-        .slice(0, startIndex)
-        .filter((message) => message.role === 'system'),
-      ...messages.slice(startIndex),
+        .slice(nextRemovableIndex)
+        .filter((message) => message.role !== 'system'),
     ],
   };
 };
@@ -201,24 +208,29 @@ export const registerRouterProvider = (
   const profileList = profileNames(state.currentConfig);
 
   // Map profiles to their capacities
-  const modelDefinitions = profileList.map((name) => {
+  const modelDefinitions = profileList.flatMap((name) => {
     const profile = state.currentConfig.profiles[name];
+    if (!profile) return [];
 
     // Report the MAX context window and max output tokens across all tiers.
     // The honesty check + truncateContext handles the case where the
     // actually routed model is smaller.
     let maxContextWindow = 0;
-    let maxMaxTokens = 0;
+    let maxOutputTokens = 0;
     for (const tier of ROUTER_TIERS) {
       if (!profile[tier]) continue;
-      const cw = resolveContextWindow(
+      const contextWindow = resolveContextWindow(
         tier,
         profile,
         state.currentModelRegistry,
       );
-      const mot = resolveMaxTokens(tier, profile, state.currentModelRegistry);
-      if (cw > maxContextWindow) maxContextWindow = cw;
-      if (mot > maxMaxTokens) maxMaxTokens = mot;
+      const maxTokens = resolveMaxTokens(
+        tier,
+        profile,
+        state.currentModelRegistry,
+      );
+      if (contextWindow > maxContextWindow) maxContextWindow = contextWindow;
+      if (maxTokens > maxOutputTokens) maxOutputTokens = maxTokens;
     }
 
     const hasReasoning = supportsReasoning(profile, state.currentModelRegistry);
@@ -233,16 +245,18 @@ export const registerRouterProvider = (
       if (Object.keys(map).length > 0) thinkingLevelMap = map;
     }
 
-    return {
-      id: name,
-      name: `Router ${name}`,
-      reasoning: hasReasoning,
-      ...(thinkingLevelMap ? { thinkingLevelMap } : {}),
-      input: ['text', 'image'] as ('text' | 'image')[],
-      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-      contextWindow: maxContextWindow || DEFAULT_CONTEXT_WINDOW,
-      maxTokens: maxMaxTokens || DEFAULT_MAX_TOKENS,
-    };
+    return [
+      {
+        id: name,
+        name: `Router ${name}`,
+        reasoning: hasReasoning,
+        ...(thinkingLevelMap ? { thinkingLevelMap } : {}),
+        input: ['text', 'image'] satisfies ('text' | 'image')[],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        contextWindow: maxContextWindow || DEFAULT_CONTEXT_WINDOW,
+        maxTokens: maxOutputTokens || DEFAULT_MAX_TOKENS,
+      },
+    ];
   });
 
   const modelsKey = JSON.stringify(modelDefinitions);
@@ -260,7 +274,7 @@ export const registerRouterProvider = (
     ): AssistantMessageEventStream {
       const stream = createAssistantMessageEventStream();
 
-      (async () => {
+      void (async () => {
         let partialMessage: AssistantMessage | undefined;
         try {
           // Wait for the router to be fully initialized (session_start sets currentModelRegistry).
@@ -447,9 +461,8 @@ export const registerRouterProvider = (
           let lastError: unknown;
           let success = false;
 
-          for (let i = 0; i < modelsToTry.length; i++) {
+          for (const [i, modelRef] of modelsToTry.entries()) {
             options?.signal?.throwIfAborted();
-            const modelRef = modelsToTry[i];
             const { provider: targetProvider, modelId: targetModelId } =
               parseCanonicalModelRef(modelRef);
 
@@ -488,15 +501,15 @@ export const registerRouterProvider = (
                 const tierConfig = profile[decision.tier];
                 if (tierConfig?.resolvedThinkingLevels) {
                   requestedReasoning = clampThinkingLevel(
-                    requestedReasoning as ThinkingLevel,
+                    requestedReasoning,
                     tierConfig.resolvedThinkingLevels,
-                  ) as typeof requestedReasoning;
+                  );
                 }
               }
 
-              const delegatedReasoning =
+              const delegatedReasoning: SimpleStreamOptions['reasoning'] =
                 targetModel.reasoning && requestedReasoning !== 'off'
-                  ? (requestedReasoning as SimpleStreamOptions['reasoning'])
+                  ? requestedReasoning
                   : undefined;
 
               try {
@@ -522,7 +535,7 @@ export const registerRouterProvider = (
                 ...delegationOptions
               } = options ?? {};
 
-              const delegatedOptions = {
+              const delegatedOptions: SimpleStreamOptions = {
                 ...delegationOptions,
                 ...(delegatedReasoning
                   ? { reasoning: delegatedReasoning }

package/extensions/routing.ts

@@ -1,7 +1,12 @@
-import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
-import type { Context, Message } from '@earendil-works/pi-ai';
-import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
-import { isRouterTier, parseCanonicalModelRef } from './config';
+import type { Context } from '@earendil-works/pi-ai';
+import { parseCanonicalModelRef } from './config';
+import {
+  containsAny,
+  countToolResults,
+  countWords,
+  getLastUserText,
+  getRecentConversationText,
+} from './context';
 import type {
   RouterPhase,
   RouterProfile,
@@ -11,67 +16,6 @@ import type {
   RoutingRule,
 } from './types';
 
-export const extractTextFromContent = (
-  content: string | Message['content'],
-): string => {
-  if (typeof content === 'string') {
-    return content;
-  }
-  return content
-    .map((part) => {
-      if (part.type === 'text') return part.text;
-      if (part.type === 'thinking') return part.thinking;
-      if (part.type === 'toolCall')
-        return `${part.name} ${JSON.stringify(part.arguments)}`;
-      return '';
-    })
-    .filter(Boolean)
-    .join('\n');
-};
-
-export const getLastUserText = (context: Context): string => {
-  for (let i = context.messages.length - 1; i >= 0; i--) {
-    const message = context.messages[i];
-    if (message.role === 'user') {
-      return extractTextFromContent(message.content).trim();
-    }
-  }
-  return '';
-};
-
-export const getRecentConversationText = (
-  context: Context,
-  limit = 6,
-): string => {
-  return context.messages
-    .slice(-limit)
-    .map((message) => extractTextFromContent(message.content).trim())
-    .filter(Boolean)
-    .join('\n')
-    .toLowerCase();
-};
-
-export const countToolResults = (context: Context): number => {
-  return context.messages.filter((message) => message.role === 'toolResult')
-    .length;
-};
-
-export const countWords = (text: string): number => {
-  return text.split(/\s+/).filter(Boolean).length;
-};
-
-export const hasImageAttachment = (context: Context): boolean => {
-  return context.messages.some(
-    (message) =>
-      Array.isArray(message.content) &&
-      message.content.some((part) => part.type === 'image'),
-  );
-};
-
-export const containsAny = (text: string, keywords: string[]): boolean => {
-  return keywords.some((keyword) => text.includes(keyword));
-};
-
 export const phaseForTier = (tier: RouterTier): RouterPhase => {
   if (tier === 'high') return 'planning';
   if (tier === 'medium') return 'implementation';
@@ -87,11 +31,13 @@ export const resolveAvailableTier = (
   const order: RouterTier[] = ['low', 'medium', 'high'];
   const startIdx = order.indexOf(preferred);
   for (let i = startIdx + 1; i < order.length; i++) {
-    if (profile[order[i]]) return order[i];
+    const tier = order[i];
+    if (tier && profile[tier]) return tier;
   }
   // Fall "down" as last resort
   for (let i = startIdx - 1; i >= 0; i--) {
-    if (profile[order[i]]) return order[i];
+    const tier = order[i];
+    if (tier && profile[tier]) return tier;
   }
   return preferred; // unreachable if profile has ≥1 tier
 };
@@ -390,92 +336,3 @@ export const decideRouting = (
   decision.isBudgetForced = isBudgetForced;
   return decision;
 };
-
-export const runClassifier = async (
-  classifierModelRef: string,
-  modelRegistry: ExtensionContext['modelRegistry'],
-  context: Context,
-  currentPhase?: RouterPhase,
-  thinking?: ThinkingLevel,
-  signal?: AbortSignal,
-): Promise<{ tier: RouterTier; reasoning: string } | undefined> => {
-  try {
-    const { provider, modelId } = parseCanonicalModelRef(classifierModelRef);
-    const model = modelRegistry.find(provider, modelId);
-    if (!model || provider === 'router') return undefined;
-    signal?.throwIfAborted();
-
-    const promptText = getLastUserText(context);
-    const historyText = getRecentConversationText(context, 4);
-
-    const classifierPrompt = `You are a model router classifier. Your job is to categorize the user's latest request into one of three tiers: "high", "medium", or "low".
-
-Tiers:
-- high: Architecture, design, planning, tradeoff analysis, broad debugging, large refactors, codebase research.
-- medium: Implementation of a known plan, multi-file edits, normal coding work, focused debugging, tests/fixes.
-- low: Summaries, changelogs, formatting, quick explanations, small bounded transforms, simple read-only lookup.
-
-${currentPhase ? `Current conversation phase: ${currentPhase}\n` : ''}
-Recent history:
-${historyText}
-
-Latest user message:
-${promptText}
-
-Return your decision in exactly two lines:
-Tier: [high|medium|low]
-Reasoning: [one short sentence]
-
-${currentPhase === 'planning' ? 'Consider that the conversation is currently in a planning phase. Bias toward "high" unless the request is clearly a simple implementation or summary.' : ''}
-${currentPhase === 'implementation' ? 'Consider that the conversation is currently in an implementation phase. Bias toward "medium" unless the request is clearly planning or a simple summary.' : ''}`;
-
-    const classifierContext: Context = {
-      messages: [
-        { role: 'user', content: classifierPrompt, timestamp: Date.now() },
-      ],
-    };
-
-    const reasoningOption =
-      model.reasoning && thinking && thinking !== 'off' ? thinking : undefined;
-
-    const timeout = AbortSignal.timeout(10_000);
-    const stream = modelRegistry.streamSimple(model, classifierContext, {
-      signal: signal ? AbortSignal.any([signal, timeout]) : timeout,
-      maxTokens: 256,
-      ...(reasoningOption ? { reasoning: reasoningOption } : {}),
-    });
-    let fullText = '';
-    let completed = false;
-    for await (const event of stream) {
-      if (event.type === 'error') return undefined;
-      if (event.type === 'text_delta') fullText += event.delta;
-      if (event.type === 'done') {
-        completed = true;
-        fullText = extractTextFromContent(event.message.content);
-        break;
-      }
-    }
-    if (!completed) return undefined;
-
-    const lines = fullText.trim().split('\n');
-    const tierLine = lines.find((l) => l.toLowerCase().startsWith('tier:'));
-    const reasoningLine = lines.find((l) =>
-      l.toLowerCase().startsWith('reasoning:'),
-    );
-
-    if (tierLine) {
-      const tierValue = tierLine.split(':')[1].trim().toLowerCase();
-      if (isRouterTier(tierValue)) {
-        return {
-          tier: tierValue,
-          reasoning: reasoningLine
-            ? reasoningLine.slice(reasoningLine.indexOf(':') + 1).trim()
-            : 'Classifier decision.',
-        };
-      }
-    }
-  } catch {
-    // Ignore classifier errors and fall back to heuristics
-  }
-  return undefined;
-};