ultimate-pi 0.18.0 → 0.19.0

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

@@ -1,580 +0,0 @@
-import {
-  createAssistantMessageEventStream,
-  streamSimple,
-  type Api,
-  type AssistantMessage,
-  type AssistantMessageEventStream,
-  type Context,
-  type Model,
-  type SimpleStreamOptions,
-  type Message,
-} from '@earendil-works/pi-ai';
-import type {
-  ExtensionAPI,
-  ExtensionContext,
-} from '@earendil-works/pi-coding-agent';
-import type {
-  RouterConfig,
-  RoutingDecision,
-  RouterTier,
-  RouterPinByProfile,
-  RouterThinkingByProfile,
-  SessionLock,
-} from './types.js';
-import { profileNames, parseCanonicalModelRef, ROUTER_TIERS } from './config.js';
-import {
-  phaseForTier,
-  buildRoutingDecision,
-  decideRouting,
-  decideSessionLock,
-  runClassifier,
-  extractTextFromContent,
-  hasImageAttachment,
-  buildSessionLockContext,
-  sessionLockToRoutingDecision,
-  routingDecisionToSessionLock,
-  applyThinkingToDecision,
-} from './routing.js';
-
-export const createErrorMessage = (
-  model: Model<Api>,
-  message: string,
-): AssistantMessage => {
-  return {
-    role: 'assistant',
-    content: [],
-    api: model.api,
-    provider: model.provider,
-    model: model.id,
-    usage: {
-      input: 0,
-      output: 0,
-      cacheRead: 0,
-      cacheWrite: 0,
-      totalTokens: 0,
-      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-    },
-    stopReason: 'error',
-    errorMessage: message,
-    timestamp: Date.now(),
-  };
-};
-
-/**
- * Heuristic token estimator (conservative: 3 characters per token)
- */
-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.
- */
-const truncateContext = (context: Context, limit: number): Context => {
-  const messages = [...context.messages];
-  if (messages.length <= 1) return context;
-
-  const getSystemTokens = () =>
-    context.systemPrompt ? estimateTokens(context.systemPrompt) : 0;
-
-  // Initial estimate
-  const totalTokens =
-    getSystemTokens() +
-    messages.reduce(
-      (sum, m) => sum + estimateTokens(extractTextFromContent(m.content)),
-      0,
-    );
-  if (totalTokens <= limit) return context;
-
-  const latestMessage = messages.pop();
-  if (!latestMessage) return context;
-
-  // Remove oldest until it fits
-  while (messages.length > 0) {
-    const currentTokens =
-      getSystemTokens() +
-      estimateTokens(extractTextFromContent(latestMessage.content)) +
-      messages.reduce(
-        (sum, m) => sum + estimateTokens(extractTextFromContent(m.content)),
-        0,
-      );
-
-    if (currentTokens <= limit) break;
-    messages.shift(); // Remove oldest
-  }
-
-  const finalMessages: Message[] = [];
-  finalMessages.push(...messages);
-  finalMessages.push(latestMessage);
-
-  return { ...context, messages: finalMessages };
-};
-
-const supportsReasoning = (
-  profile: RouterConfig['profiles'][string],
-  modelRegistry: ExtensionContext['modelRegistry'] | undefined,
-): boolean => {
-  if (!modelRegistry) return false;
-
-  for (const tier of ROUTER_TIERS) {
-    try {
-      const { provider, modelId } = parseCanonicalModelRef(profile[tier].model);
-      if (modelRegistry.find(provider, modelId)?.reasoning) {
-        return true;
-      }
-    } catch (_error) {
-      // ignore invalid model refs here; config normalization handles warnings
-    }
-  }
-
-  return false;
-};
-
-export const registerRouterProvider = (
-  pi: ExtensionAPI,
-  state: {
-    lastRegisteredModels: string;
-    readonly currentConfig: RouterConfig;
-    readonly currentModelRegistry:
-      | ExtensionContext['modelRegistry']
-      | undefined;
-    readonly lastExtensionContext: ExtensionContext | undefined;
-    selectedProfile: string;
-    routerEnabled: boolean;
-    lastDecision: RoutingDecision | undefined;
-    readonly thinkingByProfile: RouterThinkingByProfile;
-    readonly pinnedTierByProfile: RouterPinByProfile;
-    accumulatedCost: number;
-    get sessionLock(): SessionLock | undefined;
-    set sessionLock(v: SessionLock | undefined);
-  },
-  actions: {
-    persistState: () => void;
-    recordDebugDecision: (decision: RoutingDecision) => void;
-    getThinkingOverride: (profileName: string, tier: RouterTier) => any;
-    updateStatus: (ctx: ExtensionContext) => void;
-  },
-) => {
-  const profileList = profileNames(state.currentConfig);
-
-  // Map profiles to their capacities
-  const modelDefinitions = profileList.map((name) => {
-    const profile = state.currentConfig.profiles[name];
-    let contextWindow = 1_000_000;
-    let maxTokens = 64_000;
-
-    if (state.currentModelRegistry) {
-      for (const tier of ROUTER_TIERS) {
-        try {
-          const { provider, modelId } = parseCanonicalModelRef(
-            profile[tier].model,
-          );
-          const tierModel = state.currentModelRegistry.find(provider, modelId);
-          if (tierModel) {
-            if (tier === 'high') {
-              contextWindow = tierModel.contextWindow ?? contextWindow;
-              maxTokens = tierModel.maxTokens ?? maxTokens;
-            }
-          }
-        } catch (_error) {
-          // ignore
-        }
-      }
-    }
-
-    return {
-      id: name,
-      name: `Router ${name}`,
-      reasoning: supportsReasoning(profile, state.currentModelRegistry),
-      input: ['text', 'image'] as ('text' | 'image')[],
-      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-      contextWindow,
-      maxTokens,
-    };
-  });
-
-  const modelsKey = modelDefinitions
-    .map((m) => `${m.id}:${m.contextWindow}:${m.maxTokens}:${m.reasoning}`)
-    .join(',');
-  if (state.lastRegisteredModels === modelsKey) return;
-
-  pi.registerProvider('router', {
-    baseUrl: 'router://local',
-    apiKey: 'pi-model-router',
-    api: 'router-local-api',
-    models: modelDefinitions,
-    streamSimple(
-      model: Model<Api>,
-      context: Context,
-      options?: SimpleStreamOptions,
-    ): AssistantMessageEventStream {
-      const stream = createAssistantMessageEventStream();
-
-      (async () => {
-        try {
-          if (!state.currentModelRegistry) {
-            throw new Error(
-              'Router provider not initialized yet. Wait for session_start and retry.',
-            );
-          }
-          const profile = state.currentConfig.profiles[model.id];
-          if (!profile) {
-            throw new Error(`Unknown router profile: ${model.id}`);
-          }
-
-          state.selectedProfile = model.id;
-          state.routerEnabled = true;
-
-          const pinnedTier = state.pinnedTierByProfile[model.id];
-          const thinkingOverrides = state.thinkingByProfile[model.id];
-          const isBudgetExceeded =
-            state.currentConfig.maxSessionBudget !== undefined &&
-            state.accumulatedCost >= state.currentConfig.maxSessionBudget;
-
-          const checkModelSupportsImage = (modelRef: string) => {
-            try {
-              const { provider, modelId } = parseCanonicalModelRef(modelRef);
-              const m = state.currentModelRegistry?.find(provider, modelId);
-              return m?.input?.includes('image') ?? false;
-            } catch {
-              return false;
-            }
-          };
-
-          const pickVisionCapableLock = (
-            lockDecision: RoutingDecision,
-          ): RoutingDecision => {
-            if (!hasImageAttachment(context)) return lockDecision;
-            const candidates = [
-              lockDecision.targetLabel,
-              ...(profile[lockDecision.tier].fallbacks ?? []),
-            ];
-            if (candidates.some(checkModelSupportsImage)) return lockDecision;
-            for (const tier of ROUTER_TIERS) {
-              const refs = [profile[tier].model, ...(profile[tier].fallbacks ?? [])];
-              const visionRef = refs.find(checkModelSupportsImage);
-              if (visionRef) {
-                const { provider, modelId } = parseCanonicalModelRef(visionRef);
-                return {
-                  ...lockDecision,
-                  tier,
-                  targetLabel: visionRef,
-                  targetProvider: provider,
-                  targetModelId: modelId,
-                  reasoning: `${lockDecision.reasoning} | Vision-capable model for image input.`,
-                };
-              }
-            }
-            return lockDecision;
-          };
-
-          let lock = state.sessionLock;
-          if (!lock || lock.profile !== model.id) {
-            const lockContext = buildSessionLockContext(context);
-            let lockDecision = decideSessionLock(
-              lockContext,
-              model.id,
-              profile,
-              pinnedTier,
-              thinkingOverrides,
-              state.currentConfig.phaseBias,
-              state.currentConfig.rules,
-            );
-
-            if (
-              state.currentConfig.classifierModel &&
-              !pinnedTier &&
-              !lockDecision.isRuleMatched
-            ) {
-              const classifierResult = await runClassifier(
-                state.currentConfig.classifierModel,
-                state.currentModelRegistry,
-                lockContext,
-                undefined,
-              );
-              if (classifierResult) {
-                lockDecision = buildRoutingDecision(
-                  model.id,
-                  profile,
-                  classifierResult.tier,
-                  phaseForTier(classifierResult.tier),
-                  `Session lock classifier: ${classifierResult.reasoning}`,
-                  thinkingOverrides,
-                  true,
-                );
-              }
-            }
-
-            lockDecision = pickVisionCapableLock(lockDecision);
-            lock = routingDecisionToSessionLock(lockDecision);
-            state.sessionLock = lock;
-          }
-
-          const lockedBase = sessionLockToRoutingDecision(
-            lock,
-            profile,
-            thinkingOverrides,
-          );
-
-          let thinkingDecision: RoutingDecision = decideRouting(
-            context,
-            model.id,
-            profile,
-            state.lastDecision,
-            pinnedTier,
-            thinkingOverrides,
-            state.currentConfig.phaseBias,
-            state.currentConfig.rules,
-            isBudgetExceeded,
-          );
-
-          if (
-            state.currentConfig.largeContextThreshold &&
-            thinkingDecision.tier !== 'high' &&
-            state.lastExtensionContext
-          ) {
-            try {
-              const usage = await state.lastExtensionContext.getContextUsage();
-              if (
-                usage?.tokens &&
-                usage.tokens > state.currentConfig.largeContextThreshold
-              ) {
-                thinkingDecision = {
-                  ...thinkingDecision,
-                  tier: 'high',
-                  phase: 'planning',
-                  reasoning: `Context usage (${usage.tokens}) exceeds threshold (${state.currentConfig.largeContextThreshold}). Forced high thinking.`,
-                  isContextTriggered: true,
-                };
-              }
-            } catch (_e) {
-              // ignore
-            }
-          }
-
-          if (
-            state.currentConfig.classifierModel &&
-            !pinnedTier &&
-            !thinkingDecision.isContextTriggered &&
-            !thinkingDecision.isRuleMatched
-          ) {
-            const classifierResult = await runClassifier(
-              state.currentConfig.classifierModel,
-              state.currentModelRegistry,
-              context,
-              state.lastDecision?.phase,
-            );
-            if (classifierResult) {
-              thinkingDecision = buildRoutingDecision(
-                model.id,
-                profile,
-                classifierResult.tier,
-                phaseForTier(classifierResult.tier),
-                `Thinking classifier: ${classifierResult.reasoning}`,
-                thinkingOverrides,
-                true,
-              );
-              if (isBudgetExceeded && thinkingDecision.tier === 'high') {
-                thinkingDecision = {
-                  ...thinkingDecision,
-                  tier: 'medium',
-                  phase: 'implementation',
-                  reasoning: `Budget exceeded. Downgraded thinking to medium. (Original: ${thinkingDecision.reasoning})`,
-                  isBudgetForced: true,
-                };
-              }
-            }
-          }
-
-          let decision = applyThinkingToDecision(
-            lockedBase,
-            thinkingDecision,
-            profile,
-            thinkingOverrides,
-          );
-
-          const lastMessage = context.messages[context.messages.length - 1];
-          const previousDecision = state.lastDecision;
-          const isGoogleThinkingToolContinuation =
-            lastMessage?.role === 'toolResult' &&
-            previousDecision?.profile === model.id &&
-            previousDecision.targetProvider === 'google' &&
-            previousDecision.thinking !== 'off' &&
-            decision.targetProvider === 'google' &&
-            decision.thinking !== 'off' &&
-            previousDecision.targetLabel === decision.targetLabel;
-
-          if (isGoogleThinkingToolContinuation) {
-            decision = {
-              ...decision,
-              tier: previousDecision!.tier,
-              phase: previousDecision!.phase,
-              targetProvider: previousDecision!.targetProvider,
-              targetModelId: previousDecision!.targetModelId,
-              targetLabel: previousDecision!.targetLabel,
-              thinking: previousDecision!.thinking,
-              reasoning:
-                `Preserved ${previousDecision!.targetLabel} for a Google tool-result continuation ` +
-                `to avoid thought-signature replay errors. (Original: ${decision.reasoning})`,
-            };
-          }
-
-          state.lastDecision = decision;
-          actions.recordDebugDecision(decision);
-
-          if (state.lastExtensionContext) {
-            actions.updateStatus(state.lastExtensionContext);
-          }
-
-          const lockTier = lock.tier;
-          let modelsToTry = [
-            lock.modelRef,
-            ...(profile[lockTier].fallbacks ?? []),
-          ];
-          const imageAttached = hasImageAttachment(context);
-          if (imageAttached) {
-            modelsToTry = modelsToTry.filter((modelRef) => {
-              try {
-                const { provider, modelId } = parseCanonicalModelRef(modelRef);
-                const m = state.currentModelRegistry?.find(provider, modelId);
-                return m?.input?.includes('image') ?? false;
-              } catch {
-                return false;
-              }
-            });
-            if (modelsToTry.length === 0) {
-              modelsToTry = [decision.targetLabel];
-            }
-          }
-          let lastError: any;
-          let success = false;
-
-          for (let i = 0; i < modelsToTry.length; i++) {
-            const modelRef = modelsToTry[i];
-            const { provider: targetProvider, modelId: targetModelId } =
-              parseCanonicalModelRef(modelRef);
-
-            if (targetProvider === 'router') continue;
-
-            const targetModel = state.currentModelRegistry.find(
-              targetProvider,
-              targetModelId,
-            );
-            if (!targetModel) {
-              lastError = new Error(
-                `Routed model not found: ${targetProvider}/${targetModelId}`,
-              );
-              continue;
-            }
-
-            const auth =
-              await state.currentModelRegistry.getApiKeyAndHeaders(targetModel);
-            if (!auth.ok || !auth.apiKey) {
-              lastError = new Error(
-                auth.ok
-                  ? `No API key for routed model: ${targetProvider}/${targetModelId}`
-                  : `Auth failed for routed model: ${targetProvider}/${targetModelId}: ${auth.error}`,
-              );
-              continue;
-            }
-            const apiKey = auth.apiKey;
-            const headers = auth.headers;
-
-            try {
-              // HONESTY CHECK & AUTO-TRUNCATION
-              // If the picked model has a smaller context than what we reported, truncate now.
-              let effectiveContext = context;
-              const targetLimit = targetModel.contextWindow || 128_000;
-              if (targetLimit < model.contextWindow!) {
-                effectiveContext = truncateContext(context, targetLimit);
-              }
-
-              const thinkingOverride = actions.getThinkingOverride(
-                model.id,
-                decision.tier,
-              );
-              const delegatedReasoning =
-                targetModel.reasoning &&
-                (thinkingOverride ?? decision.thinking) !== 'off'
-                  ? (thinkingOverride ?? decision.thinking)
-                  : undefined;
-
-              if (state.lastExtensionContext) {
-                if (delegatedReasoning) {
-                  state.lastExtensionContext.ui.setHiddenThinkingLabel?.(
-                    `Thinking (${targetProvider}/${targetModelId})...`,
-                  );
-                } else {
-                  state.lastExtensionContext.ui.setHiddenThinkingLabel?.();
-                }
-              }
-
-              const delegatedStream = streamSimple(
-                targetModel,
-                effectiveContext,
-                {
-                  ...options,
-                  apiKey,
-                  headers,
-                  ...(delegatedReasoning
-                    ? { reasoning: delegatedReasoning }
-                    : {}),
-                },
-              );
-
-              let contentReceived = false;
-              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) {
-                  throw new Error(
-                    (event as any).error?.errorMessage ||
-                      'Model failed before sending content.',
-                  );
-                }
-                const isContent =
-                  event.type === 'text_delta' ||
-                  event.type === 'thinking_delta' ||
-                  event.type === 'toolcall_delta' ||
-                  event.type === 'toolcall_end';
-                if (isContent) contentReceived = true;
-                stream.push(event);
-              }
-              success = true;
-              if (i > 0) decision.isFallback = true;
-              break;
-            } catch (err) {
-              lastError = err;
-            }
-          }
-
-          if (!success) {
-            throw (
-              lastError ||
-              new Error('Failed to delegate to any model in the chain.')
-            );
-          }
-
-          stream.end();
-        } catch (error) {
-          stream.push({
-            type: 'error',
-            reason: 'error',
-            error: createErrorMessage(
-              model,
-              error instanceof Error ? error.message : String(error),
-            ),
-          });
-          stream.end();
-        } finally {
-          actions.persistState();
-        }
-      })();
-
-      return stream;
-    },
-  });
-
-  state.lastRegisteredModels = modelsKey;
-};