@kdejaeger/pi-model-router 0.4.2 → 0.4.5

package/README.md

@@ -23,8 +23,8 @@ This extension (forked from [yeliu84/pi-model-router](https://github.com/yeliu84
 
 The pi-model-router registers itself as a **custom logical provider** (`router`) via `pi.registerProvider`. Each profile becomes a stable model (e.g., `router/cheap`). The model shown in your footer on the left stays fixed, while the underlying LLM changes per turn based on task complexity.
 
-![Router footer](docs/footer.png)
-![Router screenshot](docs/screenshot.png)
+![Router footer](https://raw.githubusercontent.com/kdejaeger/pi-model-router/main/docs/footer.png)
+![Router screenshot](https://raw.githubusercontent.com/kdejaeger/pi-model-router/main/docs/screenshot.png)
 
 For the full decision pipeline, see [How Routing Works](#how-routing-works).
 

package/extensions/commands.ts

@@ -1,22 +1,8 @@
-import type {
-  ExtensionAPI,
-  ExtensionContext,
-} from '@earendil-works/pi-coding-agent';
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
 import type { AutocompleteItem } from '@earendil-works/pi-tui';
-import type {
-  RouterConfig,
-  RouterPinByProfile,
-  RoutingDecision,
-} from './types';
-import {
-  profileNames,
-  ROUTER_PIN_VALUES,
-  resolveModelFromRef,
-  isPinValue,
-} from './config';
-import {
-  formatPinSummary,
-} from './ui';
+import type { RouterConfig, RouterPinByProfile, RoutingDecision } from './types';
+import { profileNames, ROUTER_PIN_VALUES, resolveModelFromRef, isPinValue } from './config';
+import { formatPinSummary } from './ui';
 
 export const registerCommands = (
   pi: ExtensionAPI,
@@ -33,16 +19,9 @@ export const registerCommands = (
   actions: {
     persistState: () => void;
     updateStatus: (ctx: ExtensionContext) => void;
-    reloadConfig: (
-      ctx?: ExtensionContext,
-      options?: { preserveDebug?: boolean },
-    ) => void;
+    reloadConfig: (ctx?: ExtensionContext, options?: { preserveDebug?: boolean }) => void;
     ensureValidActiveRouterProfile: (ctx: ExtensionContext) => Promise<void>;
-    switchToRouterProfile: (
-      profileName: string,
-      ctx: ExtensionContext,
-      strict?: boolean,
-    ) => Promise<boolean>;
+    switchToRouterProfile: (profileName: string, ctx: ExtensionContext, strict?: boolean) => Promise<boolean>;
   },
 ) => {
   const SUBCOMMAND_DETAILS = [
@@ -56,12 +35,8 @@ export const registerCommands = (
     { name: 'help', desc: 'Show usage help for subcommands' },
   ];
 
-  const getSubcommandCompletions = (
-    prefix: string,
-  ): AutocompleteItem[] | null => {
-    const items = SUBCOMMAND_DETAILS.filter((s) =>
-      s.name.startsWith(prefix),
-    ).map((s) => ({
+  const getSubcommandCompletions = (prefix: string): AutocompleteItem[] | null => {
+    const items = SUBCOMMAND_DETAILS.filter((s) => s.name.startsWith(prefix)).map((s) => ({
       value: s.name,
       label: s.name,
       description: s.desc,
@@ -73,9 +48,7 @@ export const registerCommands = (
     // pin [profile] <tier|clear>
     if (args.length <= 1) {
       const token = args[0] ?? '';
-      const pinItems = ROUTER_PIN_VALUES.filter((value) =>
-        value.startsWith(token),
-      ).map((value) => ({ value, label: value }));
+      const pinItems = ROUTER_PIN_VALUES.filter((value) => value.startsWith(token)).map((value) => ({ value, label: value }));
       const profileItems = profileNames(state.currentConfig)
         .filter((name) => name.startsWith(token))
         .map((name) => ({ value: name, label: `router/${name}` }));
@@ -86,9 +59,7 @@ export const registerCommands = (
     const profileToken = args[0];
     if (!state.currentConfig.profiles[profileToken]) return null;
     const pinPrefix = args[1] ?? '';
-    const items = ROUTER_PIN_VALUES.filter((value) =>
-      value.startsWith(pinPrefix),
-    ).map((value) => ({
+    const items = ROUTER_PIN_VALUES.filter((value) => value.startsWith(pinPrefix)).map((value) => ({
       value: `${profileToken} ${value}`,
       label: `${profileToken} ${value}`,
     }));
@@ -107,9 +78,7 @@ export const registerCommands = (
       ctx.ui.notify('Usage: /router status (no arguments)', 'error');
       return;
     }
-    const profilePin = state.selectedProfile
-      ? state.pinnedTierByProfile[state.selectedProfile] ?? 'none'
-      : 'none';
+    const profilePin = state.selectedProfile ? (state.pinnedTierByProfile[state.selectedProfile] ?? 'none') : 'none';
     const lines = [
       `Router enabled: ${state.routerEnabled ? 'yes' : 'off'}`,
       `Selected profile: ${state.selectedProfile ?? 'none'}`,
@@ -193,10 +162,7 @@ export const registerCommands = (
     }
 
     if (!isPinValue(pinValue)) {
-      ctx.ui.notify(
-        `Invalid router pin: ${pinValue}. Use high, medium, low, or clear`,
-        'error',
-      );
+      ctx.ui.notify(`Invalid router pin: ${pinValue}. Use high, medium, low, or clear`, 'error');
       return;
     }
 
@@ -235,10 +201,7 @@ export const registerCommands = (
     }
     const targetModel = resolveModelFromRef(state.lastNonRouterModel, ctx.modelRegistry);
     if (!targetModel) {
-      ctx.ui.notify(
-        `Recorded non-router model is unavailable: ${state.lastNonRouterModel}`,
-        'error',
-      );
+      ctx.ui.notify(`Recorded non-router model is unavailable: ${state.lastNonRouterModel}`, 'error');
       return;
     }
     const success = await pi.setModel(targetModel);
@@ -250,10 +213,7 @@ export const registerCommands = (
     actions.persistState();
     pi.setThinkingLevel('off');
     actions.updateStatus(ctx);
-    ctx.ui.notify(
-      `Router disabled. Restored ${state.lastNonRouterModel}`,
-      'info',
-    );
+    ctx.ui.notify(`Router disabled. Restored ${state.lastNonRouterModel}`, 'info');
   };
 
   const handleDebug = async (args: string[], ctx: ExtensionContext) => {
@@ -269,10 +229,7 @@ export const registerCommands = (
     }
     actions.persistState();
     actions.updateStatus(ctx);
-    ctx.ui.notify(
-      `Router debug ${state.debugEnabled ? 'enabled' : 'disabled'}.`,
-      'info',
-    );
+    ctx.ui.notify(`Router debug ${state.debugEnabled ? 'enabled' : 'disabled'}.`, 'info');
   };
 
   const handleReload = async (args: string[], ctx: ExtensionContext) => {
@@ -288,10 +245,7 @@ export const registerCommands = (
       ctx.ui.notify(`Router reload warnings:\n${state.lastConfigWarnings.join('\n')}`, 'warning');
     }
 
-    ctx.ui.notify(
-      `Router config reloaded. Profiles: ${profileNames(state.currentConfig).join(', ')}`,
-      'info',
-    );
+    ctx.ui.notify(`Router config reloaded. Profiles: ${profileNames(state.currentConfig).join(', ')}`, 'info');
   };
 
   pi.registerCommand('router', {
@@ -412,22 +366,13 @@ export const registerCommands = (
             // Check if subcommand is actually a profile name (backwards compatible-ish with /router-on)
             if (state.currentConfig.profiles[subcommand]) {
               if (subArgs.length > 0) {
-                ctx.ui.notify(
-                  `Usage: /router ${subcommand} (no extra arguments allowed)`,
-                  'error',
-                );
+                ctx.ui.notify(`Usage: /router ${subcommand} (no extra arguments allowed)`, 'error');
                 return;
               }
               await actions.switchToRouterProfile(subcommand, ctx);
-              ctx.ui.notify(
-                `Router enabled with profile: ${state.selectedProfile}`,
-                'info',
-              );
+              ctx.ui.notify(`Router enabled with profile: ${state.selectedProfile}`, 'info');
             } else {
-              ctx.ui.notify(
-                `Unknown router subcommand: ${subcommand}. Try /router help`,
-                'error',
-              );
+              ctx.ui.notify(`Unknown router subcommand: ${subcommand}. Try /router help`, 'error');
             }
           } else {
             await handleStatus(subArgs, ctx);

package/extensions/config.ts

@@ -1,26 +1,23 @@
 import { existsSync, readFileSync } from 'node:fs';
 import { join } from 'node:path';
-import { type Api, type Model } from '@earendil-works/pi-ai';
+import {
+  type Api,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+} from '@earendil-works/pi-ai';
+import { streamSimple } from '@earendil-works/pi-ai/compat';
 import { getAgentDir, type ExtensionContext } from '@earendil-works/pi-coding-agent';
 import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
-import type {
-  RouterConfig,
-  RouterProfile,
-  RoutedTierConfig,
-  ConfigLoadResult,
-  ParsedConfigFile,
-  RouterTier,
-} from './types';
+import type { RouterConfig, RouterProfile, RoutedTierConfig, ConfigLoadResult, ParsedConfigFile, RouterTier } from './types';
 
 export const ROUTER_TIERS = ['high', 'medium', 'low'] as const;
 
-const THINKING_LEVELS: readonly ThinkingLevel[] = ['off', 'minimal', 'low', 'medium', 'high', 'xhigh', 'max'];  
+const THINKING_LEVELS: readonly ThinkingLevel[] = ['off', 'minimal', 'low', 'medium', 'high', 'xhigh', 'max'];
 export const ROUTER_PIN_VALUES = ['clear', 'high', 'medium', 'low'] as const;
 
-const isObjectRecord = (
-  value: unknown,
-): value is Record<string, unknown> =>
-  typeof value === 'object' && value !== null;
+const isObjectRecord = (value: unknown): value is Record<string, unknown> => typeof value === 'object' && value !== null;
 
 const isThinkingLevel = (value: unknown): value is ThinkingLevel =>
   typeof value === 'string' && THINKING_LEVELS.includes(value as ThinkingLevel);
@@ -32,7 +29,12 @@ const isThinkingLevel = (value: unknown): value is ThinkingLevel =>
 export const isPinValue = (value: string): value is (typeof ROUTER_PIN_VALUES)[number] =>
   (ROUTER_PIN_VALUES as readonly string[]).includes(value);
 
-const validateNonNegativeInt = (val: unknown, label: string, fallback: number | undefined, warnings: string[]): number | undefined => {
+const validateNonNegativeInt = (
+  val: unknown,
+  label: string,
+  fallback: number | undefined,
+  warnings: string[],
+): number | undefined => {
   if (val === undefined || val === null) return fallback;
   if (typeof val !== 'number' || !Number.isInteger(val) || val < 0) {
     warnings.push(`Invalid ${label} (${JSON.stringify(val)}). Must be a non-negative integer.`);
@@ -58,26 +60,18 @@ const parseConfigFile = (path: string): ParsedConfigFile => {
   } catch (error) {
     return {
       config: {},
-      warnings: [
-        `Failed to parse router config at ${path}: ${error instanceof Error ? error.message : String(error)}`,
-      ],
+      warnings: [`Failed to parse router config at ${path}: ${error instanceof Error ? error.message : String(error)}`],
     };
   }
 };
 
-const mergeTier = (
-  existing?: RoutedTierConfig,
-  next?: Partial<RoutedTierConfig>,
-): RoutedTierConfig | undefined => {
+const mergeTier = (existing?: RoutedTierConfig, next?: Partial<RoutedTierConfig>): RoutedTierConfig | undefined => {
   if (!next) return existing;
   if (!existing) return next?.model ? (next as RoutedTierConfig) : undefined;
   return { ...existing, ...next };
 };
 
-const mergeConfig = (
-  base: RouterConfig,
-  override: Partial<RouterConfig>,
-): RouterConfig => {
+const mergeConfig = (base: RouterConfig, override: Partial<RouterConfig>): RouterConfig => {
   const mergedProfiles: Record<string, RouterProfile> = { ...base.profiles };
   for (const [name, profile] of Object.entries(override.profiles ?? {})) {
     if (!isObjectRecord(profile)) continue;
@@ -92,37 +86,25 @@ const mergeConfig = (
   return {
     debug: override.debug ?? base.debug,
     classifierModels: override.classifierModels ?? base.classifierModels,
-    classifierModelThinking:
-      override.classifierModelThinking ?? base.classifierModelThinking,
-    classifierRunOnceAfterToolCount:
-      override.classifierRunOnceAfterToolCount ?? base.classifierRunOnceAfterToolCount,
-    classifierRunAfterToolFailures:
-      override.classifierRunAfterToolFailures ?? base.classifierRunAfterToolFailures,
-    classifierInterval:
-      override.classifierInterval ?? base.classifierInterval,
-    defaultContextThresholdPercent:
-      override.defaultContextThresholdPercent ?? base.defaultContextThresholdPercent,
-    contextThresholdPercentOverrides:
-      override.contextThresholdPercentOverrides ?? base.contextThresholdPercentOverrides,
+    classifierModelThinking: override.classifierModelThinking ?? base.classifierModelThinking,
+    classifierRunOnceAfterToolCount: override.classifierRunOnceAfterToolCount ?? base.classifierRunOnceAfterToolCount,
+    classifierRunAfterToolFailures: override.classifierRunAfterToolFailures ?? base.classifierRunAfterToolFailures,
+    classifierInterval: override.classifierInterval ?? base.classifierInterval,
+    defaultContextThresholdPercent: override.defaultContextThresholdPercent ?? base.defaultContextThresholdPercent,
+    contextThresholdPercentOverrides: override.contextThresholdPercentOverrides ?? base.contextThresholdPercentOverrides,
     profiles: mergedProfiles,
   };
 };
 
-export const parseCanonicalModelRef = (
-  value: string,
-): { provider: string; modelId: string } => {
+export const parseCanonicalModelRef = (value: string): { provider: string; modelId: string } => {
   const slashIndex = value.indexOf('/');
   if (slashIndex === -1) {
-    throw new Error(
-      `Invalid model reference "${value}". Expected "provider/model".`,
-    );
+    throw new Error(`Invalid model reference "${value}". Expected "provider/model".`);
   }
   const provider = value.slice(0, slashIndex).trim();
   const modelId = value.slice(slashIndex + 1).trim();
   if (!provider || !modelId) {
-    throw new Error(
-      `Invalid model reference "${value}". Expected "provider/model".`,
-    );
+    throw new Error(`Invalid model reference "${value}". Expected "provider/model".`);
   }
   return { provider, modelId };
 };
@@ -155,9 +137,7 @@ const normalizeTierConfig = (
 
   const model = typeof value.model === 'string' ? value.model.trim() : '';
   if (!model) {
-    warnings.push(
-      `Profile "${profileName}" ${tier} tier is missing a model. Tier disabled.`,
-    );
+    warnings.push(`Profile "${profileName}" ${tier} tier is missing a model. Tier disabled.`);
     return undefined;
   }
 
@@ -170,13 +150,9 @@ const normalizeTierConfig = (
     return undefined;
   }
 
-  const thinking = isThinkingLevel(value.thinking)
-    ? value.thinking
-    : 'medium';
+  const thinking = isThinkingLevel(value.thinking) ? value.thinking : 'medium';
   if (value.thinking !== undefined && !isThinkingLevel(value.thinking)) {
-    warnings.push(
-      `Profile "${profileName}" ${tier} tier has invalid thinking level. Defaulting to medium.`,
-    );
+    warnings.push(`Profile "${profileName}" ${tier} tier has invalid thinking level. Defaulting to medium.`);
   }
 
   let fallbacks: string[] | undefined;
@@ -213,29 +189,12 @@ const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
     if (trimmedName !== name) {
       warnings.push(`Profile name "${name}" has leading/trailing whitespace. Using "${trimmedName}".`);
     }
-    const high = normalizeTierConfig(
-      profile?.high,
-      trimmedName,
-      'high',
-      warnings,
-    );
-    const medium = normalizeTierConfig(
-      profile?.medium,
-      trimmedName,
-      'medium',
-      warnings,
-    );
-    const low = normalizeTierConfig(
-      profile?.low,
-      trimmedName,
-      'low',
-      warnings,
-    );
+    const high = normalizeTierConfig(profile?.high, trimmedName, 'high', warnings);
+    const medium = normalizeTierConfig(profile?.medium, trimmedName, 'medium', warnings);
+    const low = normalizeTierConfig(profile?.low, trimmedName, 'low', warnings);
 
     if (!high && !medium && !low) {
-      warnings.push(
-        `Profile "${trimmedName}" has no valid tiers. Skipped.`,
-      );
+      warnings.push(`Profile "${trimmedName}" has no valid tiers. Skipped.`);
       continue;
     }
     if (!high) {
@@ -271,9 +230,7 @@ const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
         `defaultContextThresholdPercent (${raw.defaultContextThresholdPercent}) is not a positive number. Falling back to 90.`,
       );
     } else if (raw.defaultContextThresholdPercent > 100) {
-      warnings.push(
-        `defaultContextThresholdPercent (${raw.defaultContextThresholdPercent}) exceeds 100. Falling back to 90.`,
-      );
+      warnings.push(`defaultContextThresholdPercent (${raw.defaultContextThresholdPercent}) exceeds 100. Falling back to 90.`);
     } else {
       defaultContextThresholdPercent = raw.defaultContextThresholdPercent;
     }
@@ -287,11 +244,15 @@ const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
           try {
             parseCanonicalModelRef(trimmed);
           } catch (error) {
-            warnings.push(`Ignored contextThresholdPercentOverride "${key}": invalid model reference — ${error instanceof Error ? error.message : String(error)}`);
+            warnings.push(
+              `Ignored contextThresholdPercentOverride "${key}": invalid model reference — ${error instanceof Error ? error.message : String(error)}`,
+            );
             return [];
           }
           if (typeof val !== 'number' || val <= 0) {
-            warnings.push(`Ignored contextThresholdPercentOverride "${key}" (${JSON.stringify(val)}): expected a positive number.`);
+            warnings.push(
+              `Ignored contextThresholdPercentOverride "${key}" (${JSON.stringify(val)}): expected a positive number.`,
+            );
             return [];
           }
           return [[trimmed, val] as [string, number]];
@@ -310,9 +271,7 @@ const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
           parseCanonicalModelRef(trimmedCM);
           classifierModels.push(trimmedCM);
         } catch (error) {
-          warnings.push(
-            `Invalid classifierModels entry "${rawCM}": ${error instanceof Error ? error.message : String(error)}`,
-          );
+          warnings.push(`Invalid classifierModels entry "${rawCM}": ${error instanceof Error ? error.message : String(error)}`);
         }
       } else {
         warnings.push(`Ignored non-string classifierModels entry: ${JSON.stringify(rawCM)}`);
@@ -328,8 +287,18 @@ const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
     warnings.push(`Invalid classifierModelThinking value "${raw.classifierModelThinking}". Falling back to "off".`);
   }
 
-  const classifierRunOnceAfterToolCount = validateNonNegativeInt(raw.classifierRunOnceAfterToolCount, 'classifierRunOnceAfterToolCount', 3, warnings);
-  const classifierRunAfterToolFailures = validateNonNegativeInt(raw.classifierRunAfterToolFailures, 'classifierRunAfterToolFailures', 2, warnings);
+  const classifierRunOnceAfterToolCount = validateNonNegativeInt(
+    raw.classifierRunOnceAfterToolCount,
+    'classifierRunOnceAfterToolCount',
+    3,
+    warnings,
+  );
+  const classifierRunAfterToolFailures = validateNonNegativeInt(
+    raw.classifierRunAfterToolFailures,
+    'classifierRunAfterToolFailures',
+    2,
+    warnings,
+  );
   const classifierInterval = validateNonNegativeInt(raw.classifierInterval, 'classifierInterval', 10, warnings);
 
   return {
@@ -354,18 +323,11 @@ export const loadRouterConfig = (cwd: string): ConfigLoadResult => {
   const globalResult = parseConfigFile(globalPath);
   const projectResult = parseConfigFile(projectPath);
   const baseConfig: RouterConfig = { profiles: {} };
-  const merged = mergeConfig(
-    mergeConfig(baseConfig, globalResult.config),
-    projectResult.config,
-  );
+  const merged = mergeConfig(mergeConfig(baseConfig, globalResult.config), projectResult.config);
   const normalized = normalizeConfig(merged);
   return {
     config: normalized.config,
-    warnings: [
-      ...globalResult.warnings,
-      ...projectResult.warnings,
-      ...normalized.warnings,
-    ],
+    warnings: [...globalResult.warnings, ...projectResult.warnings, ...normalized.warnings],
   };
 };
 
@@ -386,20 +348,57 @@ export const OPENROUTER_ATTR_HEADERS: Readonly<Record<string, string>> = {
 /** Create an onPayload handler that injects session_id for OpenRouter session tracking. */
 export const createOpenRouterOnPayload = (
   sessionProvider?: { getSessionId(): string; getSessionName(): string | undefined },
-  origOnPayload?: (p: any, m: any) => any,
-): ((p: any, m: any) => Promise<any>) | undefined => {
+  origOnPayload?: (payload: unknown, model: Model<Api>) => unknown | Promise<unknown>,
+): ((payload: unknown, model: Model<Api>) => Promise<unknown>) | undefined => {
   const rawId = sessionProvider?.getSessionId();
   const name = sessionProvider?.getSessionName();
   const sessionId = name && rawId ? `${name.replace(/\s+/g, '-')}-${rawId.slice(0, 8)}` : rawId;
   if (!sessionId) return undefined;
-  return async (p: any, m: any) => {
-    const payload = origOnPayload ? await origOnPayload(p, m) : p;
-    return { ...payload, session_id: sessionId };
+  return async (p: unknown, m: Model<Api>) => {
+    const payload = origOnPayload ? ((await origOnPayload(p, m)) ?? p) : p;
+    return typeof payload === 'object' && payload !== null ? { ...payload, session_id: sessionId } : payload;
   };
 };
 
-export const resolveProfileName = (
-  config: RouterConfig,
-  requested?: string,
-): string | undefined =>
+export const resolveProfileName = (config: RouterConfig, requested?: string): string | undefined =>
   requested && config.profiles[requested] ? requested : undefined;
+
+/**
+ * Duck-typing interface for forward-compatible dispatch with Pi ModelRegistry
+ * runtimes that register custom provider stream handlers.
+ */
+interface RegistryWithProviderDispatch {
+  getRegisteredProviderConfig?(providerName: string):
+    | {
+        api?: Api;
+        streamSimple?: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
+      }
+    | undefined;
+  getRegisteredNativeProvider?(providerName: string):
+    | {
+        streamSimple?: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
+      }
+    | undefined;
+}
+
+/**
+ * Dispatch streaming to a registered custom provider if available,
+ * falling back to the global compatibility streamSimple.
+ */
+export const dispatchStream = (
+  model: Model<Api>,
+  context: Context,
+  options: SimpleStreamOptions,
+  modelRegistry?: ExtensionContext['modelRegistry'],
+): AssistantMessageEventStream => {
+  const registry = modelRegistry as unknown as RegistryWithProviderDispatch | undefined;
+  const providerConfig = registry?.getRegisteredProviderConfig?.(model.provider);
+  if (providerConfig?.streamSimple && (!providerConfig.api || providerConfig.api === model.api)) {
+    return providerConfig.streamSimple(model, context, options);
+  }
+  const nativeProvider = registry?.getRegisteredNativeProvider?.(model.provider);
+  if (nativeProvider?.streamSimple) {
+    return nativeProvider.streamSimple(model, context, options);
+  }
+  return streamSimple(model, context, options);
+};

package/extensions/index.ts

@@ -1,18 +1,6 @@
-import type {
-  ExtensionAPI,
-  ExtensionContext,
-} from '@earendil-works/pi-coding-agent';
-import {
-  type RouterConfig,
-  type RoutingDecision,
-  type RouterPinByProfile,
-  type CustomSessionEntry,
-} from './types';
-import {
-  loadRouterConfig,
-  profileNames,
-  resolveProfileName,
-} from './config';
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
+import { type RouterConfig, type RoutingDecision, type RouterPinByProfile, type CustomSessionEntry } from './types';
+import { loadRouterConfig, profileNames, resolveProfileName } from './config';
 import { isRouterPersistedState, buildPersistedState } from './state';
 import { updateStatus } from './ui';
 import { registerCommands } from './commands';
@@ -35,12 +23,13 @@ const routerExtension = (pi: ExtensionAPI) => {
   let isInitialized = false;
   let isRouterDelegating = false;
 
-  const setModelInternally = async (
-    model: NonNullable<ExtensionContext['model']>,
-  ) => {
+  const setModelInternally = async (model: NonNullable<ExtensionContext['model']>) => {
     isRouterDelegating = true;
     try {
       return await pi.setModel(model);
+    } catch {
+      // Extension context may be stale after session teardown.
+      return false;
     } finally {
       isRouterDelegating = false;
     }
@@ -58,31 +47,24 @@ const routerExtension = (pi: ExtensionAPI) => {
     const snapshot = JSON.stringify({
       ...state,
       timestamp: 0,
-      lastDecision: state.lastDecision
-        ? { ...state.lastDecision, timestamp: 0 }
-        : undefined,
+      lastDecision: state.lastDecision ? { ...state.lastDecision, timestamp: 0 } : undefined,
     });
     if (snapshot === lastPersistedSnapshot) {
       return;
     }
-    pi.appendEntry('router-state', state);
+    try {
+      pi.appendEntry('router-state', state);
+    } catch {
+      // Defensive fallback: session_shutdown or teardown may have invalidated context
+      return;
+    }
     lastPersistedSnapshot = snapshot;
   };
 
   const actions = {
     persistState,
-    updateStatus: (ctx: ExtensionContext) =>
-      updateStatus(
-        ctx,
-        routerEnabled,
-        selectedProfile,
-        pinnedTierByProfile,
-        lastDecision,
-      ),
-    reloadConfig: (
-      ctx?: ExtensionContext,
-      options?: { preserveDebug?: boolean },
-    ) => {
+    updateStatus: (ctx: ExtensionContext) => updateStatus(ctx, routerEnabled, selectedProfile, pinnedTierByProfile, lastDecision),
+    reloadConfig: (ctx?: ExtensionContext, options?: { preserveDebug?: boolean }) => {
       const loaded = loadRouterConfig(currentCwd);
       currentConfig = loaded.config;
       lastConfigWarnings = loaded.warnings;
@@ -114,11 +96,7 @@ const routerExtension = (pi: ExtensionAPI) => {
       routerEnabled = false;
       selectedProfile = undefined;
     },
-    switchToRouterProfile: async (
-      profileName: string,
-      ctx: ExtensionContext,
-      strict = true,
-    ) => {
+    switchToRouterProfile: async (profileName: string, ctx: ExtensionContext, strict = true) => {
       if (!currentConfig.profiles[profileName]) {
         if (strict) {
           ctx.ui.notify(`Unknown router profile: ${profileName}`, 'error');
@@ -146,7 +124,11 @@ const routerExtension = (pi: ExtensionAPI) => {
       selectedProfile = profileName;
       routerEnabled = true;
       persistState();
-      pi.setThinkingLevel('off');
+      try {
+        pi.setThinkingLevel('off');
+      } catch {
+        // Stale context
+      }
       actions.updateStatus(ctx);
       return true;
     },
@@ -217,39 +199,26 @@ const routerExtension = (pi: ExtensionAPI) => {
     await new Promise((resolve) => setTimeout(resolve, 50));
 
     routerEnabled = ctx.model?.provider === 'router';
-    selectedProfile = resolveProfileName(
-      currentConfig,
-      ctx.model?.provider === 'router' ? ctx.model.id : selectedProfile,
-    );
+    selectedProfile = resolveProfileName(currentConfig, ctx.model?.provider === 'router' ? ctx.model.id : selectedProfile);
     pinnedTierByProfile = {};
     lastNonRouterModel =
-      ctx.model && ctx.model.provider !== 'router'
-        ? `${ctx.model.provider}/${ctx.model.id}`
-        : lastNonRouterModel;
+      ctx.model && ctx.model.provider !== 'router' ? `${ctx.model.provider}/${ctx.model.id}` : lastNonRouterModel;
 
     const entries = ctx.sessionManager.getBranch() as CustomSessionEntry[];
     const savedState = entries
-      .filter(
-        (entry) =>
-          entry.type === 'custom' && entry.customType === 'router-state',
-      )
+      .filter((entry) => entry.type === 'custom' && entry.customType === 'router-state')
       .map((entry) => entry.data)
       .findLast((data) => isRouterPersistedState(data));
 
     if (isRouterPersistedState(savedState)) {
-      selectedProfile = resolveProfileName(
-        currentConfig,
-        savedState.selectedProfile,
-      );
+      selectedProfile = resolveProfileName(currentConfig, savedState.selectedProfile);
       if (!selectedProfile) {
         routerEnabled = false;
       } else {
         routerEnabled = savedState.enabled;
       }
       lastDecision = savedState.lastDecision;
-      pinnedTierByProfile = savedState.pinByProfile
-        ? { ...savedState.pinByProfile }
-        : {};
+      pinnedTierByProfile = savedState.pinByProfile ? { ...savedState.pinByProfile } : {};
       if (savedState.pinTier && selectedProfile) {
         pinnedTierByProfile[selectedProfile] = savedState.pinTier;
       }
@@ -319,7 +288,7 @@ const routerExtension = (pi: ExtensionAPI) => {
       set debugEnabled(v) {
         debugEnabled = v;
       },
-      
+
       get lastConfigWarnings() {
         return lastConfigWarnings;
       },
@@ -327,6 +296,15 @@ const routerExtension = (pi: ExtensionAPI) => {
     actions,
   );
 
+  const ensureInitializedFromContext = (ctx: ExtensionContext) => {
+    if (!currentModelRegistry) {
+      currentModelRegistry = ctx.modelRegistry;
+      lastExtensionContext = ctx;
+      currentCwd = ctx.cwd;
+      actions.reloadConfig(ctx);
+    }
+  };
+
   pi.on('session_start', async (_event, ctx) => {
     await restoreStateFromSession(ctx);
     isInitialized = true;
@@ -336,14 +314,16 @@ const routerExtension = (pi: ExtensionAPI) => {
     }
 
     if (debugEnabled) {
-      ctx.ui.notify(
-        `Router initialized with profiles: ${profileNames(currentConfig).join(', ')}`,
-        'info',
-      );
+      ctx.ui.notify(`Router initialized with profiles: ${profileNames(currentConfig).join(', ')}`, 'info');
     }
   });
 
+  pi.on('turn_start', async (_event, ctx) => {
+    ensureInitializedFromContext(ctx);
+  });
+
   pi.on('model_select', async (event, ctx) => {
+    ensureInitializedFromContext(ctx);
     if (!isInitialized || isRouterDelegating) return;
     if (event.model.provider === 'router') {
       const profileName = resolveProfileName(currentConfig, event.model.id);
@@ -357,8 +337,7 @@ const routerExtension = (pi: ExtensionAPI) => {
       const registryModel = ctx.modelRegistry.find('router', profileName);
       if (
         registryModel &&
-        (registryModel.contextWindow !== event.model.contextWindow ||
-          registryModel.maxTokens !== event.model.maxTokens)
+        (registryModel.contextWindow !== event.model.contextWindow || registryModel.maxTokens !== event.model.maxTokens)
       ) {
         await setModelInternally(registryModel);
       }
@@ -375,19 +354,28 @@ const routerExtension = (pi: ExtensionAPI) => {
   });
 
   pi.on('turn_end', async (_event, ctx) => {
+    ensureInitializedFromContext(ctx);
     if (routerEnabled && selectedProfile && ctx.model?.provider !== 'router') {
       const routerModel = ctx.modelRegistry.find('router', selectedProfile);
       if (routerModel) {
         const success = await setModelInternally(routerModel);
         if (!success) {
-          ctx.ui.notify('Failed to re-assert router model after turn. Router disabled.', 'warning');
+          try {
+            ctx.ui.notify('Failed to re-assert router model after turn. Router disabled.', 'warning');
+          } catch {
+            // Stale context
+          }
           routerEnabled = false;
           selectedProfile = undefined;
         }
       }
     }
     persistState();
-    actions.updateStatus(ctx);
+    try {
+      actions.updateStatus(ctx);
+    } catch {
+      // Stale context
+    }
   });
 };
 

package/extensions/provider.ts

@@ -5,15 +5,56 @@ import {
   type Context,
   createAssistantMessageEventStream,
   type Model,
+  type ProviderHeaders,
   type SimpleStreamOptions,
 } from '@earendil-works/pi-ai';
-import { streamSimple } from '@earendil-works/pi-ai/compat';
 import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
 import type { RouterConfig, RouterPinByProfile, RouterTier, RoutingDecision } from './types';
-import { createOpenRouterOnPayload, OPENROUTER_ATTR_HEADERS, parseCanonicalModelRef, profileNames, resolveModelFromRef, ROUTER_TIERS } from './config';
-import { buildRoutingDecision, countToolResultsSinceLastUserPrompt, extractTextFromContent, runClassifier, shouldRunClassifier } from './routing';
+import {
+  createOpenRouterOnPayload,
+  dispatchStream,
+  OPENROUTER_ATTR_HEADERS,
+  parseCanonicalModelRef,
+  profileNames,
+  resolveModelFromRef,
+  ROUTER_TIERS,
+} from './config';
+import {
+  buildRoutingDecision,
+  countToolResultsSinceLastUserPrompt,
+  extractTextFromContent,
+  runClassifier,
+  shouldRunClassifier,
+} from './routing';
 import { formatDecision } from './ui';
 
+const REGISTRY_WAIT_TIMEOUT_MS = 5000;
+const REGISTRY_WAIT_INITIAL_DELAY_MS = 50;
+const REGISTRY_WAIT_MAX_DELAY_MS = 500;
+
+/**
+ * Wait for the model registry to become available with exponential backoff.
+ * This handles race conditions where subagents invoke the router provider
+ * before session_start has fired in their context.
+ */
+export const waitForRegistry = async (
+  state: {
+    readonly currentModelRegistry: ExtensionContext['modelRegistry'] | undefined;
+  },
+  timeoutMs: number = REGISTRY_WAIT_TIMEOUT_MS,
+): Promise<ExtensionContext['modelRegistry'] | undefined> => {
+  if (state.currentModelRegistry) return state.currentModelRegistry;
+
+  const start = Date.now();
+  let delay = REGISTRY_WAIT_INITIAL_DELAY_MS;
+  while (Date.now() - start < timeoutMs) {
+    await new Promise((resolve) => setTimeout(resolve, delay));
+    if (state.currentModelRegistry) return state.currentModelRegistry;
+    delay = Math.min(delay * 2, REGISTRY_WAIT_MAX_DELAY_MS);
+  }
+  return undefined;
+};
+
 const createErrorMessage = (model: Model<Api>, message: string): AssistantMessage => {
   return {
     role: 'assistant',
@@ -40,7 +81,7 @@ const estimateTokens = (text: string): number => Math.ceil(text.length / 4);
 
 /**
  * Truncate context to fit within a target token limit by removing oldest messages.
- * Always preserves the first system message and the latest user message.
+ * Preserves the first system message and the latest user message, ensuring no orphaned tool results.
  */
 const truncateContext = (context: Context, limit: number): Context => {
   const messages = [...context.messages];
@@ -53,12 +94,19 @@ const truncateContext = (context: Context, limit: number): Context => {
   const latestMessage = messages.pop()!;
   const latestTokens = estimateTokens(extractTextFromContent(latestMessage.content));
 
-  let runningTotal = systemTokens + latestTokens + messages.reduce((sum, m) => sum + estimateTokens(extractTextFromContent(m.content)), 0);
+  let runningTotal =
+    systemTokens + latestTokens + messages.reduce((sum, m) => sum + estimateTokens(extractTextFromContent(m.content)), 0);
   while (messages.length > 0 && runningTotal > limit) {
     const shifted = messages.shift()!;
     runningTotal -= estimateTokens(extractTextFromContent(shifted.content));
   }
 
+  // Ensure the truncated context starts cleanly at a user prompt, rather than an orphaned toolResult
+  while (messages.length > 0 && messages[0].role === 'toolResult') {
+    const dropped = messages.shift()!;
+    runningTotal -= estimateTokens(extractTextFromContent(dropped.content));
+  }
+
   return { ...context, messages: [...messages, latestMessage] };
 };
 
@@ -123,9 +171,7 @@ export const registerRouterProvider = (
           if (!tierConfig) return false;
           return (
             resolveModelFromRef(tierConfig.model, state.currentModelRegistry)?.reasoning ||
-            tierConfig.fallbacks?.some(
-              (fb) => resolveModelFromRef(fb, state.currentModelRegistry)?.reasoning,
-            )
+            tierConfig.fallbacks?.some((fb) => resolveModelFromRef(fb, state.currentModelRegistry)?.reasoning)
           );
         })
       : false;
@@ -142,9 +188,9 @@ export const registerRouterProvider = (
   });
 
   if (state.currentModelRegistry) {
-    const invalidOverrides = Object.keys(
-      state.currentConfig.contextThresholdPercentOverrides ?? {},
-    ).filter((modelRef) => !resolveModelFromRef(modelRef, state.currentModelRegistry));
+    const invalidOverrides = Object.keys(state.currentConfig.contextThresholdPercentOverrides ?? {}).filter(
+      (modelRef) => !resolveModelFromRef(modelRef, state.currentModelRegistry),
+    );
 
     if (invalidOverrides.length > 0) {
       state.lastExtensionContext?.ui.notify(
@@ -157,25 +203,21 @@ export const registerRouterProvider = (
   const loadedModelKeys = models.map((m) => `${m.id}:${m.contextWindow}:${m.maxTokens}:${m.reasoning}`).join(',');
   if (state.lastLoadedModelKeys === loadedModelKeys) return; // models did not change, no need to re-register
 
-  pi.registerProvider(
-    'router', { // config (baseUrl, apiKey, ...)
+  pi.registerProvider('router', {
+    // config (baseUrl, apiKey, ...)
     baseUrl: 'router://local',
     apiKey: 'pi-model-router',
     api: 'router-local-api',
     models,
-    streamSimple: (
-      model: Model<Api>,
-      context: Context,
-      options?: SimpleStreamOptions,
-    ): AssistantMessageEventStream => {
+    streamSimple: (model: Model<Api>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream => {
       const stream = createAssistantMessageEventStream();
       const ctx = state.lastExtensionContext;
-      const modelRegistry = state.currentModelRegistry;
 
       (async () => {
         try {
+          const modelRegistry = await waitForRegistry(state);
           if (!modelRegistry) {
-            throw new Error('Router provider not initialized yet. Wait for session_start and retry.');
+            throw new Error('Router provider initialization timed out. session_start may not have fired.');
           }
           const profile = currentConfig.profiles[model.id];
           if (!profile) {
@@ -200,7 +242,8 @@ export const registerRouterProvider = (
             lastDecision?.profile === model.id &&
             lastDecision?.targetProvider === 'google' &&
             lastDecision?.thinking !== 'off';
-          if (isGoogleContinuation) { // Google thinking lock — preserve exact model on tool-result continuations
+          if (isGoogleContinuation) {
+            // Google thinking lock — preserve exact model on tool-result continuations
             const toolResultsCount = countToolResultsSinceLastUserPrompt(context);
             decision = {
               ...lastDecision!,
@@ -216,7 +259,15 @@ export const registerRouterProvider = (
             if (currentConfig.classifierModels?.length && !pinnedTier) {
               const toolResultsCount = countToolResultsSinceLastUserPrompt(context);
 
-              bShouldRunClassifier = shouldRunClassifier(currentConfig, context, lastDecision, lastMsgWasTool, toolResultsCount, state.debugEnabled, ctx);
+              bShouldRunClassifier = shouldRunClassifier(
+                currentConfig,
+                context,
+                lastDecision,
+                lastMsgWasTool,
+                toolResultsCount,
+                state.debugEnabled,
+                ctx,
+              );
               const classifierResult = bShouldRunClassifier
                 ? await runClassifier(currentConfig, modelRegistry, context, lastDecision, ctx, state.debugEnabled)
                 : null;
@@ -295,7 +346,8 @@ export const registerRouterProvider = (
 
                 const thresholdPercent =
                   currentConfig.contextThresholdPercentOverrides?.[modelRef] ??
-                  currentConfig.defaultContextThresholdPercent ?? 90;
+                  currentConfig.defaultContextThresholdPercent ??
+                  90;
                 const targetContextWindow = targetModel.contextWindow || 200_000;
                 const targetContextLimit = Math.floor((thresholdPercent / 100) * targetContextWindow);
                 const fitsContext = tokensUsed <= targetContextLimit;
@@ -315,7 +367,9 @@ export const registerRouterProvider = (
 
                   if (tier !== decision.tier) {
                     decision = buildRoutingDecision(
-                      model.id, profile, tier,
+                      model.id,
+                      profile,
+                      tier,
                       `Forced ${tier} tier because ${decision.tier} tier lacks models${triggerReasons ? ` for ${triggerReasons}` : ''}.`,
                       decision.lastClassifierRunToolCount,
                     );
@@ -330,10 +384,14 @@ export const registerRouterProvider = (
                 decision.isContextTriggered = !fitsContext;
 
                 if (ctx) {
-                  if (state.debugEnabled && bShouldRunClassifier) {
-                    ctx.ui.notify(`Decision ${formatDecision(decision)}`, 'info');
+                  try {
+                    if (state.debugEnabled && bShouldRunClassifier) {
+                      ctx.ui.notify(`Decision ${formatDecision(decision)}`, 'info');
+                    }
+                    actions.updateStatus(ctx);
+                  } catch {
+                    // Stale extension context
                   }
-                  actions.updateStatus(ctx);
                 }
 
                 const auth = await modelRegistry.getApiKeyAndHeaders(targetModel);
@@ -351,20 +409,28 @@ export const registerRouterProvider = (
                 let effectiveContext = context;
                 if (!fitsContext) {
                   effectiveContext = truncateContext(context, targetContextLimit);
-                  ctx?.ui.notify(`Memory too large for ${modelRef} — trimmed ${context.messages.length - effectiveContext.messages.length} messages. Run /compact to reduce context size.`, 'warning');
+                  ctx?.ui.notify(
+                    `Memory too large for ${modelRef} — trimmed ${context.messages.length - effectiveContext.messages.length} messages. Run /compact to reduce context size.`,
+                    'warning',
+                  );
                 }
 
                 // Attention: stripping reasoning from pi's incoming options so it doesn't leak into ...baseOptions. The router controls this via delegatedReasoning below.
                 const { onPayload, headers: originalHeaders, reasoning: _incomingReasoning, ...baseOptions } = options ?? {};
 
-                const effectiveHeaders: Record<string, string> = {
-                  ...(originalHeaders as Record<string, string> | undefined),
-                  ...(auth.headers ?? {}),
+                const effectiveHeaders: ProviderHeaders = {
+                  ...targetModel.headers,
                   ...(targetProvider === 'openrouter' ? OPENROUTER_ATTR_HEADERS : {}),
+                  ...auth.headers,
+                  ...originalHeaders,
                 };
 
                 const delegatedReasoning = targetModel.reasoning && decision.thinking !== 'off' ? decision.thinking : undefined;
-                pi.setThinkingLevel(delegatedReasoning ?? 'off');
+                try {
+                  pi.setThinkingLevel(delegatedReasoning ?? 'off');
+                } catch {
+                  // Stale extension context after session teardown
+                }
 
                 const effectiveOptions: SimpleStreamOptions = {
                   ...baseOptions,
@@ -379,10 +445,16 @@ export const registerRouterProvider = (
                   effectiveOptions.onPayload = onPayload;
                 }
 
+                // Apply credential-specific baseUrl override if resolved by getApiKeyAndHeaders
+                // (e.g. GitHub Copilot business/enterprise tenant endpoints) to avoid 421 Misdirected Request.
+                const authBaseUrl = (auth as { baseUrl?: string }).baseUrl;
+                const requestModel = authBaseUrl ? { ...targetModel, baseUrl: authBaseUrl } : targetModel;
+
+                let eventsPushed = 0;
                 const MAX_ATTEMPTS_PER_MODEL = 2;
                 for (let attempt = 1; attempt <= MAX_ATTEMPTS_PER_MODEL; attempt++) {
                   try {
-                    const delegatedStream = streamSimple(targetModel, effectiveContext, effectiveOptions);
+                    const delegatedStream = dispatchStream(requestModel, effectiveContext, effectiveOptions, modelRegistry);
                     let contentReceived = false;
                     for await (const event of delegatedStream) {
                       if (event.type === 'error' && !contentReceived) {
@@ -397,15 +469,23 @@ export const registerRouterProvider = (
                         contentReceived = true;
                       }
                       stream.push(event);
+                      eventsPushed++;
                     }
                     success = true;
                     state.lastDecision = decision;
                     break attemptLoop;
                   } catch (err) {
                     lastError = err;
+                    if (eventsPushed > 0) {
+                      // Stream was already partially transmitted to the consumer; cannot retry or fall back safely.
+                      throw err;
+                    }
                     const remaining = MAX_ATTEMPTS_PER_MODEL - attempt;
                     const retryMsg = remaining > 0 ? ` — ${remaining} ${remaining === 1 ? 'retry' : 'retries'} left` : '';
-                    ctx?.ui.notify(`Failed to reach model ${modelRef} (attempt ${attempt}/${MAX_ATTEMPTS_PER_MODEL}): ${err}${retryMsg}`, 'warning');
+                    ctx?.ui.notify(
+                      `Failed to reach model ${modelRef} (attempt ${attempt}/${MAX_ATTEMPTS_PER_MODEL}): ${err}${retryMsg}`,
+                      'warning',
+                    );
                   }
                 }
               }
@@ -414,20 +494,34 @@ export const registerRouterProvider = (
 
           if (!success) {
             const errorMsg = `Failed to delegate to any model in the chain.${failureReasons.length > 0 ? ' Reasons: ' + failureReasons.filter(Boolean).join('; ') + '.' : ''}`;
-            const combinedError = lastError ? new Error(`${(lastError as Error).message} — ${errorMsg}`) : new Error(errorMsg);
+            const lastErrorText = lastError instanceof Error ? lastError.message : String(lastError ?? '');
+            const combinedError = lastError ? new Error(`${lastErrorText} — ${errorMsg}`) : new Error(errorMsg);
             throw combinedError;
           }
 
           stream.end();
         } catch (error) {
-          stream.push({
-            type: 'error',
-            reason: 'error',
-            error: createErrorMessage(model, error instanceof Error ? error.message : String(error)),
-          });
+          const isStaleCtx = error instanceof Error && error.message.includes('stale');
+          if (isStaleCtx) {
+            stream.push({
+              type: 'done',
+              reason: 'stop',
+              message: createErrorMessage(model, ''),
+            });
+          } else {
+            stream.push({
+              type: 'error',
+              reason: 'error',
+              error: createErrorMessage(model, error instanceof Error ? error.message : String(error)),
+            });
+          }
           stream.end();
         } finally {
-          actions.persistState();
+          try {
+            actions.persistState();
+          } catch {
+            // Ignore: extension context may be stale after session teardown.
+          }
         }
       })();
 

package/extensions/routing.ts

@@ -1,8 +1,13 @@
-import { streamSimple } from '@earendil-works/pi-ai/compat';
-import type { Context, Message } from '@earendil-works/pi-ai';
+import type { Context, Message, ProviderHeaders, SimpleStreamOptions } from '@earendil-works/pi-ai';
 import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
 import type { RouterConfig, RouterProfile, RouterTier, RoutingDecision } from './types';
-import { createOpenRouterOnPayload, OPENROUTER_ATTR_HEADERS, parseCanonicalModelRef, resolveModelFromRef } from './config';
+import {
+  createOpenRouterOnPayload,
+  dispatchStream,
+  OPENROUTER_ATTR_HEADERS,
+  parseCanonicalModelRef,
+  resolveModelFromRef,
+} from './config';
 
 export const extractTextFromContent = (content: string | Message['content']): string => {
   if (typeof content === 'string') {
@@ -27,7 +32,8 @@ const lastUserMessage = (context: Context): Message | undefined => {
   return undefined;
 };
 
-const escapeXML = (s: string): string => s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;').replace(/"/g, '&quot;');
+const escapeXML = (s: string): string =>
+  s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;').replace(/"/g, '&quot;');
 
 const getRecentConversationText = (context: Context, limit: number): string => {
   const messages = context.messages.slice(-limit);
@@ -239,25 +245,32 @@ export const runClassifier = async (
 
     const isOpenRouter = parseCanonicalModelRef(classifierModelRef).provider === 'openrouter';
 
-    const classifierOptions: Record<string, unknown> = {
+    const effectiveClassifierHeaders: ProviderHeaders = {
+      ...model.headers,
+      ...(isOpenRouter ? OPENROUTER_ATTR_HEADERS : {}),
+      ...auth.headers,
+    };
+
+    const classifierOptions: SimpleStreamOptions = {
       apiKey: auth.apiKey,
-      headers: {
-        ...(auth.headers ?? {}),
-        ...(isOpenRouter ? OPENROUTER_ATTR_HEADERS : {}),
-      },
+      headers: effectiveClassifierHeaders,
       ...(thinking && thinking !== 'off' ? { reasoning: thinking } : {}),
     };
 
-
     if (isOpenRouter) {
       const onPayload = createOpenRouterOnPayload(extCtx?.sessionManager);
       if (onPayload) classifierOptions.onPayload = onPayload;
     }
 
+    // Apply credential-specific baseUrl override if resolved by getApiKeyAndHeaders
+    // (e.g. GitHub Copilot business/enterprise tenant endpoints) to avoid 421 Misdirected Request.
+    const authBaseUrl = (auth as { baseUrl?: string }).baseUrl;
+    const requestModel = authBaseUrl ? { ...model, baseUrl: authBaseUrl } : model;
+
     const MAX_CLASSIFIER_ATTEMPTS = 3;
     for (let attempt = 1; attempt <= MAX_CLASSIFIER_ATTEMPTS; attempt++) {
       try {
-        const stream = streamSimple(model, classifierContext, classifierOptions);
+        const stream = dispatchStream(requestModel, classifierContext, classifierOptions, modelRegistry);
         let fullText = '';
         for await (const event of stream) {
           if (event.type === 'error') throw new Error(event.error?.errorMessage ?? 'Unknown classifier error');

package/extensions/state.ts

@@ -1,21 +1,11 @@
-import type {
-  RouterPinByProfile,
-  RoutingDecision,
-  RouterPersistedState,
-} from './types';
+import type { RouterPinByProfile, RoutingDecision, RouterPersistedState } from './types';
 
-export const isRouterPersistedState = (
-  value: unknown,
-): value is RouterPersistedState => {
+export const isRouterPersistedState = (value: unknown): value is RouterPersistedState => {
   if (typeof value !== 'object' || value === null) {
     return false;
   }
   const v = value as Record<string, unknown>;
-  return (
-    typeof v.enabled === 'boolean' &&
-    typeof v.selectedProfile === 'string' &&
-    typeof v.timestamp === 'number'
-  );
+  return typeof v.enabled === 'boolean' && typeof v.selectedProfile === 'string' && typeof v.timestamp === 'number';
 };
 
 export const buildPersistedState = (

package/extensions/ui.ts

@@ -1,8 +1,5 @@
 import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
-import type {
-  RoutingDecision,
-  RouterPinByProfile,
-} from './types';
+import type { RoutingDecision, RouterPinByProfile } from './types';
 
 const getDecisionFlags = (decision: RoutingDecision): string[] => {
   const flags: string[] = [];
@@ -15,9 +12,7 @@ export const formatDecision = (d: RoutingDecision): string => {
   return `[${new Date(d.timestamp).toLocaleTimeString()}] ${d.tier} -> ${d.targetProvider}/${d.targetModelId} (${d.thinking}) - ${d.reasoning}`;
 };
 
-export const formatPinSummary = (
-  pinnedTierByProfile: RouterPinByProfile,
-): string => {
+export const formatPinSummary = (pinnedTierByProfile: RouterPinByProfile): string => {
   const entries = Object.entries(pinnedTierByProfile)
     .sort(([a], [b]) => a.localeCompare(b))
     .map(([profile, tier]) => `${profile}:${tier}`);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@kdejaeger/pi-model-router",
-  "version": "0.4.2",
+  "version": "0.4.5",
   "type": "module",
   "description": "Intelligent per-turn model router extension for the pi coding agent",
   "keywords": [
@@ -42,14 +42,14 @@
     "prepublishOnly": "npm run tsc"
   },
   "peerDependencies": {
-    "@earendil-works/pi-agent-core": "^0.80.6",
-    "@earendil-works/pi-ai": "^0.80.6",
-    "@earendil-works/pi-coding-agent": "^0.80.6",
-    "@earendil-works/pi-tui": "^0.80.6",
+    "@earendil-works/pi-agent-core": "^0.80",
+    "@earendil-works/pi-ai": "^0.80",
+    "@earendil-works/pi-coding-agent": "^0.80",
+    "@earendil-works/pi-tui": "^0.80",
     "@sinclair/typebox": "*"
   },
   "devDependencies": {
-    "prettier": "^3.8.1",
-    "typescript": "^6.0.2"
+    "prettier": "^3.8",
+    "typescript": "^6"
   }
 }