@alexeiled/pi-model-router 0.5.0 → 0.5.2

package/extensions/routing.ts

@@ -1,13 +1,12 @@
-import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
-import type { Context, Message } from '@earendil-works/pi-ai';
-import { streamSimple } from '@earendil-works/pi-ai/compat';
-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 {
-  hasUsableRequestAuth,
-  type RegistryWithProviderAuth,
-  resolveDelegatedModel,
-} from './constants';
+  containsAny,
+  countToolResults,
+  countWords,
+  getLastUserText,
+  getRecentConversationText,
+} from './context';
 import type {
   RouterPhase,
   RouterProfile,
@@ -17,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';
@@ -93,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
 };
@@ -148,6 +88,7 @@ export const decideRouting = (
   rules?: RoutingRule[],
   isBudgetExceeded = false,
 ): RoutingDecision => {
+  if (previousDecision?.profile !== profileName) previousDecision = undefined;
   const prompt = getLastUserText(context).toLowerCase();
   const recentConversation = getRecentConversationText(context);
   const toolResultCount = countToolResults(context);
@@ -372,7 +313,10 @@ export const decideRouting = (
   }
 
   // Resolve to nearest available tier if the selected tier is disabled
-  const resolvedTier = resolveAvailableTier(profile, tier);
+  const resolvedTier =
+    isBudgetForced && !profile.medium && profile.low
+      ? 'low'
+      : resolveAvailableTier(profile, tier);
   if (resolvedTier !== tier) {
     reasoning = `Resolved from ${tier} to ${resolvedTier} tier (${tier} tier is not configured). Original: ${reasoning}`;
     phase = phaseForTier(resolvedTier);
@@ -392,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,
-): Promise<{ tier: RouterTier; reasoning: string } | undefined> => {
-  try {
-    const { provider, modelId } = parseCanonicalModelRef(classifierModelRef);
-    const model = modelRegistry.find(provider, modelId);
-    if (!model) return undefined;
-
-    const auth = await modelRegistry.getApiKeyAndHeaders(model);
-    if (!auth.ok || !hasUsableRequestAuth(auth)) return undefined;
-    const apiKey = auth.apiKey;
-    const headers = auth.headers;
-    const requestModel = await resolveDelegatedModel(
-      modelRegistry as unknown as RegistryWithProviderAuth,
-      model,
-    );
-
-    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 stream = streamSimple(requestModel, classifierContext, {
-      apiKey,
-      headers,
-      ...(reasoningOption ? { reasoning: reasoningOption } : {}),
-    });
-    let fullText = '';
-    for await (const event of stream) {
-      if (event.type === 'text_delta' && typeof event.delta === 'string') {
-        fullText += event.delta;
-      }
-    }
-
-    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.split(':')[1].trim()
-            : 'Classifier decision.',
-        };
-      }
-    }
-  } catch {
-    // Ignore classifier errors and fall back to heuristics
-  }
-  return undefined;
-};

package/extensions/state.ts

@@ -1,23 +1,61 @@
 import { readFileSync, writeFileSync } from 'node:fs';
 import { join } from 'node:path';
 import { getAgentDir } from '@earendil-works/pi-coding-agent';
+import {
+  isObjectRecord,
+  isRouterTier,
+  isThinkingLevel,
+  parseCanonicalModelRef,
+} from './config';
 import type {
+  PersistedStateInput,
   RouterLastProfileState,
   RouterPersistedState,
   RouterPinByProfile,
-  RouterThinkingByProfile,
   RoutingDecision,
 } from './types';
 
 const LAST_PROFILE_STATE_FILE = 'model-router-state.json';
 
