@alexeiled/pi-model-router 0.5.0 → 0.5.1

package/CHANGELOG.md

@@ -1,6 +1,18 @@
 # Changelog
 
-## 0.5.0 (unreleased)
+## 0.5.1 — 2026-09-20
+
+- Delegate generation and classification through Pi's native model registry instead of duplicating auth/dispatch logic. Cover keyless and headers-only auth, native providers and credential URLs with in-memory SDK integration tests.
+- Retry only before output; preserve aborts and partial output, reject unterminated streams, and record the actual fallback model.
+- Use each attempted model's context limit and trim complete turns without discarding system messages or orphaning tool results.
+- Resolve classifier choices against partial profiles, bound classifier requests, and reject malformed/error responses.
+- Deep-copy and validate persisted state; reset snapshot deduplication across branches.
+- Validate rule keywords and malformed profiles; avoid inherited-name lookups and respect non-reasoning tier declarations.
+- Reject invalid debug/widget options, fix thinking completions, and avoid success notifications after failed switches.
+- Replace sleep-based stream tests with real event streams; share typed fixtures and keep test helpers out of npm artifacts.
+- Fail closed on npm registry lookup errors other than a missing version. Document the release and provenance verification procedure.
+
+## 0.5.0 — 2026-09-20
 
 This is the first release of the independently maintained `@alexeiled/pi-model-router` fork.
 

package/README.md

@@ -1,6 +1,13 @@
 # pi-model-router
 
