@alexeiled/pi-model-router 0.5.0 → 0.5.2

package/extensions/provider.ts

@@ -1,20 +1,20 @@
+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,
 } from '@earendil-works/pi-coding-agent';
+import { runClassifier } from './classifier';
 import {
   clampThinkingLevel,
   collectProfileThinkingLevels,
@@ -25,13 +25,14 @@ import {
   resolveContextWindow,
   resolveMaxTokens,
 } from './config';
+import { DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from './constants';
+import { extractTextFromContent, hasImageAttachment } from './context';
 import {
-  DEFAULT_CONTEXT_WINDOW,
-  DEFAULT_MAX_TOKENS,
-  hasUsableRequestAuth,
-  type RegistryWithProviderAuth,
-  resolveDelegatedModel,
-} from './constants';
+  buildRoutingDecision,
+  decideRouting,
+  phaseForTier,
+  resolveAvailableTier,
+} from './routing';
 import type {
   RouterConfig,
   RouterPinByProfile,
@@ -44,17 +45,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,28 +57,25 @@ 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;
 };
 
-import {
-  buildRoutingDecision,
-  decideRouting,
-  extractTextFromContent,
-  hasImageAttachment,
-  phaseForTier,
-  runClassifier,
-} from './routing';
-
 export const createErrorMessage = (
   model: Model<Api>,
   message: string,
@@ -120,7 +107,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 +126,32 @@ 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);
-
-  let startIndex = 0;
-  while (startIndex < messages.length) {
-    const currentTokens = systemTokens + latestTokens + activeMessagesTokensSum;
-    if (currentTokens <= limit) break;
-
-    activeMessagesTokensSum -= messageTokens[startIndex];
-    startIndex++;
+  // Drop only complete turns. Splitting an assistant/tool-result pair corrupts transcripts.
+  // This text estimate cannot guarantee a fit for images/tools or one oversized active turn.
+  const systemMessages = messages.filter(
+    (message) => message.role === 'system',
+  );
+  let remaining = totalTokens;
+  let nextRemovableIndex = 0;
+  for (let i = 0; i < messages.length && remaining > limit; i += 1) {
+    if (messages[i]?.role !== 'user') continue;
+    while (nextRemovableIndex < i) {
+      const candidate = messages[nextRemovableIndex];
+      if (candidate?.role !== 'system') {
+        remaining -= messageTokens[nextRemovableIndex] ?? 0;
+      }
+      nextRemovableIndex += 1;
+    }
   }
-
-  const finalMessages = [...messages.slice(startIndex), latestMessage];
-  return { ...context, messages: finalMessages };
+  return {
+    ...context,
+    messages: [
+      ...systemMessages,
+      ...messages
+        .slice(nextRemovableIndex)
+        .filter((message) => message.role !== 'system'),
+    ],
+  };
 };
 
 const supportsReasoning = (
@@ -213,24 +208,29 @@ export const registerRouterProvider = (
   const profileList = profileNames(state.currentConfig);
 
   // Map profiles to their capacities
-  const modelDefinitions = profileList.map((name) => {
+  const modelDefinitions = profileList.flatMap((name) => {
     const profile = state.currentConfig.profiles[name];
+    if (!profile) return [];
 
     // Report the MAX context window and max output tokens across all tiers.
     // The honesty check + truncateContext handles the case where the
     // actually routed model is smaller.
-    let maxContextWindow = DEFAULT_CONTEXT_WINDOW;
-    let maxMaxTokens = DEFAULT_MAX_TOKENS;
+    let maxContextWindow = 0;
+    let maxOutputTokens = 0;
     for (const tier of ROUTER_TIERS) {
       if (!profile[tier]) continue;
-      const cw = resolveContextWindow(
+      const contextWindow = resolveContextWindow(
         tier,
         profile,
         state.currentModelRegistry,
       );
-      const mot = resolveMaxTokens(tier, profile, state.currentModelRegistry);
-      if (cw > maxContextWindow) maxContextWindow = cw;
-      if (mot > maxMaxTokens) maxMaxTokens = mot;
+      const maxTokens = resolveMaxTokens(
+        tier,
+        profile,
+        state.currentModelRegistry,
+      );
+      if (contextWindow > maxContextWindow) maxContextWindow = contextWindow;
+      if (maxTokens > maxOutputTokens) maxOutputTokens = maxTokens;
     }
 
     const hasReasoning = supportsReasoning(profile, state.currentModelRegistry);
@@ -245,21 +245,21 @@ export const registerRouterProvider = (
       if (Object.keys(map).length > 0) thinkingLevelMap = map;
     }
 
-    return {
-      id: name,
-      name: `Router ${name}`,
-      reasoning: hasReasoning,
-      ...(thinkingLevelMap ? { thinkingLevelMap } : {}),
-      input: ['text', 'image'] as ('text' | 'image')[],
-      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-      contextWindow: maxContextWindow,
-      maxTokens: maxMaxTokens,
-    };
+    return [
+      {
+        id: name,
+        name: `Router ${name}`,
+        reasoning: hasReasoning,
+        ...(thinkingLevelMap ? { thinkingLevelMap } : {}),
+        input: ['text', 'image'] satisfies ('text' | 'image')[],
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        contextWindow: maxContextWindow || DEFAULT_CONTEXT_WINDOW,
+        maxTokens: maxOutputTokens || DEFAULT_MAX_TOKENS,
+      },
+    ];
   });
 
-  const modelsKey = 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', {
@@ -274,7 +274,8 @@ export const registerRouterProvider = (
     ): AssistantMessageEventStream {
       const stream = createAssistantMessageEventStream();
 
-      (async () => {
+      void (async () => {
+        let partialMessage: AssistantMessage | undefined;
         try {
           // Wait for the router to be fully initialized (session_start sets currentModelRegistry).
           // This handles the race where subagents (e.g. from pi-dynamic-workflows) invoke
@@ -282,6 +283,7 @@ export const registerRouterProvider = (
           const registry = await waitForRegistry(
             state,
             state.registryTimeoutMs,
+            options?.signal,
           );
           if (!registry) {
             throw new Error(
@@ -325,15 +327,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,14 +453,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++) {
-            const modelRef = modelsToTry[i];
+          for (const [i, modelRef] of modelsToTry.entries()) {
+            options?.signal?.throwIfAborted();
             const { provider: targetProvider, modelId: targetModelId } =
               parseCanonicalModelRef(modelRef);
 
@@ -467,31 +476,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
@@ -509,15 +501,15 @@ export const registerRouterProvider = (
                 const tierConfig = profile[decision.tier];
                 if (tierConfig?.resolvedThinkingLevels) {
                   requestedReasoning = clampThinkingLevel(
-                    requestedReasoning as ThinkingLevel,
+                    requestedReasoning,
                     tierConfig.resolvedThinkingLevels,
-                  ) as typeof requestedReasoning;
+                  );
                 }
               }
 
-              const delegatedReasoning =
+              const delegatedReasoning: SimpleStreamOptions['reasoning'] =
                 targetModel.reasoning && requestedReasoning !== 'off'
-                  ? (requestedReasoning as SimpleStreamOptions['reasoning'])
+                  ? requestedReasoning
                   : undefined;
 
               try {
@@ -534,42 +526,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 = {
+              const delegatedOptions: SimpleStreamOptions = {
                 ...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 +579,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 +626,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 {