pi-provider-kimi-code 0.6.7 → 0.6.9

package/src/models.ts

@@ -1,8 +1,14 @@
 // Model identity: discovery against the server's /v1/models endpoint, plus the
 // extras-merging helpers used by both registration and the OAuth modifyModels hook.
 
-import type { Api, Model, OAuthCredentials } from "@earendil-works/pi-ai";
-import type { KimiInputModality, KimiResolvedModelConfig, ModelConfig } from "./config.ts";
+import type { Api, Model, OAuthCredentials, ThinkingLevelMap } from "@earendil-works/pi-ai";
+import type {
+  KimiInputModality,
+  KimiResolvedModelConfig,
+  ModelConfig,
+  ModelReasoningMap,
+} from "./config.ts";
+import { DEFAULT_KIMI_CODE_CONFIG } from "./config.ts";
 
 import { type KimiWireProtocol, getBaseUrl } from "./constants.ts";
 import { getKimiProviderHeaders } from "./device.ts";
@@ -31,6 +37,10 @@ const DEFAULT_DISCOVERY_TIMEOUT_MS = 2500;
 
 export interface DiscoverKimiModelMetadataOptions {
   timeoutMs?: number;
+  // 401 recovery for runtime callers (extension startup, settings menu).
+  // oauth.ts's login/refresh flows must NOT pass this: they call discovery
+  // while holding the OAuth lock, and refreshKimiAuthToken would re-enter it.
+  refreshAccessToken?: (currentToken: string) => Promise<string | null>;
 }
 
 function mergeInputModalities(
@@ -50,62 +60,63 @@ function mergeInputModalities(
   return (["text", "image", "video"] as const).filter((modality) => next.has(modality));
 }
 
-// Pricing per million tokens in USD (CNY converted at ~7.25).
-// Source: https://platform.kimi.com/docs/pricing/chat-k27-code
-const COST_STANDARD = { input: 0.897, output: 3.724, cacheRead: 0.179, cacheWrite: 0.897 };
-const COST_HIGH_SPEED = { input: 1.793, output: 7.448, cacheRead: 0.359, cacheWrite: 1.793 };
+// Pricing per million tokens in USD.
+// Sources: https://www.kimi.com/resources/kimi-k2-7-code-pricing
+//          https://www.kimi.com/resources/kimi-k3-pricing
+const COST_STANDARD = { input: 0.95, output: 4, cacheRead: 0.19, cacheWrite: 0.95 };
+const COST_HIGH_SPEED = { input: 1.9, output: 8, cacheRead: 0.38, cacheWrite: 1.9 };
+const COST_K3 = { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3 };
 
 export const KIMI_CODING_MODEL_ID = "kimi-for-coding";
 export const KIMI_CODING_HIGHSPEED_MODEL_ID = "kimi-for-coding-highspeed";
 export const KIMI_K3_MODEL_ID = "k3";
 export const KIMI_MODEL_CATALOG_VERSION = 1;
 
-const KIMI_K3_MODERATO_CONTEXT_WINDOW = 262144;
-const KIMI_MEMBERSHIP_RANK: Readonly<Record<string, number>> = {
-  LEVEL_FREE: 0,
-  LEVEL_BASIC: 1,
-  LEVEL_STANDARD: 2,
-  LEVEL_INTERMEDIATE: 3,
-  LEVEL_ADVANCED: 4,
-  LEVEL_PREMIUM: 5,
-};
-const KIMI_MODERATO_RANK = KIMI_MEMBERSHIP_RANK.LEVEL_STANDARD;
-const KIMI_ALLEGRETTO_RANK = KIMI_MEMBERSHIP_RANK.LEVEL_INTERMEDIATE;
-
-export function isKimiModelAvailableForMembership(
+function resolveModelCost(
   modelId: string,
-  membershipLevel: string | null | undefined,
-): boolean | undefined {
-  if (!membershipLevel) return undefined;
-  const rank = KIMI_MEMBERSHIP_RANK[membershipLevel];
-  if (rank === undefined) return undefined;
-  if (modelId === KIMI_K3_MODEL_ID) return rank >= KIMI_MODERATO_RANK;
-  if (modelId === KIMI_CODING_HIGHSPEED_MODEL_ID) return rank >= KIMI_ALLEGRETTO_RANK;
-  return true;
-}
-
-export function applyKimiMembershipLimitsToModel(
-  model: Model<Api>,
-  membershipLevel: string | null | undefined,
-): Model<Api> {
-  const rank = membershipLevel ? KIMI_MEMBERSHIP_RANK[membershipLevel] : undefined;
-  if (model.id !== KIMI_K3_MODEL_ID || rank === undefined || rank >= KIMI_ALLEGRETTO_RANK) {
-    return model;
-  }
-  return {
-    ...model,
-    contextWindow: Math.min(model.contextWindow, KIMI_K3_MODERATO_CONTEXT_WINDOW),
-  };
-}
-
-function resolveModelCost(modelDisplay: string | undefined): {
+  modelDisplay?: string,
+): {
   input: number;
   output: number;
   cacheRead: number;
   cacheWrite: number;
 } {
-  if (modelDisplay && /high\s*speed/i.test(modelDisplay)) return COST_HIGH_SPEED;
-  return COST_STANDARD;
+  if (modelId === KIMI_K3_MODEL_ID) return COST_K3;
+  if (modelId === KIMI_CODING_MODEL_ID) return COST_STANDARD;
+  if (modelId === KIMI_CODING_HIGHSPEED_MODEL_ID || /high\s*speed/i.test(modelDisplay ?? "")) {
+    return COST_HIGH_SPEED;
+  }
+  return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
+}
+
+// Policy: a model's output cap tracks its context window unless the user
+// explicitly configured model.maxTokens. pi-ai clamps the request cap to
+// contextWindow - used - 4096 (clampMaxTokensToContext in
+// @earendil-works/pi-ai), so setting maxTokens to the window yields the official
+// kimi-code behavior (computeCompletionBudgetCap: maxCtx, clamped to
+// maxCtx - usedContext) instead of a fixed small cap that truncates long
+// max-effort reasoning (reasoning_content counts toward the cap). A deliberate
+// per-request cap still wins (options.maxTokens ?? model.maxTokens). Reconcile
+// wherever contextWindow is authored so the two stay in lockstep.
+//
+// "Explicitly configured" means any value other than the built-in default:
+// ensureKimiCodeConfig materializes the full default config (including
+// maxTokens) into the home config file, so config-source tracking cannot
+// distinguish a deliberate 32000 from a materialized default. A user who
+// explicitly sets the default value gets the window cap — indistinguishable
+// from the materialized-default case, and the intended new behavior anyway.
+const DEFAULT_MODEL_MAX_TOKENS = DEFAULT_KIMI_CODE_CONFIG.model.maxTokens;
+
+// True while the cap is "unset": still the built-in default, or already
+// tracking the context window (set by an earlier withWindowCappedMaxTokens
+// call, before discovery grew the window).
+function isWindowTrackedMaxTokens(model: Model<Api>): boolean {
+  return model.maxTokens === DEFAULT_MODEL_MAX_TOKENS || model.maxTokens === model.contextWindow;
+}
+
+function withWindowCappedMaxTokens<M extends Model<Api>>(model: M): M {
+  if (!isWindowTrackedMaxTokens(model)) return model;
+  return { ...model, maxTokens: model.contextWindow };
 }
 
 export function buildKimiModelFromConfig(
@@ -118,16 +129,18 @@ export function buildKimiModelFromConfig(
       ? "Kimi K3"
       : isHighSpeed
         ? "Kimi for Coding High Speed"
-        : "Kimi for Coding";
-  return {
+        : modelId === KIMI_CODING_MODEL_ID
+          ? "Kimi for Coding"
+          : modelId;
+  return withWindowCappedMaxTokens({
     id: modelId,
     name,
     reasoning: config.reasoning,
     input: [...config.input] as unknown as ("text" | "image" | "video")[],
-    cost: { ...(isHighSpeed ? COST_HIGH_SPEED : COST_STANDARD) },
+    cost: { ...resolveModelCost(modelId) },
     contextWindow: config.contextWindow,
     maxTokens: config.maxTokens,
-  } as Model<Api>;
+  } as Model<Api>);
 }
 
 export function resolveKimiModelConfig(
@@ -226,38 +239,58 @@ function getModelsUrl(protocol?: KimiWireProtocol): string {
   return buildModelsUrl(getBaseUrl(protocol));
 }
 
+export function isOfficialKimiModelsUrl(url: string): boolean {
+  try {
+    const parsed = new URL(url);
+    return parsed.origin === "https://api.kimi.com" && parsed.pathname === "/coding/v1/models";
+  } catch {
+    return false;
+  }
+}
+
 export async function discoverKimiModelMetadata(
   accessToken: string,
   protocol?: KimiWireProtocol,
   options: DiscoverKimiModelMetadataOptions = {},
 ): Promise<KimiOAuthExtras> {
   if (!accessToken) return {};
+  const modelsUrl = getModelsUrl(protocol);
+  if (!isOfficialKimiModelsUrl(modelsUrl)) return {};
   const timeoutMs = options.timeoutMs ?? DEFAULT_DISCOVERY_TIMEOUT_MS;
   const controller = new AbortController();
   const timeout =
     timeoutMs > 0 ? setTimeout(() => controller.abort(), timeoutMs).unref() : undefined;
-  try {
-    const response = await fetch(getModelsUrl(protocol), {
-      signal: controller.signal,
+  const requestModels = (token: string, signal: AbortSignal) =>
+    fetch(modelsUrl, {
+      signal,
       headers: {
         ...getKimiProviderHeaders(),
-        Authorization: `Bearer ${accessToken}`,
+        Authorization: `Bearer ${token}`,
         Accept: "application/json",
       },
     });
+  try {
+    let response = await requestModels(accessToken, controller.signal);
+    if (response.status === 401 && options.refreshAccessToken) {
+      const refreshed = await options.refreshAccessToken(accessToken);
+      if (refreshed && refreshed !== accessToken) {
+        const retryController = new AbortController();
+        const retryTimeout =
+          timeoutMs > 0 ? setTimeout(() => retryController.abort(), timeoutMs).unref() : undefined;
+        try {
+          response = await requestModels(refreshed, retryController.signal);
+        } finally {
+          if (retryTimeout) clearTimeout(retryTimeout);
+        }
+      }
+    }
     if (!response.ok) return {};
     const json = (await response.json()) as { data?: unknown };
     const list = Array.isArray(json.data) ? (json.data as KimiServerModel[]) : [];
-    const supportedIds = new Set([
-      KIMI_CODING_MODEL_ID,
-      KIMI_CODING_HIGHSPEED_MODEL_ID,
-      KIMI_K3_MODEL_ID,
-    ]);
     const modelCatalog: Record<string, KimiModelMetadata> = {};
     for (const model of list) {
-      if (typeof model.id !== "string" || !supportedIds.has(model.id)) continue;
       const metadata = parseKimiModelMetadata(model);
-      if (metadata) modelCatalog[model.id] = metadata;
+      if (metadata) modelCatalog[metadata.wireModelId!] = metadata;
     }
     const preferred = modelCatalog[KIMI_CODING_MODEL_ID];
     if (!preferred) {
@@ -279,10 +312,53 @@ export function getKimiModelMetadata(extras: KimiOAuthExtras, modelId: string):
   return modelId === KIMI_CODING_MODEL_ID ? extras : {};
 }
 
+const KIMI_DISCOVERY_KEYS = [
+  "wireModelId",
+  "modelDisplay",
+  "contextLength",
+  "supportsReasoning",
+  "supportsImageIn",
+  "supportsVideoIn",
+  "supportsThinkingType",
+  "protocol",
+  "supportEfforts",
+  "defaultEffort",
+  "modelCatalog",
+] as const;
+
+// Credentials always carry access/refresh/expires, so discovery presence must
+// be checked on the metadata fields themselves. Empty means either a legacy
+// pre-discovery credential or a failed discovery during login/refresh.
+export function hasKimiModelMetadata(extras: KimiOAuthExtras): boolean {
+  return KIMI_DISCOVERY_KEYS.some((key) => extras[key] !== undefined);
+}
+
+export function buildKimiThinkingLevelMap(
+  reasoningMap: ModelReasoningMap,
+  extras: Pick<KimiModelMetadata, "supportEfforts" | "supportsThinkingType">,
+): ThinkingLevelMap | undefined {
+  if (!extras.supportEfforts?.length) return undefined;
+  const map = {
+    off: extras.supportsThinkingType === "only" ? null : "off",
+  } as ThinkingLevelMap & Record<string, string | null>;
+  for (const level of ["minimal", "low", "medium", "high", "xhigh", "max"] as const) {
+    const entry = reasoningMap[level];
+    map[level] =
+      entry?.enabled && entry.effort && extras.supportEfforts.includes(entry.effort)
+        ? entry.effort
+        : null;
+  }
+  return map;
+}
+
 export function applyKimiOAuthExtrasToModel(
   model: Model<Api>,
   extras: KimiModelMetadata,
+  reasoningMap?: ModelReasoningMap,
 ): Model<Api> {
+  // Capture before extras grow the window: a window-tracked cap must follow
+  // the new window, an explicit cap must survive unchanged.
+  const windowTracked = isWindowTrackedMaxTokens(model);
   const next: Model<Api> & {
     wireModelId?: string;
     supportsThinkingType?: "only" | "no" | "both";
@@ -295,7 +371,7 @@ export function applyKimiOAuthExtrasToModel(
       model.id === KIMI_K3_MODEL_ID && /^k3$/i.test(extras.modelDisplay)
         ? "Kimi K3"
         : extras.modelDisplay;
-    next.cost = resolveModelCost(extras.modelDisplay);
+    next.cost = resolveModelCost(model.id, extras.modelDisplay);
   }
   if (typeof extras.contextLength === "number" && extras.contextLength > 0) {
     next.contextWindow = extras.contextLength;
@@ -316,6 +392,15 @@ export function applyKimiOAuthExtrasToModel(
   }
   if (extras.protocol) next.wireProtocol = extras.protocol;
   if (extras.supportEfforts) next.supportEfforts = [...extras.supportEfforts];
+  else delete next.supportEfforts;
   if (extras.defaultEffort) next.defaultEffort = extras.defaultEffort;
+  else delete next.defaultEffort;
+  if (reasoningMap) {
+    const thinkingLevelMap = buildKimiThinkingLevelMap(reasoningMap, extras);
+    if (thinkingLevelMap) next.thinkingLevelMap = thinkingLevelMap;
+    else delete next.thinkingLevelMap;
+  }
+  // Reconcile the output cap with the (possibly discovery-updated) window.
+  if (windowTracked) next.maxTokens = next.contextWindow;
   return next;
 }