-Smart per-turn model router extension for the [pi-coding-agent](https://github.com/earendil-works/pi/tree/main/packages/coding-agent) that optimizes your AI budget and usage limits without sacrificing quality by dynamically routing each turn to the optimal LLM tier. It automatically selects between high, medium, and low-tier models based on task intent, session budget, context size, and custom rules — complete with automatic fallbacks and phase awareness.
+[![CI](https://github.com/alexei-led/pi-model-router/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/alexei-led/pi-model-router/actions/workflows/ci.yml)
+[![npm version](https://img.shields.io/npm/v/%40alexeiled%2Fpi-model-router?logo=npm)](https://www.npmjs.com/package/@alexeiled/pi-model-router)
+[![npm downloads](https://img.shields.io/npm/dm/%40alexeiled%2Fpi-model-router?logo=npm)](https://www.npmjs.com/package/@alexeiled/pi-model-router)
+[![Latest release](https://img.shields.io/github/v/release/alexei-led/pi-model-router?display_name=tag&sort=semver)](https://github.com/alexei-led/pi-model-router/releases)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
+[![Node.js >=22.19](https://img.shields.io/badge/node-%3E%3D22.19-339933?logo=node.js&logoColor=white)](package.json)
+
+Per-turn model router for [Pi](https://github.com/earendil-works/pi/tree/main/packages/coding-agent). Selects high, medium or low-tier models using task intent, a soft budget policy and custom rules, while keeping the selected `router/<profile>` model stable.
 
 > **Independent fork:** This project is an independently maintained fork of [yeliu84/pi-model-router](https://github.com/yeliu84/pi-model-router), originally created by Ye Liu. It is not an official upstream release. The original MIT license and copyright notice are preserved.
 
@@ -82,6 +89,15 @@ Or load directly for one run:
 pi -e ./extensions/index.ts
 ```
 
+## Reliability
+
+- Both generation and classification use Pi's provider registry, including native/custom providers and credential-specific URLs.
+- Fallbacks run only before content is emitted; cancellation does not retry.
+- Classifier requests use isolated context, a 10-second cancellation deadline and a 256-token output limit. Failures retain local routing.
+- Context trimming preserves system instructions and whole active tool turns. It is a text estimate, not a guarantee that images or a large active turn fit.
+
+See [architecture](https://github.com/alexei-led/pi-model-router/blob/main/docs/ARCHITECTURE.md) and [release procedure](https://github.com/alexei-led/pi-model-router/blob/main/docs/RELEASING.md).
+
 ## Configuration
 
 Copy the example config to one of:
@@ -112,7 +128,7 @@ The extension stores the last selected profile in `~/.pi/agent/model-router-stat
 | Field                   | Description                                                                       |
 | ----------------------- | --------------------------------------------------------------------------------- |
 | `classifierModel`       | (Optional) Model used to categorize intent. Supports model aliases. If omitted, fast heuristics are used. |
-| `maxSessionBudget`      | (Optional) USD budget for the session. Forces `medium` tier once exceeded.        |
+| `maxSessionBudget`      | (Optional) Soft generation-cost threshold in USD. Downgrades high to medium, or low if medium is absent. Not a spending cap; classifier cost is excluded. |
 | `phaseBias`             | (0.0 - 1.0) Stickiness of the current phase. Higher = more stable. Default `0.5`. |
 | `rules`                 | List of custom keyword rules (e.g. `{ "matches": "deploy", "tier": "high" }`).    |
 | `models`                | (Optional) Map of model aliases to definitions with `model`, `contextWindow`, `maxTokens`. |

package/extensions/commands.ts

@@ -136,10 +136,6 @@ export const registerCommands = (
       ];
     }
 
-    if (levelValues.includes(args[0])) {
-      return null;
-    }
-
     if ((tierValues as string[]).includes(args[0])) {
       const tier = args[0];
       const levelPrefix = args[1] ?? '';
@@ -443,6 +439,10 @@ export const registerCommands = (
       return;
     }
     const cmd = args[0]?.toLowerCase();
+    if (cmd && !['on', 'off', 'toggle'].includes(cmd)) {
+      ctx.ui.notify('Usage: /router widget <on|off|toggle>', 'error');
+      return;
+    }
     if (cmd === 'on') state.widgetEnabled = true;
     else if (cmd === 'off') state.widgetEnabled = false;
     else state.widgetEnabled = !state.widgetEnabled;
@@ -460,6 +460,10 @@ export const registerCommands = (
       return;
     }
     const cmd = args[0]?.toLowerCase();
+    if (cmd && !['on', 'off', 'toggle', 'clear', 'show'].includes(cmd)) {
+      ctx.ui.notify('Usage: /router debug <on|off|toggle|show|clear>', 'error');
+      return;
+    }
     if (cmd === 'on') state.debugEnabled = true;
     else if (cmd === 'off') state.debugEnabled = false;
     else if (cmd === 'clear') state.debugHistory.length = 0;
@@ -656,11 +660,12 @@ export const registerCommands = (
                 );
                 return;
               }
-              await actions.switchToRouterProfile(subcommand, ctx);
-              ctx.ui.notify(
-                `Router enabled with profile: ${state.selectedProfile}`,
-                'info',
-              );
+              if (await actions.switchToRouterProfile(subcommand, ctx)) {
+                ctx.ui.notify(
+                  `Router enabled with profile: ${state.selectedProfile}`,
+                  'info',
+                );
+              }
             } else {
               ctx.ui.notify(
                 `Unknown router subcommand: ${subcommand}. Try /router help`,

package/extensions/config.ts

@@ -82,7 +82,8 @@ export const resolveModelRef = (
   ref: string,
   models: Record<string, ModelDefinition> | undefined,
 ): { canonicalRef: string; definition?: ModelDefinition } => {
-  const definition = models?.[ref];
+  const definition =
+    models && Object.hasOwn(models, ref) ? models[ref] : undefined;
   if (definition) {
     return { canonicalRef: definition.model, definition };
   }
@@ -104,8 +105,13 @@ export const mergeConfig = (
   override: Partial<RouterConfig>,
 ): RouterConfig => {
   const mergedProfiles: Record<string, RouterProfile> = { ...base.profiles };
-  for (const [name, profile] of Object.entries(override.profiles ?? {})) {
-    const existing = mergedProfiles[name];
+  for (const [name, profile] of Object.entries(
+    isObjectRecord(override.profiles) ? override.profiles : {},
+  )) {
+    if (!isObjectRecord(profile) || name === '__proto__') continue;
+    const existing = Object.hasOwn(mergedProfiles, name)
+      ? mergedProfiles[name]
+      : undefined;
     const nextProfile = profile as Partial<RouterProfile>;
     mergedProfiles[name] = {
       high: mergeTier(existing?.high, nextProfile.high),
@@ -116,7 +122,7 @@ export const mergeConfig = (
 
   const mergedModels: Record<string, ModelDefinition> = {
     ...(base.models ?? {}),
-    ...(override.models ?? {}),
+    ...(isObjectRecord(override.models) ? override.models : {}),
   };
 
   return {
@@ -160,6 +166,7 @@ export const normalizeModelsMap = (
   if (!raw || !isObjectRecord(raw)) return result;
 
   for (const [alias, entry] of Object.entries(raw)) {
+    if (alias === '__proto__') continue;
     if (!isObjectRecord(entry)) {
       warnings.push(
         `Ignored invalid model definition "${alias}": expected an object.`,
@@ -331,6 +338,7 @@ export const normalizeTierConfig = (
   const resolvedThinkingLevels: ThinkingLevel[] = [...baseThinkingLevels];
   if (
     !explicitThinkingLevels &&
+    effectiveReasoning !== false &&
     thinking !== 'off' &&
     !resolvedThinkingLevels.includes(thinking)
   ) {
@@ -363,7 +371,10 @@ export const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
 
   const normalizedProfiles: Record<string, RouterProfile> = {};
 
-  for (const [name, profile] of Object.entries(raw.profiles ?? {})) {
+  for (const [name, profile] of Object.entries(
+    isObjectRecord(raw.profiles) ? raw.profiles : {},
+  )) {
+    if (name === '__proto__') continue;
     const high = normalizeTierConfig(
       profile?.high,
       name,
@@ -411,7 +422,12 @@ export const normalizeConfig = (raw: RouterConfig): ConfigLoadResult => {
         const matches = rule.matches;
         const tier = rule.tier;
         if (
-          (typeof matches === 'string' || Array.isArray(matches)) &&
+          ((typeof matches === 'string' && matches.trim().length > 0) ||
+            (Array.isArray(matches) &&
+              matches.length > 0 &&
+              matches.every(
+                (m) => typeof m === 'string' && m.trim().length > 0,
+              ))) &&
           isRouterTier(tier)
         ) {
           rules.push({
@@ -518,7 +534,7 @@ export const resolveProfileName = (
   config: RouterConfig,
   requested?: string,
 ): string | undefined => {
-  if (requested && config.profiles[requested]) {
+  if (requested && Object.hasOwn(config.profiles, requested)) {
     return requested;
   }
   return undefined;

package/extensions/constants.ts

@@ -1,49 +1,3 @@
 export const MAX_DEBUG_HISTORY = 12;
 export const DEFAULT_CONTEXT_WINDOW = 128_000;
 export const DEFAULT_MAX_TOKENS = 16_384;
-
-const AUTH_HEADERS = new Set([
-  'authorization',
-  'x-api-key',
-  'cf-aig-authorization',
-]);
-
-export interface RegistryWithProviderAuth {
-  getProviderAuth?: (
-    provider: string,
-  ) => Promise<{ auth: { baseUrl?: string } } | undefined>;
-}
-
-export const hasUsableRequestAuth = (auth: {
-  apiKey?: string;
-  headers?: Record<string, string | null | undefined>;
-}): boolean => {
-  if (typeof auth.apiKey === 'string' && auth.apiKey.trim().length > 0) {
-    return true;
-  }
-
-  return Object.entries(auth.headers ?? {}).some(
-    ([name, value]) =>
-      AUTH_HEADERS.has(name.toLowerCase()) &&
-      typeof value === 'string' &&
-      value.trim().length > 0,
-  );
-};
-
-export const resolveDelegatedModel = async <
-  TModel extends { provider: string; baseUrl: string },
->(
-  registry: RegistryWithProviderAuth,
-  model: TModel,
-): Promise<TModel> => {
-  try {
-    const providerAuth = await registry.getProviderAuth?.(model.provider);
-    const authBaseUrl = providerAuth?.auth.baseUrl;
-    if (authBaseUrl && authBaseUrl !== model.baseUrl) {
-      return { ...model, baseUrl: authBaseUrl };
-    }
-  } catch {
-    // Older Pi versions and unavailable credentials use the model's static URL.
-  }
-  return model;
-};

package/extensions/index.ts

@@ -199,7 +199,6 @@ const routerExtension = (pi: ExtensionAPI) => {
 
       // Ensure the provider is registered with current capacities for this profile
       actions.registerRouterProvider();
-      await new Promise((resolve) => setTimeout(resolve, 50));
 
       const routerModel = ctx.modelRegistry.find('router', profileName);
       if (!routerModel) {
@@ -284,6 +283,7 @@ const routerExtension = (pi: ExtensionAPI) => {
     ctx: ExtensionContext,
     startReason: SessionStartEvent['reason'],
   ) => {
+    lastPersistedSnapshot = undefined;
     ignoreStartupThinkingEvent =
       startReason === 'startup' ||
       startReason === 'new' ||
@@ -295,9 +295,6 @@ const routerExtension = (pi: ExtensionAPI) => {
     const hasExplicitStartupModel =
       startReason === 'startup' && hasExplicitCliModel();
 
-    // Give the registry a moment to synchronize after re-registration
-    await new Promise((resolve) => setTimeout(resolve, 50));
-
     routerEnabled = ctx.model?.provider === 'router';
     selectedProfile =
       ctx.model?.provider === 'router'
@@ -342,7 +339,10 @@ const routerExtension = (pi: ExtensionAPI) => {
         Object.assign(pinnedTierByProfile, savedState.pinByProfile);
       }
       if (savedState.thinkingByProfile) {
-        Object.assign(thinkingByProfile, savedState.thinkingByProfile);
+        Object.assign(
+          thinkingByProfile,
+          structuredClone(savedState.thinkingByProfile),
+        );
       }
       if (savedState.pinTier && selectedProfile) {
         pinnedTierByProfile[selectedProfile] = savedState.pinTier;
@@ -350,12 +350,14 @@ const routerExtension = (pi: ExtensionAPI) => {
       debugEnabled = savedState.debugEnabled ?? debugEnabled;
       widgetEnabled = savedState.widgetEnabled ?? widgetEnabled;
       debugHistory = savedState.debugHistory
-        ? [...savedState.debugHistory].slice(-MAX_DEBUG_HISTORY)
+        ? structuredClone(savedState.debugHistory).slice(-MAX_DEBUG_HISTORY)
         : [];
       if (!hasExplicitStartupModel) {
         lastNonRouterModel =
           savedState.lastNonRouterModel ?? lastNonRouterModel;
-        lastDecision = savedState.lastDecision;
+        lastDecision = savedState.lastDecision
+          ? structuredClone(savedState.lastDecision)
+          : undefined;
       }
       accumulatedCost = savedState.accumulatedCost ?? 0;
     } else if (

package/extensions/provider.ts

@@ -1,16 +1,15 @@
+import { setTimeout as delay } from 'node:timers/promises';
 import type { ThinkingLevel } from '@earendil-works/pi-agent-core';
 import {
   type Api,
   type AssistantMessage,
+  type AssistantMessageEvent,
   type AssistantMessageEventStream,
   type Context,
   createAssistantMessageEventStream,
   type Model,
-  normalizeContext,
   type SimpleStreamOptions,
-  type TranscriptContext,
 } from '@earendil-works/pi-ai';
-import { streamSimple } from '@earendil-works/pi-ai/compat';
 import type {
   ExtensionAPI,
   ExtensionContext,
@@ -25,13 +24,7 @@ import {
   resolveContextWindow,
   resolveMaxTokens,
 } from './config';
-import {
-  DEFAULT_CONTEXT_WINDOW,
-  DEFAULT_MAX_TOKENS,
-  hasUsableRequestAuth,
-  type RegistryWithProviderAuth,
-  resolveDelegatedModel,
-} from './constants';
+import { DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from './constants';
 import type {
   RouterConfig,
   RouterPinByProfile,
@@ -44,17 +37,6 @@ const REGISTRY_WAIT_TIMEOUT_MS = 5000;
 const REGISTRY_WAIT_INITIAL_DELAY_MS = 50;
 const REGISTRY_WAIT_MAX_DELAY_MS = 500;
 
-type ProviderAwareRegistry = ExtensionContext['modelRegistry'] & {
-  getRegisteredProviderConfig?: (provider: string) => {
-    api?: Api;
-    streamSimple?: (
-      model: Model<Api>,
-      context: TranscriptContext,
-      options?: SimpleStreamOptions,
-    ) => AssistantMessageEventStream;
-  };
-};
-
 /**
  * Wait for the model registry to become available with exponential backoff.
  * This handles the race condition where subagents (e.g. from pi-dynamic-workflows)
@@ -67,15 +49,21 @@ export const waitForRegistry = async (
       | undefined;
   },
   timeoutMs: number = REGISTRY_WAIT_TIMEOUT_MS,
+  signal?: AbortSignal,
 ): Promise<ExtensionContext['modelRegistry'] | undefined> => {
+  signal?.throwIfAborted();
   if (state.currentModelRegistry) return state.currentModelRegistry;
 
   const start = Date.now();
-  let delay = REGISTRY_WAIT_INITIAL_DELAY_MS;
+  let interval = REGISTRY_WAIT_INITIAL_DELAY_MS;
   while (Date.now() - start < timeoutMs) {
-    await new Promise((resolve) => setTimeout(resolve, delay));
+    await delay(
+      Math.min(interval, timeoutMs - (Date.now() - start)),
+      undefined,
+      { signal },
+    );
     if (state.currentModelRegistry) return state.currentModelRegistry;
-    delay = Math.min(delay * 2, REGISTRY_WAIT_MAX_DELAY_MS);
+    interval = Math.min(interval * 2, REGISTRY_WAIT_MAX_DELAY_MS);
   }
   return undefined;
 };
@@ -86,6 +74,7 @@ import {
   extractTextFromContent,
   hasImageAttachment,
   phaseForTier,
+  resolveAvailableTier,
   runClassifier,
 } from './routing';
 
@@ -120,7 +109,7 @@ const estimateTokens = (text: string): number => Math.ceil(text.length / 3);
 
 /**
  * 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 system prompt and the complete latest user/tool turn.
  */
 const truncateContext = (context: Context, limit: number): Context => {
   const messages = [...context.messages];
@@ -139,24 +128,23 @@ const truncateContext = (context: Context, limit: number): Context => {
 
   if (totalTokens <= limit) return context;
 
-  const latestMessage = messages.pop();
-  if (!latestMessage) return context;
-  const latestTokens = messageTokens.pop() ?? 0;
-
-  // Keep shifting oldest messages from the start of the list
-  let activeMessagesTokensSum = messageTokens.reduce((sum, t) => sum + t, 0);
-
+  // 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.
+  let remaining = totalTokens;
   let startIndex = 0;
-  while (startIndex < messages.length) {
-    const currentTokens = systemTokens + latestTokens + activeMessagesTokensSum;
-    if (currentTokens <= limit) break;
-
-    activeMessagesTokensSum -= messageTokens[startIndex];
-    startIndex++;
+  for (let i = 1; i < messages.length && remaining > limit; i++) {
+    if (messages[i].role !== 'user') continue;
+    while (startIndex < i) remaining -= messageTokens[startIndex++];
   }
-
-  const finalMessages = [...messages.slice(startIndex), latestMessage];
-  return { ...context, messages: finalMessages };
+  return {
+    ...context,
+    messages: [
+      ...messages
+        .slice(0, startIndex)
+        .filter((message) => message.role === 'system'),
+      ...messages.slice(startIndex),
+    ],
+  };
 };
 
 const supportsReasoning = (
@@ -219,8 +207,8 @@ export const registerRouterProvider = (
     // 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 = DEFAULT_CONTEXT_WINDOW;
-    let maxMaxTokens = DEFAULT_MAX_TOKENS;
+    let maxContextWindow = 0;
+    let maxMaxTokens = 0;
     for (const tier of ROUTER_TIERS) {
       if (!profile[tier]) continue;
       const cw = resolveContextWindow(
@@ -252,14 +240,12 @@ export const registerRouterProvider = (
       ...(thinkingLevelMap ? { thinkingLevelMap } : {}),
       input: ['text', 'image'] as ('text' | 'image')[],
       cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-      contextWindow: maxContextWindow,
-      maxTokens: maxMaxTokens,
+      contextWindow: maxContextWindow || DEFAULT_CONTEXT_WINDOW,
+      maxTokens: maxMaxTokens || DEFAULT_MAX_TOKENS,
     };
   });
 
-  const modelsKey = modelDefinitions
-    .map((m) => `${m.id}:${m.contextWindow}:${m.maxTokens}:${m.reasoning}`)
-    .join(',');
+  const modelsKey = JSON.stringify(modelDefinitions);
   if (state.lastRegisteredModels === modelsKey) return;
 
   pi.registerProvider('router', {
@@ -275,6 +261,7 @@ export const registerRouterProvider = (
       const stream = createAssistantMessageEventStream();
 
       (async () => {
+        let partialMessage: AssistantMessage | undefined;
         try {
           // Wait for the router to be fully initialized (session_start sets currentModelRegistry).
           // This handles the race where subagents (e.g. from pi-dynamic-workflows) invoke
@@ -282,6 +269,7 @@ export const registerRouterProvider = (
           const registry = await waitForRegistry(
             state,
             state.registryTimeoutMs,
+            options?.signal,
           );
           if (!registry) {
             throw new Error(
@@ -325,15 +313,20 @@ export const registerRouterProvider = (
               state.currentConfig.classifierModel.model,
               registry,
               context,
-              state.lastDecision?.phase,
+              state.lastDecision?.profile === model.id
+                ? state.lastDecision.phase
+                : undefined,
               state.currentConfig.classifierModel.thinking,
+              options?.signal,
             );
+            options?.signal?.throwIfAborted();
             if (classifierResult) {
+              const tier = resolveAvailableTier(profile, classifierResult.tier);
               decision = buildRoutingDecision(
                 model.id,
                 profile,
-                classifierResult.tier,
-                phaseForTier(classifierResult.tier),
+                tier,
+                phaseForTier(tier),
                 `Classifier: ${classifierResult.reasoning}`,
                 state.thinkingByProfile[model.id],
                 true,
@@ -446,13 +439,16 @@ export const registerRouterProvider = (
           if (imageAttached) {
             modelsToTry = modelsToTry.filter(checkModelSupportsImage);
             if (modelsToTry.length === 0) {
-              modelsToTry = [decision.targetLabel];
+              throw new Error(
+                'No configured model supports image attachments.',
+              );
             }
           }
           let lastError: unknown;
           let success = false;
 
           for (let i = 0; i < modelsToTry.length; i++) {
+            options?.signal?.throwIfAborted();
             const modelRef = modelsToTry[i];
             const { provider: targetProvider, modelId: targetModelId } =
               parseCanonicalModelRef(modelRef);
@@ -467,31 +463,14 @@ export const registerRouterProvider = (
               continue;
             }
 
-            const auth = await registry.getApiKeyAndHeaders(targetModel);
-            if (!auth.ok || !hasUsableRequestAuth(auth)) {
-              lastError = new Error(
-                auth.ok
-                  ? `No API key or authentication headers for routed model: ${targetProvider}/${targetModelId}`
-                  : `Auth failed for routed model: ${targetProvider}/${targetModelId}: ${auth.error}`,
-              );
-              continue;
-            }
-            const apiKey = auth.apiKey;
-            const headers = auth.headers;
-            const requestModel = await resolveDelegatedModel(
-              registry as unknown as RegistryWithProviderAuth,
-              targetModel,
-            );
-
+            let contentReceived = false;
             try {
               // HONESTY CHECK & AUTO-TRUNCATION
               // If the picked model has a smaller context than what we reported, truncate now.
               let effectiveContext = context;
-              const targetLimit = resolveContextWindow(
-                decision.tier,
-                profile,
-                registry,
-              );
+              const targetLimit =
+                targetModel.contextWindow ??
+                resolveContextWindow(decision.tier, profile, registry);
               if (
                 model.contextWindow !== undefined &&
                 targetLimit < model.contextWindow
@@ -534,42 +513,42 @@ export const registerRouterProvider = (
                 // Stale extension context — skip non-critical UI updates.
               }
 
-              // Strip pi's reasoning from options — the router controls thinking
-              const { reasoning: _piReasoning, ...delegationOptions } =
-                options ?? {};
+              // Router credentials must not override the concrete provider's request auth.
+              const {
+                reasoning: _piReasoning,
+                apiKey: _routerKey,
+                headers: _routerHeaders,
+                env: _routerEnv,
+                ...delegationOptions
+              } = options ?? {};
 
               const delegatedOptions = {
                 ...delegationOptions,
-                apiKey,
-                headers,
                 ...(delegatedReasoning
                   ? { reasoning: delegatedReasoning }
                   : {}),
               };
-              const registeredProvider = (
-                registry as ProviderAwareRegistry
-              ).getRegisteredProviderConfig?.(targetProvider);
-              const delegatedStream =
-                registeredProvider?.streamSimple &&
-                registeredProvider.api === requestModel.api
-                  ? registeredProvider.streamSimple(
-                      requestModel,
-                      normalizeContext(effectiveContext),
-                      delegatedOptions,
-                    )
-                  : streamSimple(
-                      requestModel,
-                      effectiveContext,
-                      delegatedOptions,
-                    );
-
-              let contentReceived = false;
+              // Pi owns request-time auth, custom/native providers, URLs and transcript normalization.
+              const delegatedStream = registry.streamSimple(
+                targetModel,
+                effectiveContext,
+                delegatedOptions,
+              );
+              let terminalReceived = false;
+              const pendingEvents: AssistantMessageEvent[] = [];
+              const recordTarget = () => {
+                decision.targetProvider = targetProvider;
+                decision.targetModelId = targetModelId;
+                decision.targetLabel = modelRef;
+                decision.thinking = delegatedReasoning ?? 'off';
+                decision.isFallback = i > 0;
+              };
               for await (const event of delegatedStream) {
-                if (event.type === 'done') {
-                  const cost = event.message.usage?.cost?.total ?? 0;
-                  state.accumulatedCost += cost;
-                }
-                if (event.type === 'error' && !contentReceived) {
+                if (
+                  event.type === 'error' &&
+                  event.reason !== 'aborted' &&
+                  !contentReceived
+                ) {
                   const errorMessage =
                     'error' in event &&
                     event.error &&
@@ -587,13 +566,37 @@ export const registerRouterProvider = (
                   event.type === 'thinking_delta' ||
                   event.type === 'toolcall_delta' ||
                   event.type === 'toolcall_end';
-                if (isContent) contentReceived = true;
-                stream.push(event);
+                if (isContent) {
+                  contentReceived = true;
+                  recordTarget();
+                }
+                if (event.type === 'done' || event.type === 'error') {
+                  terminalReceived = true;
+                  recordTarget();
+                  const cost = (
+                    event.type === 'done' ? event.message : event.error
+                  ).usage.cost.total;
+                  if (Number.isFinite(cost) && cost > 0)
+                    state.accumulatedCost += cost;
+                }
+                if (contentReceived || terminalReceived) {
+                  for (const pending of pendingEvents.splice(0))
+                    stream.push(pending);
+                  stream.push(event);
+                  if ('partial' in event) partialMessage = event.partial;
+                } else {
+                  pendingEvents.push(event);
+                }
+                if (terminalReceived) break;
               }
+              if (!terminalReceived)
+                throw new Error(
+                  'Provider stream ended without a terminal event.',
+                );
               success = true;
-              if (i > 0) decision.isFallback = true;
               break;
             } catch (err) {
+              if (contentReceived || options?.signal?.aborted) throw err;
               lastError = err;
             }
           }
@@ -610,28 +613,17 @@ export const registerRouterProvider = (
 
           stream.end();
         } catch (error) {
-          // When a subagent session is torn down (e.g. by pi-dynamic-workflows),
-          // the extension runtime is invalidated and any pi/ctx call throws a
-          // stale-context error. Push a graceful done event so the stream's
-          // result() promise resolves (required by AssistantMessageEventStream).
-          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,
+          const reason = options?.signal?.aborted ? 'aborted' : 'error';
+          stream.push({
+            type: 'error',
+            reason,
+            error: {
+              ...(partialMessage ?? createErrorMessage(model, '')),
+              errorMessage:
                 error instanceof Error ? error.message : String(error),
-              ),
-            });
-          }
+              stopReason: reason,
+            },
+          });
           stream.end();
         } finally {
           try {

package/extensions/routing.ts

@@ -1,13 +1,7 @@
 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 {
-  hasUsableRequestAuth,
-  type RegistryWithProviderAuth,
-  resolveDelegatedModel,
-} from './constants';
 import type {
   RouterPhase,
   RouterProfile,
@@ -148,6 +142,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 +367,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);
@@ -399,20 +397,13 @@ export const runClassifier = async (
   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) 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,
-    );
+    if (!model || provider === 'router') return undefined;
+    signal?.throwIfAborted();
 
     const promptText = getLastUserText(context);
     const historyText = getRecentConversationText(context, 4);
@@ -447,17 +438,24 @@ ${currentPhase === 'implementation' ? 'Consider that the conversation is current
     const reasoningOption =
       model.reasoning && thinking && thinking !== 'off' ? thinking : undefined;
 
-    const stream = streamSimple(requestModel, classifierContext, {
-      apiKey,
-      headers,
+    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 === 'text_delta' && typeof event.delta === 'string') {
-        fullText += event.delta;
+      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:'));
@@ -471,7 +469,7 @@ ${currentPhase === 'implementation' ? 'Consider that the conversation is current
         return {
           tier: tierValue,
           reasoning: reasoningLine
-            ? reasoningLine.split(':')[1].trim()
+            ? reasoningLine.slice(reasoningLine.indexOf(':') + 1).trim()
             : 'Classifier decision.',
         };
       }

package/extensions/state.ts

@@ -1,6 +1,12 @@
 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 {
   RouterLastProfileState,
   RouterPersistedState,
@@ -11,13 +17,45 @@ import type {
 
 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,16 +96,35 @@ 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',
+    )
   );
 };
 
@@ -83,7 +140,7 @@ export const buildPersistedState = (
   lastNonRouterModel: string | undefined,
   accumulatedCost: number,
 ): RouterPersistedState => {
-  return {
+  return structuredClone({
     enabled: routerEnabled,
     selectedProfile: selectedProfile ?? '',
     pinTier: selectedProfile ? pinnedTierByProfile[selectedProfile] : undefined,
@@ -97,5 +154,5 @@ export const buildPersistedState = (
     lastNonRouterModel,
     accumulatedCost,
     timestamp: Date.now(),
-  };
+  });
 };

package/package.json

@@ -1,10 +1,11 @@
 {
   "name": "@alexeiled/pi-model-router",
-  "version": "0.5.0",
+  "version": "0.5.1",
   "type": "module",
   "files": [
     "extensions",
-    "!extensions/*.test.ts",
+    "!extensions/**/*.test.ts",
+    "!extensions/test",
     "LICENSE",
     "model-router.example.json",
     "README.md",