@alexeiled/pi-model-router 0.5.0 → 0.5.2

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/context.ts

@@ -0,0 +1,58 @@
+import type { Context, Message } from '@earendil-works/pi-ai';
+
+export const extractTextFromContent = (
+  content: string | Message['content'],
+): string => {
+  if (typeof content === 'string') return content;
+  return content
+    .map((part) => {
+      if (part.type === 'text') return part.text;
+      if (part.type === 'thinking') return part.thinking;
+      if (part.type === 'toolCall') {
+        return `${part.name} ${JSON.stringify(part.arguments)}`;
+      }
+      return '';
+    })
+    .filter(Boolean)
+    .join('\n');
+};
+
+export const getLastUserText = (context: Context): string => {
+  for (let i = context.messages.length - 1; i >= 0; i -= 1) {
+    const message = context.messages[i];
+    if (message?.role === 'user') {
+      return extractTextFromContent(message.content).trim();
+    }
+  }
+  return '';
+};
+
+export const getRecentConversationText = (
+  context: Context,
+  limit = 6,
+): string =>
+  context.messages
+    .slice(-limit)
+    .map((message) =>
+      message ? extractTextFromContent(message.content).trim() : '',
+    )
+    .filter(Boolean)
+    .join('\n')
+    .toLowerCase();
+
+export const countToolResults = (context: Context): number =>
+  context.messages.filter((message) => message?.role === 'toolResult').length;
+
+export const countWords = (text: string): number =>
+  text.split(/\s+/).filter(Boolean).length;
+
+export const hasImageAttachment = (context: Context): boolean =>
+  context.messages.some(
+    (message) =>
+      message &&
+      Array.isArray(message.content) &&
+      message.content.some((part) => part.type === 'image'),
+  );
+
+export const containsAny = (text: string, keywords: string[]): boolean =>
+  keywords.some((keyword) => text.includes(keyword));

package/extensions/index.ts

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