-const isRecord = (value: unknown): value is Record<string, unknown> =>
-  typeof value === 'object' && value !== null && !Array.isArray(value);
+const isPhase = (value: unknown) =>
+  value === 'planning' || value === 'implementation' || value === 'lightweight';
+const isFiniteNumber = (value: unknown): value is number =>
+  typeof value === 'number' && Number.isFinite(value);
+const isModelRef = (value: unknown) => {
+  if (typeof value !== 'string') return false;
+  try {
+    parseCanonicalModelRef(value);
+    return true;
+  } catch {
+    return false;
+  }
+};
+const isDecision = (value: unknown): value is RoutingDecision =>
+  isObjectRecord(value) &&
+  isRouterTier(value.tier) &&
+  isPhase(value.phase) &&
+  isThinkingLevel(value.thinking) &&
+  isFiniteNumber(value.timestamp) &&
+  [
+    'profile',
+    'targetProvider',
+    'targetModelId',
+    'targetLabel',
+    'reasoning',
+  ].every((key) => typeof value[key] === 'string') &&
+  ['isClassifier', 'isFallback', 'isBudgetForced', 'isRuleMatched'].every(
+    (key) => value[key] === undefined || typeof value[key] === 'boolean',
+  );
+const isMap = (value: unknown, validate: (entry: unknown) => boolean) =>
+  isObjectRecord(value) &&
+  Object.entries(value).every(
+    ([key, entry]) => key !== '__proto__' && validate(entry),
+  );
 
 export const isRouterLastProfileState = (
   value: unknown,
 ): value is RouterLastProfileState =>
-  isRecord(value) &&
+  isObjectRecord(value) &&
   typeof value.selectedProfile === 'string' &&
   value.selectedProfile.length > 0 &&
   typeof value.timestamp === 'number';
