ultimate-pi 0.18.1 → 0.19.0

package/vendor/pi-model-router/extensions/index.ts

@@ -1,478 +0,0 @@
-import type {
-  ExtensionAPI,
-  ExtensionContext,
-} from '@earendil-works/pi-coding-agent';
-import {
-  type RouterConfig,
-  type RouterPersistedState,
-  type RoutingDecision,
-  type RouterPinByProfile,
-  type RouterThinkingByProfile,
-  type RouterTier,
-  type CustomSessionEntry,
-  type SessionLock,
-} from './types.js';
-import {
-  FALLBACK_CONFIG,
-  loadRouterConfig,
-  profileNames,
-  resolveProfileName,
-  parseCanonicalModelRef,
-} from './config.js';
-import { MAX_DEBUG_HISTORY } from './constants.js';
-import { isRouterPersistedState, buildPersistedState } from './state.js';
-import { updateStatus, formatModelRef } from './ui.js';
-import { registerCommands } from './commands.js';
-import { registerRouterProvider } from './provider.js';
-
-const routerExtension = (pi: ExtensionAPI) => {
-  let currentConfig: RouterConfig = FALLBACK_CONFIG;
-  let currentModelRegistry: ExtensionContext['modelRegistry'] | undefined;
-  let currentCwd = process.cwd();
-  let lastDecision: RoutingDecision | undefined;
-  let debugEnabled = false;
-  let routerEnabled = false;
-  let selectedProfile = resolveProfileName(
-    FALLBACK_CONFIG,
-    FALLBACK_CONFIG.defaultProfile,
-  );
-  let widgetEnabled = false;
-  let lastRegisteredModels = '';
-  let pinnedTierByProfile: RouterPinByProfile = {};
-  let thinkingByProfile: RouterThinkingByProfile = {};
-  let debugHistory: RoutingDecision[] = [];
-  let lastNonRouterModel: string | undefined;
-  let accumulatedCost = 0;
-  let lastExtensionContext: ExtensionContext | undefined;
-  let lastConfigWarnings: string[] = [];
-  let lastPersistedSnapshot: string | undefined;
-  let isInitialized = false;
-  let isInternalModelSwitch = false;
-  let sessionLock: SessionLock | undefined;
-
-  const setModelInternally = async (
-    model: NonNullable<ExtensionContext['model']>,
-  ) => {
-    isInternalModelSwitch = true;
-    try {
-      return await pi.setModel(model);
-    } finally {
-      isInternalModelSwitch = false;
-    }
-  };
-
-  const getPinnedTierForProfile = (
-    profileName: string,
-  ): RouterTier | undefined => pinnedTierByProfile[profileName];
-
-  const setPinnedTierForProfile = (
-    profileName: string,
-    tier: RouterTier | undefined,
-  ) => {
-    if (tier) {
-      pinnedTierByProfile[profileName] = tier;
-    } else {
-      delete pinnedTierByProfile[profileName];
-    }
-  };
-
-  const recordDebugDecision = (decision: RoutingDecision) => {
-    debugHistory = [...debugHistory, decision].slice(-MAX_DEBUG_HISTORY);
-  };
-
-  const getThinkingOverride = (profileName: string, tier: RouterTier) => {
-    return thinkingByProfile[profileName]?.[tier];
-  };
-
-  const persistState = () => {
-    const state = buildPersistedState(
-      routerEnabled,
-      selectedProfile,
-      pinnedTierByProfile,
-      thinkingByProfile,
-      debugEnabled,
-      widgetEnabled,
-      debugHistory,
-      lastDecision,
-      lastNonRouterModel,
-      accumulatedCost,
-      sessionLock,
-    );
-    const snapshot = JSON.stringify({
-      ...state,
-      timestamp: 0,
-      lastDecision: state.lastDecision
-        ? { ...state.lastDecision, timestamp: 0 }
-        : undefined,
-      debugHistory: state.debugHistory?.map((decision) => ({
-        ...decision,
-        timestamp: 0,
-      })),
-    });
-    if (snapshot === lastPersistedSnapshot) {
-      return;
-    }
-    pi.appendEntry('router-state', state);
-    lastPersistedSnapshot = snapshot;
-  };
-
-  const actions = {
-    persistState,
-    updateStatus: (ctx: ExtensionContext) =>
-      updateStatus(
-        ctx,
-        routerEnabled,
-        selectedProfile,
-        pinnedTierByProfile,
-        thinkingByProfile,
-        lastDecision,
-        lastNonRouterModel,
-        accumulatedCost,
-        widgetEnabled,
-        currentConfig,
-        sessionLock,
-      ),
-    reloadConfig: (
-      ctx?: ExtensionContext,
-      options?: { preserveDebug?: boolean },
-    ) => {
-      const loaded = loadRouterConfig(currentCwd);
-      currentConfig = loaded.config;
-      lastConfigWarnings = loaded.warnings;
-      if (!options?.preserveDebug) {
-        debugEnabled = currentConfig.debug ?? false;
-      }
-      selectedProfile = resolveProfileName(currentConfig, selectedProfile);
-      actions.registerRouterProvider();
-      if (ctx) {
-        actions.updateStatus(ctx);
-      }
-    },
-    ensureValidActiveRouterProfile: async (ctx: ExtensionContext) => {
-      if (ctx.model?.provider !== 'router') {
-        return;
-      }
-      if (currentConfig.profiles[ctx.model.id]) {
-        selectedProfile = ctx.model.id;
-        routerEnabled = true;
-        return;
-      }
-
-      const fallbackProfile = resolveProfileName(
-        currentConfig,
-        selectedProfile,
-      );
-      const routerModel = ctx.modelRegistry.find('router', fallbackProfile);
-      selectedProfile = fallbackProfile;
-      if (!routerModel) {
-        ctx.ui.notify(
-          `Router profile "${ctx.model.id}" is no longer configured.`,
-          'warning',
-        );
-        return;
-      }
-
-      await setModelInternally(routerModel);
-      ctx.ui.notify(
-        `Router profile "${ctx.model.id}" is no longer configured. Switched to router/${fallbackProfile}.`,
-        'warning',
-      );
-    },
-    switchToRouterProfile: async (
-      profileName: string,
-      ctx: ExtensionContext,
-      strict = true,
-    ) => {
-      if (strict && !currentConfig.profiles[profileName]) {
-        ctx.ui.notify(`Unknown router profile: ${profileName}`, 'error');
-        return false;
-      }
-      const resolvedProfile = resolveProfileName(currentConfig, profileName);
-
-      // 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', resolvedProfile);
-      if (!routerModel) {
-        ctx.ui.notify(`Unknown router profile: ${profileName}`, 'error');
-        return false;
-      }
-      if (ctx.model && ctx.model.provider !== 'router') {
-        lastNonRouterModel = `${ctx.model.provider}/${ctx.model.id}`;
-      }
-      const success = await setModelInternally(routerModel);
-      if (!success) {
-        ctx.ui.notify(`Failed to switch to router/${resolvedProfile}`, 'error');
-        return false;
-      }
-      selectedProfile = resolvedProfile;
-      routerEnabled = true;
-      sessionLock = undefined;
-      persistState();
-      actions.updateStatus(ctx);
-      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;
-          },
-          get sessionLock() {
-            return sessionLock;
-          },
-          set sessionLock(v) {
-            sessionLock = v;
-          },
-        },
-        {
-          persistState,
-          recordDebugDecision,
-          getThinkingOverride,
-          updateStatus: actions.updateStatus,
-        },
-      );
-    },
-  };
-
-  actions.reloadConfig();
-
-  const restoreStateFromSession = async (ctx: ExtensionContext) => {
-    lastExtensionContext = ctx;
-    currentModelRegistry = ctx.modelRegistry;
-    currentCwd = ctx.cwd;
-    actions.reloadConfig();
-
-    // Give the registry a moment to synchronize after re-registration
-    await new Promise((resolve) => setTimeout(resolve, 50));
-
-    routerEnabled = ctx.model?.provider === 'router';
-    selectedProfile = resolveProfileName(
-      currentConfig,
-      ctx.model?.provider === 'router' ? ctx.model.id : selectedProfile,
-    );
-    pinnedTierByProfile = {};
-    thinkingByProfile = {};
-    widgetEnabled = false;
-    debugHistory = [];
-    accumulatedCost = 0;
-    lastNonRouterModel =
-      ctx.model && ctx.model.provider !== 'router'
-        ? `${ctx.model.provider}/${ctx.model.id}`
-        : lastNonRouterModel;
-    lastDecision = undefined;
-    sessionLock = undefined;
-
-    const entries = ctx.sessionManager.getBranch() as CustomSessionEntry[];
-    const savedState = entries
-      .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,
-      );
-      routerEnabled = savedState.enabled;
-      pinnedTierByProfile = savedState.pinByProfile
-        ? { ...savedState.pinByProfile }
-        : {};
-      thinkingByProfile = savedState.thinkingByProfile
-        ? { ...savedState.thinkingByProfile }
-        : {};
-      if (savedState.pinTier) {
-        pinnedTierByProfile[selectedProfile] = savedState.pinTier;
-      }
-      debugEnabled = savedState.debugEnabled ?? debugEnabled;
-      widgetEnabled = savedState.widgetEnabled ?? widgetEnabled;
-      debugHistory = savedState.debugHistory
-        ? [...savedState.debugHistory].slice(-MAX_DEBUG_HISTORY)
-        : [];
-      lastNonRouterModel = savedState.lastNonRouterModel ?? lastNonRouterModel;
-      accumulatedCost = savedState.accumulatedCost ?? 0;
-      if (
-        savedState.sessionLock &&
-        savedState.sessionLock.profile === selectedProfile
-      ) {
-        sessionLock = { ...savedState.sessionLock };
-      }
-    }
-
-    await actions.ensureValidActiveRouterProfile(ctx);
-
-    if (routerEnabled) {
-      const routerModel = ctx.modelRegistry.find('router', selectedProfile);
-      if (routerModel) {
-        const success = await setModelInternally(routerModel);
-        if (!success) {
-          ctx.ui.notify(
-            `Failed to restore router/${selectedProfile} after relaunch.`,
-            'warning',
-          );
-          routerEnabled = false;
-        }
-      } else {
-        ctx.ui.notify(
-          `Unable to restore router/${selectedProfile}; model is unavailable.`,
-          'warning',
-        );
-        routerEnabled = false;
-        ctx.ui.setHiddenThinkingLabel?.();
-      }
-    } else {
-      ctx.ui.setHiddenThinkingLabel?.();
-    }
-
-    persistState();
-    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;
-      },
-    },
-    actions,
-  );
-
-  pi.on('session_start', async (_event, ctx) => {
-    isInitialized = true;
-    await restoreStateFromSession(ctx);
-    if (debugEnabled) {
-      ctx.ui.notify(
-        `Router initialized with profiles: ${profileNames(currentConfig).join(', ')}`,
-        'info',
-      );
-    }
-  });
-
-  pi.on('model_select', async (event, ctx) => {
-    if (!isInitialized || isInternalModelSwitch) return;
-    if (event.model.provider === 'router') {
-      const profileName = resolveProfileName(currentConfig, event.model.id);
-
-      // If the selected model has stale capacities (e.g. from the initial registration),
-      // re-apply the model from the registry to force a TUI refresh.
-      const registryModel = ctx.modelRegistry.find('router', profileName);
-      if (
-        registryModel &&
-        (registryModel.contextWindow !== event.model.contextWindow ||
-          registryModel.maxTokens !== event.model.maxTokens)
-      ) {
-        await setModelInternally(registryModel);
-      }
-
-      routerEnabled = true;
-      if (selectedProfile !== profileName) {
-        sessionLock = undefined;
-      }
-      selectedProfile = profileName;
-    } else {
-      routerEnabled = false;
-      lastNonRouterModel = `${event.model.provider}/${event.model.id}`;
-      ctx.ui.setHiddenThinkingLabel?.();
-    }
-    persistState();
-    actions.updateStatus(ctx);
-  });
-
-  pi.on('turn_end', async (_event, ctx) => {
-    if (routerEnabled && ctx.model?.provider !== 'router') {
-      const routerModel = ctx.modelRegistry.find('router', selectedProfile);
-      if (routerModel) {
-        await setModelInternally(routerModel);
-      }
-    }
-    persistState();
-    actions.updateStatus(ctx);
-  });
-};
-
-export default routerExtension;