@@ -58,36 +96,61 @@ export const saveLastRouterProfile = (
 export const isRouterPersistedState = (
   value: unknown,
 ): value is RouterPersistedState => {
-  if (typeof value !== 'object' || value === null) {
-    return false;
-  }
-  if (!isRecord(value)) {
-    return false;
-  }
+  if (!isObjectRecord(value)) return false;
   return (
     typeof value.enabled === 'boolean' &&
     typeof value.selectedProfile === 'string' &&
-    typeof value.timestamp === 'number'
+    isFiniteNumber(value.timestamp) &&
+    (value.pinTier === undefined || isRouterTier(value.pinTier)) &&
+    (value.pinByProfile === undefined ||
+      isMap(value.pinByProfile, isRouterTier)) &&
+    (value.thinkingByProfile === undefined ||
+      isMap(
+        value.thinkingByProfile,
+        (tiers) =>
+          isObjectRecord(tiers) &&
+          Object.entries(tiers).every(
+            ([tier, level]) => isRouterTier(tier) && isThinkingLevel(level),
+          ),
+      )) &&
+    (value.lastDecision === undefined || isDecision(value.lastDecision)) &&
+    (value.debugHistory === undefined ||
+      (Array.isArray(value.debugHistory) &&
+        value.debugHistory.every(isDecision))) &&
+    (value.lastPhase === undefined || isPhase(value.lastPhase)) &&
+    (value.lastNonRouterModel === undefined ||
+      isModelRef(value.lastNonRouterModel)) &&
+    (value.accumulatedCost === undefined ||
+      (isFiniteNumber(value.accumulatedCost) && value.accumulatedCost >= 0)) &&
+    ['debugEnabled', 'widgetEnabled'].every(
+      (key) => value[key] === undefined || typeof value[key] === 'boolean',
+    )
   );
 };
 
-export const buildPersistedState = (
-  routerEnabled: boolean,
-  selectedProfile: string | undefined,
-  pinnedTierByProfile: RouterPinByProfile,
-  thinkingByProfile: RouterThinkingByProfile,
-  debugEnabled: boolean,
-  widgetEnabled: boolean,
-  debugHistory: RoutingDecision[],
-  lastDecision: RoutingDecision | undefined,
-  lastNonRouterModel: string | undefined,
-  accumulatedCost: number,
-): RouterPersistedState => {
-  return {
+export const buildPersistedState = ({
+  routerEnabled,
+  selectedProfile,
+  pinnedTierByProfile,
+  thinkingByProfile,
+  debugEnabled,
+  widgetEnabled,
+  debugHistory,
+  lastDecision,
+  lastNonRouterModel,
+  accumulatedCost,
+}: PersistedStateInput): RouterPersistedState => {
+  const pinByProfile: RouterPinByProfile = {};
+  for (const [profile, tier] of Object.entries(pinnedTierByProfile)) {
+    if (tier) pinByProfile[profile] = tier;
+  }
+  return structuredClone({
     enabled: routerEnabled,
     selectedProfile: selectedProfile ?? '',
-    pinTier: selectedProfile ? pinnedTierByProfile[selectedProfile] : undefined,
-    pinByProfile: { ...pinnedTierByProfile },
+    ...(selectedProfile && pinnedTierByProfile[selectedProfile]
+      ? { pinTier: pinnedTierByProfile[selectedProfile] }
+      : {}),
+    pinByProfile,
     thinkingByProfile: { ...thinkingByProfile },
     debugEnabled,
     widgetEnabled,
@@ -97,5 +160,5 @@ export const buildPersistedState = (
     lastNonRouterModel,
     accumulatedCost,
     timestamp: Date.now(),
-  };
+  });
 };

package/extensions/types.ts

@@ -10,49 +10,60 @@ export type RouterThinkingByProfile = Record<string, RouterThinkingByTier>;
 export interface RoutingRule {
   matches: string | string[];
   tier: RouterTier;
-  reason?: string;
+  reason?: string | undefined;
 }
 
 export interface ModelDefinition {
   model: string;
-  contextWindow?: number;
-  maxTokens?: number;
-  reasoning?: boolean;
-  thinkingLevels?: ThinkingLevel[];
+  contextWindow?: number | undefined;
+  maxTokens?: number | undefined;
+  reasoning?: boolean | undefined;
+  thinkingLevels?: ThinkingLevel[] | undefined;
 }
 
 export interface ClassifierConfig {
   model: string;
-  thinking?: ThinkingLevel;
+  thinking?: ThinkingLevel | undefined;
 }
 
 export interface RoutedTierConfig {
   model: string;
-  thinking?: ThinkingLevel;
-  fallbacks?: string[];
-  contextWindow?: number;
-  maxTokens?: number;
-  reasoning?: boolean;
-  thinkingLevels?: ThinkingLevel[];
-  resolvedContextWindow?: number;
-  resolvedMaxTokens?: number;
-  resolvedThinkingLevels?: ThinkingLevel[];
+  thinking?: ThinkingLevel | undefined;
+  fallbacks?: string[] | undefined;
+  contextWindow?: number | undefined;
+  maxTokens?: number | undefined;
+  reasoning?: boolean | undefined;
+  thinkingLevels?: ThinkingLevel[] | undefined;
+  resolvedContextWindow?: number | undefined;
+  resolvedMaxTokens?: number | undefined;
+  resolvedThinkingLevels?: ThinkingLevel[] | undefined;
 }
 
 export interface RouterProfile {
-  high?: RoutedTierConfig;
-  medium?: RoutedTierConfig;
-  low?: RoutedTierConfig;
+  high?: RoutedTierConfig | undefined;
+  medium?: RoutedTierConfig | undefined;
+  low?: RoutedTierConfig | undefined;
 }
 
 export interface RouterConfig {
-  debug?: boolean;
-  classifierModel?: ClassifierConfig;
-  phaseBias?: number;
-  maxSessionBudget?: number;
-  rules?: RoutingRule[];
+  debug?: boolean | undefined;
+  classifierModel?: ClassifierConfig | undefined;
+  phaseBias?: number | undefined;
+  maxSessionBudget?: number | undefined;
+  rules?: RoutingRule[] | undefined;
   profiles: Record<string, RouterProfile>;
-  models?: Record<string, ModelDefinition>;
+  models?: Record<string, ModelDefinition> | undefined;
+}
+
+export interface RouterStatusState {
+  routerEnabled: boolean;
+  selectedProfile: string | undefined;
+  pinnedTierByProfile: RouterPinByProfile;
+  lastDecision: RoutingDecision | undefined;
+  lastNonRouterModel: string | undefined;
+  accumulatedCost: number;
+  widgetEnabled: boolean;
+  currentConfig: RouterConfig;
 }
 
 export interface RoutingDecision {
@@ -65,10 +76,10 @@ export interface RoutingDecision {
   reasoning: string;
   thinking: ThinkingLevel;
   timestamp: number;
-  isClassifier?: boolean;
-  isFallback?: boolean;
-  isBudgetForced?: boolean;
-  isRuleMatched?: boolean;
+  isClassifier?: boolean | undefined;
+  isFallback?: boolean | undefined;
+  isBudgetForced?: boolean | undefined;
+  isRuleMatched?: boolean | undefined;
 }
 
 export interface RouterLastProfileState {
@@ -76,34 +87,51 @@ export interface RouterLastProfileState {
   timestamp: number;
 }
 
+export interface PersistedStateInput {
+  routerEnabled: boolean;
+  selectedProfile: string | undefined;
+  pinnedTierByProfile: RouterPinByProfile;
+  thinkingByProfile: RouterThinkingByProfile;
+  debugEnabled: boolean;
+  widgetEnabled: boolean;
+  debugHistory: RoutingDecision[];
+  lastDecision: RoutingDecision | undefined;
+  lastNonRouterModel: string | undefined;
+  accumulatedCost: number;
+}
+
 export interface RouterPersistedState {
   enabled: boolean;
   selectedProfile: string;
-  pinTier?: RouterTier;
-  pinByProfile?: RouterPinByProfile;
-  thinkingByProfile?: RouterThinkingByProfile;
-  debugEnabled?: boolean;
-  widgetEnabled?: boolean;
-  debugHistory?: RoutingDecision[];
-  lastPhase?: RouterPhase;
-  lastDecision?: RoutingDecision;
-  lastNonRouterModel?: string;
-  accumulatedCost?: number;
+  pinTier?: RouterTier | undefined;
+  pinByProfile?: RouterPinByProfile | undefined;
+  thinkingByProfile?: RouterThinkingByProfile | undefined;
+  debugEnabled?: boolean | undefined;
+  widgetEnabled?: boolean | undefined;
+  debugHistory?: RoutingDecision[] | undefined;
+  lastPhase?: RouterPhase | undefined;
+  lastDecision?: RoutingDecision | undefined;
+  lastNonRouterModel?: string | undefined;
+  accumulatedCost?: number | undefined;
   timestamp: number;
 }
 
+export interface RawRouterConfig {
+  debug?: unknown;
+  classifierModel?: unknown;
+  phaseBias?: unknown;
+  maxSessionBudget?: unknown;
+  rules?: unknown;
+  profiles?: unknown;
+  models?: unknown;
+}
+
 export interface ConfigLoadResult {
   config: RouterConfig;
   warnings: string[];
 }
 
 export interface ParsedConfigFile {
-  config: Partial<RouterConfig>;
+  config: RawRouterConfig;
   warnings: string[];
 }
-
-export interface CustomSessionEntry {
-  type: string;
-  customType?: string;
-  data?: unknown;
-}

package/extensions/ui.ts

@@ -1,17 +1,11 @@
 import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
 import type {
-  RouterConfig,
   RouterPinByProfile,
+  RouterStatusState,
   RouterThinkingByProfile,
   RoutingDecision,
 } from './types';
 
-const getEffectiveThinking = (
-  thinkingByProfile: RouterThinkingByProfile,
-  profileName: string,
-  decision: RoutingDecision,
-) => thinkingByProfile[profileName]?.[decision.tier] ?? decision.thinking;
-
 const getDecisionFlags = (decision: RoutingDecision): string[] => {
   const flags: string[] = [];
   if (decision.isFallback) flags.push('fallback');
@@ -53,16 +47,18 @@ export const formatModelRef = (ref: string | undefined): string => {
 
 export const updateStatus = (
   ctx: ExtensionContext,
-  routerEnabled: boolean,
-  selectedProfile: string | undefined,
-  pinnedTierByProfile: RouterPinByProfile,
-  thinkingByProfile: RouterThinkingByProfile,
-  lastDecision: RoutingDecision | undefined,
-  lastNonRouterModel: string | undefined,
-  accumulatedCost: number,
-  widgetEnabled: boolean,
-  currentConfig: RouterConfig,
+  state: RouterStatusState,
 ) => {
+  const {
+    routerEnabled,
+    selectedProfile,
+    pinnedTierByProfile,
+    lastDecision,
+    lastNonRouterModel,
+    accumulatedCost,
+    widgetEnabled,
+    currentConfig,
+  } = state;
   const activeRouterProfile = routerEnabled ? selectedProfile : undefined;
   const statusProfile = selectedProfile ?? 'none';
   const activePin = selectedProfile
@@ -77,12 +73,7 @@ export const updateStatus = (
 
     let statusText: string;
     if (lastDecision && matchesProfile && matchesPin) {
-      const effectiveThinking = getEffectiveThinking(
-        thinkingByProfile,
-        activeRouterProfile,
-        lastDecision,
-      );
-      statusText = `router:${activeRouterProfile}${pinLabel} -> ${lastDecision.tier} -> ${lastDecision.targetProvider}/${lastDecision.targetModelId} (${effectiveThinking})`;
+      statusText = `router:${activeRouterProfile}${pinLabel} -> ${lastDecision.tier} -> ${lastDecision.targetProvider}/${lastDecision.targetModelId} (${lastDecision.thinking})`;
     } else {
       statusText = `router:${activeRouterProfile}${pinLabel} -> waiting`;
     }
@@ -106,16 +97,11 @@ export const updateStatus = (
         : ''),
   ];
   if (lastDecision && lastDecision.profile === statusProfile) {
-    const effectiveThinking = getEffectiveThinking(
-      thinkingByProfile,
-      statusProfile,
-      lastDecision,
-    );
     const flags = getDecisionFlags(lastDecision);
     const flagsStr = flags.length > 0 ? ` [${flags.join(',')}]` : '';
 
     widgetLines.push(
-      `Route: ${lastDecision.tier}${flagsStr} -> ${lastDecision.targetProvider}/${lastDecision.targetModelId} (${effectiveThinking})`,
+      `Route: ${lastDecision.tier}${flagsStr} -> ${lastDecision.targetProvider}/${lastDecision.targetModelId} (${lastDecision.thinking})`,
       `Phase: ${lastDecision.phase}`,
     );
   } else if (!routerEnabled && lastNonRouterModel) {

package/package.json

@@ -1,16 +1,17 @@
 {
   "name": "@alexeiled/pi-model-router",
-  "version": "0.5.0",
+  "version": "0.5.2",
   "type": "module",
   "files": [
     "extensions",
-    "!extensions/*.test.ts",
+    "!extensions/**/*.test.ts",
+    "!extensions/test",
     "LICENSE",
     "model-router.example.json",
     "README.md",
     "CHANGELOG.md"
   ],
-  "description": "Smart per-turn model router for Pi that optimizes your AI budget and usage limits by dynamically switching LLM tiers based on task complexity, budget, and phase awareness.",
+  "description": "Independently maintained Pi model router with tiered routing, soft budget controls, classifier fallback, and custom-provider support.",
   "keywords": [
     "pi-package",
     "pi",
@@ -23,7 +24,7 @@
   "engines": {
     "node": ">=22.19.0"
   },
-  "packageManager": "npm@11.18.0",
+  "packageManager": "npm@12.0.2",
   "author": "Alexei Ledenev",
   "contributors": [
     "Ye Liu"
@@ -51,7 +52,7 @@
     "tsc": "tsc --noEmit",
     "build": "tsc",
     "prepublishOnly": "npm run check && npm test",
-    "test": "vitest run",
+    "test": "vitest run --pool=threads",
     "check": "biome check --error-on-warnings . && npm run tsc",
     "lint": "biome lint --error-on-warnings .",
     "lint:fix": "biome lint --write --error-on-warnings .",