pi-provider-kimi-code 0.6.3 → 0.6.6

package/src/payload.ts

@@ -6,7 +6,7 @@ import type { CacheRetention, ThinkingLevel } from "@earendil-works/pi-ai";
 import type { KimiResolvedModelConfig, ModelReasoningEntry } from "./config.ts";
 
 import { getBaseUrl } from "./constants.ts";
-import { getCommonHeaders } from "./device.ts";
+import { getKimiProviderHeaders } from "./device.ts";
 import { optimizeToolSchemas } from "./schema-dedup.ts";
 
 // =============================================================================
@@ -116,7 +116,7 @@ export async function uploadKimiFile(
   try {
     const response = await fetch(uploadUrl, {
       method: "POST",
-      headers: { Authorization: `Bearer ${apiKey}`, ...getCommonHeaders() },
+      headers: { ...getKimiProviderHeaders(), Authorization: `Bearer ${apiKey}` },
       body: formData,
     });
     if (!response.ok) throw new Error(`${response.status} ${await response.text()}`);
@@ -425,7 +425,19 @@ export async function applyKimiPayloadMutations(
       : { include_usage: true };
   }
 
-  // 5. Normalize deprecated max_tokens (OpenAI path only — Anthropic
+  // 5. Spread extra_body into the top-level payload before normalization and
+  //    config caps. Top-level fields retain precedence over extra_body.
+  if (isRecord(payload.extra_body)) {
+    const extraBody = payload.extra_body as JsonRecord;
+    delete payload.extra_body;
+    for (const [key, value] of Object.entries(extraBody)) {
+      if (payload[key] === undefined) {
+        payload[key] = value;
+      }
+    }
+  }
+
+  // 6. Normalize deprecated max_tokens (OpenAI path only — Anthropic
   //    /messages uses max_tokens natively).
   if (ctx.api === "openai-completions") {
     if (payload.max_completion_tokens === undefined && typeof payload.max_tokens === "number") {
@@ -438,44 +450,67 @@ export async function applyKimiPayloadMutations(
   if (generation.temperature !== undefined) payload.temperature = generation.temperature;
   if (generation.topP !== undefined) payload.top_p = generation.topP;
   if (generation.maxCompletionTokens !== undefined) {
-    payload.max_completion_tokens = generation.maxCompletionTokens;
+    const maxTokensKey = ctx.api === "anthropic-messages" ? "max_tokens" : "max_completion_tokens";
+    const currentMaxTokens = payload[maxTokensKey];
+    payload[maxTokensKey] =
+      typeof currentMaxTokens === "number"
+        ? Math.min(currentMaxTokens, generation.maxCompletionTokens)
+        : generation.maxCompletionTokens;
   }
 
-  // 5. Spread extra_body into top-level payload before reasoning mapping,
-  //    matching upstream (Kimi expects thinking etc. at the root).
-  if (isRecord(payload.extra_body)) {
-    const extraBody = payload.extra_body as JsonRecord;
-    delete payload.extra_body;
-    for (const [key, value] of Object.entries(extraBody)) {
-      if (payload[key] === undefined) {
-        payload[key] = value;
-      }
-    }
+  // 7. Reasoning effort mapping. Kimi now accepts effort only inside the
+  //    thinking object, and only for values advertised by the model catalog.
+  delete payload.reasoning_effort;
+  if (ctx.modelConfig.supportsThinkingType === "no") {
+    delete payload.thinking;
   }
-
-  // 6. Reasoning effort mapping. Merges into any existing top-level thinking
-  //    field (which may have come from extra_body above).
   const resolvedReasoning = resolveThinkingLevel(ctx);
   if (resolvedReasoning) {
     const mapped = resolveReasoningForLevel(resolvedReasoning, ctx.modelConfig);
     if (mapped) {
-      if (mapped.effort !== null) {
-        payload.reasoning_effort = mapped.effort;
-      } else {
-        delete payload.reasoning_effort;
-      }
       const oldThinking = isRecord(payload.thinking) ? payload.thinking : {};
-      payload.thinking = {
+      const thinking: JsonRecord = {
         ...oldThinking,
         type: mapped.enabled ? "enabled" : "disabled",
       };
+      delete thinking.effort;
+      if (!mapped.enabled) delete thinking.keep;
+      const effort = ctx.reasoning
+        ? mapped.effort
+        : (ctx.modelConfig.defaultEffort ?? mapped.effort);
+      if (mapped.enabled && effort !== null && ctx.modelConfig.supportEfforts?.includes(effort)) {
+        thinking.effort = effort;
+      }
       if (mapped.enabled && ctx.modelConfig.thinkingKeep) {
-        (payload.thinking as JsonRecord).keep = ctx.modelConfig.thinkingKeep;
+        thinking.keep = ctx.modelConfig.thinkingKeep;
+      }
+      payload.thinking = thinking;
+    }
+  }
+
+  // Preserved-thinking Kimi endpoints require every replayed assistant turn to
+  // carry reasoning_content, including turns whose reasoning delta was empty.
+  // pi-ai drops empty thinking blocks while building Chat Completions history,
+  // so restore the explicit empty field after the final thinking mode is known.
+  if (
+    ctx.api === "openai-completions" &&
+    isRecord(payload.thinking) &&
+    payload.thinking.type !== "disabled" &&
+    payload.thinking.keep === "all" &&
+    Array.isArray(payload.messages)
+  ) {
+    for (const message of payload.messages) {
+      if (
+        isRecord(message) &&
+        message.role === "assistant" &&
+        message.reasoning_content === undefined
+      ) {
+        message.reasoning_content = "";
       }
     }
   }
 
-  // 7. K2.7 Code API constraints: the server rejects non-default values for
+  // 8. K2.7 Code API constraints: the server rejects non-default values for
   //    temperature (must be 1.0) and top_p (must be 0.95), and tool_choice
   //    "required" / function-specific when thinking is enabled (always-on).
   if (payload.temperature !== undefined && payload.temperature !== 1) {

package/src/stream.ts

@@ -24,7 +24,7 @@ import {
 } from "./config.ts";
 
 import { ENV_KIMI_CODE_PROTOCOL, getApiProtocol, getBaseUrl } from "./constants.ts";
-import { getCommonHeaders } from "./device.ts";
+import { getKimiProviderHeaders } from "./device.ts";
 import { isKimiAuthErrorMessage, refreshKimiAuthToken } from "./oauth.ts";
 import {
   type Uploader,
@@ -163,7 +163,7 @@ export function resolveKimiApiKey(apiKey: string | undefined): string {
 }
 
 export function mergeKimiRequestHeaders(headers?: Record<string, string>): Record<string, string> {
-  return { ...getCommonHeaders(), ...headers };
+  return { ...getKimiProviderHeaders(), ...headers };
 }
 
 export function streamSimpleKimi(
@@ -192,8 +192,27 @@ export function streamSimpleKimi(
     protocol: ENV_KIMI_CODE_PROTOCOL,
     uploads: DEFAULT_KIMI_CODE_CONFIG.uploads,
   };
-  const modelConfig = streamConfig.model;
-  const apiProtocol = getApiProtocol(streamConfig.protocol);
+  const discoveredModel = model as Model<Api> & {
+    supportsThinkingType?: "only" | "no" | "both";
+    wireProtocol?: KimiCodeConfig["protocol"];
+    supportEfforts?: string[];
+    defaultEffort?: string;
+  };
+  const modelConfig: KimiResolvedModelConfig = {
+    ...streamConfig.model,
+    contextWindow: model.contextWindow,
+    maxTokens: model.maxTokens,
+    input: [...model.input],
+    reasoning: model.reasoning,
+    supportsThinkingType:
+      discoveredModel.supportsThinkingType ?? (model.reasoning ? undefined : "no"),
+    supportEfforts: discoveredModel.supportEfforts
+      ? [...discoveredModel.supportEfforts]
+      : undefined,
+    defaultEffort: discoveredModel.defaultEffort,
+  };
+  const wireProtocol = discoveredModel.wireProtocol ?? streamConfig.protocol;
+  const apiProtocol = getApiProtocol(wireProtocol);
   const originalOnPayload = options?.onPayload;
   // The pi-side model id ("kimi-for-coding") is what users select via /model
   // and what gets persisted into sessions. The wire model id discovered at
@@ -257,10 +276,10 @@ export function streamSimpleKimi(
       const runtimeModel = {
         ...model,
         api: apiProtocol,
-        baseUrl: getBaseUrl(streamConfig.protocol),
+        baseUrl: getBaseUrl(wireProtocol),
       } as Model<Api>;
       const upstream =
-        streamConfig.protocol === "openai"
+        wireProtocol === "openai"
           ? streamSimpleOpenAICompletions(
               runtimeModel as Model<"openai-completions">,
               context,

package/src/tools/common.ts

@@ -2,8 +2,8 @@ import { AuthStorage, type AgentToolResult } from "@earendil-works/pi-coding-age
 import { truncateToWidth } from "@earendil-works/pi-tui";
 
 import { PROVIDER_ID, getBaseUrl } from "../constants.ts";
-import { getCommonHeaders } from "../device.ts";
-import { refreshKimiAuthToken } from "../oauth.ts";
+import { getKimiProviderHeaders } from "../device.ts";
+import { KIMI_LOGIN_REQUIRED_MESSAGE, refreshKimiAuthToken } from "../oauth.ts";
 
 export const KIMI_TOOL_TIMEOUT_MS = 180_000;
 
@@ -67,7 +67,7 @@ export function getKimiDatasourceUrl(): string {
 
 export function buildHeaders(accessToken: string, toolCallId: string): Record<string, string> {
   return {
-    ...getCommonHeaders(),
+    ...getKimiProviderHeaders(),
     Authorization: `Bearer ${accessToken}`,
     "Content-Type": "application/json",
     "X-Msh-Tool-Call-Id": toolCallId,
@@ -94,7 +94,9 @@ export async function fetchWithAuthRetry(
   if (response.status !== 401) return response;
 
   const refreshed = await deps.refreshAccessToken(accessToken);
-  if (!refreshed || refreshed === accessToken) return response;
+  if (!refreshed || refreshed === accessToken) {
+    throw new Error(KIMI_LOGIN_REQUIRED_MESSAGE);
+  }
   return request(refreshed);
 }
 

package/src/tools/moonshot.ts

@@ -17,20 +17,6 @@ import {
 
 export const moonshotSearchSchema = Type.Object({
   query: Type.String({ description: "The query text to search for." }),
-  limit: Type.Optional(
-    Type.Number({
-      description: "The number of results to return.",
-      default: 5,
-      minimum: 1,
-      maximum: 20,
-    }),
-  ),
-  include_content: Type.Optional(
-    Type.Boolean({
-      description: "Whether to include fetched page content in the search results.",
-      default: false,
-    }),
-  ),
 });
 
 export const moonshotFetchSchema = Type.Object({
@@ -44,7 +30,8 @@ export interface MoonshotSearchResult {
   url: string;
   title: string;
   snippet: string;
-  content?: string;
+  date?: string;
+  siteName?: string;
 }
 
 export interface MoonshotFetchResult {
@@ -52,11 +39,6 @@ export interface MoonshotFetchResult {
   content: string;
 }
 
-function clampLimit(value: number | undefined): number {
-  if (typeof value !== "number" || !Number.isFinite(value)) return 5;
-  return Math.max(1, Math.min(20, Math.trunc(value)));
-}
-
 interface MoonshotSearchResponse {
   search_results?: unknown;
 }
@@ -66,7 +48,7 @@ function stringField(record: Record<string, unknown>, key: string): string {
   return typeof value === "string" ? value : "";
 }
 
-function mapSearchResults(value: unknown, includeContent: boolean): MoonshotSearchResult[] {
+function mapSearchResults(value: unknown): MoonshotSearchResult[] {
   if (!Array.isArray(value)) return [];
   return value
     .filter((item): item is Record<string, unknown> => typeof item === "object" && item !== null)
@@ -76,9 +58,10 @@ function mapSearchResults(value: unknown, includeContent: boolean): MoonshotSear
         title: stringField(item, "title"),
         snippet: stringField(item, "snippet"),
       };
-      if (includeContent) {
-        result.content = stringField(item, "content");
-      }
+      const date = stringField(item, "date");
+      if (date) result.date = date;
+      const siteName = stringField(item, "site_name");
+      if (siteName) result.siteName = siteName;
       return result;
     });
 }
@@ -130,8 +113,10 @@ export function buildMoonshotSearchTool(options: BuildKimiToolOptions = {}) {
   return defineTool({
     name: "moonshot_search",
     label: "Moonshot Search",
-    description: "Search the web through Kimi Coding's server-side Moonshot search service.",
-    promptSnippet: "Search the web with Kimi Coding's Moonshot search service",
+    description:
+      "Search the web through Kimi Coding. Results are summaries; use moonshot_fetch to read relevant pages.",
+    promptSnippet:
+      "Search the web with Kimi Coding, then fetch the few relevant result URLs before answering",
     parameters: moonshotSearchSchema,
 
     async execute(toolCallId, params, signal) {
@@ -142,20 +127,13 @@ export function buildMoonshotSearchTool(options: BuildKimiToolOptions = {}) {
         );
       }
 
-      const limit = clampLimit(params.limit);
-      const includeContent = params.include_content === true;
       const timeout = buildTimeoutSignal(signal);
       try {
         const response = await fetchWithAuthRetry(deps, accessToken, (token) =>
           deps.fetch(`${getKimiBaseV1()}/search`, {
             method: "POST",
             headers: buildHeaders(token, toolCallId),
-            body: JSON.stringify({
-              text_query: params.query,
-              limit,
-              enable_page_crawling: includeContent,
-              timeout_seconds: 30,
-            }),
+            body: JSON.stringify({ text_query: params.query }),
             signal: timeout.signal,
           }),
         );
@@ -168,7 +146,7 @@ export function buildMoonshotSearchTool(options: BuildKimiToolOptions = {}) {
         }
 
         const data = (await response.json()) as MoonshotSearchResponse;
-        const results = mapSearchResults(data.search_results, includeContent);
+        const results = mapSearchResults(data.search_results);
         return {
           content: [{ type: "text", text: searchResultsText(results) }],
           details: results,

package/src/usage.ts

@@ -1,7 +1,7 @@
 import { AuthStorage } from "@earendil-works/pi-coding-agent";
 
 import { PROVIDER_ID, getBaseUrl } from "./constants.ts";
-import { getCommonHeaders } from "./device.ts";
+import { getKimiProviderHeaders } from "./device.ts";
 import { refreshKimiAuthToken } from "./oauth.ts";
 
 const MEMBERSHIP_LEVEL_NAMES: Record<string, string> = {
@@ -36,6 +36,21 @@ export interface UsageFormatOptions {
   timeZone?: string;
 }
 
+export interface KimiUsageSnapshot {
+  summary: string;
+  membershipLevel: string | null;
+}
+
+interface BoosterWalletInfo {
+  balanceCents: bigint;
+  monthlyChargeLimitEnabled: boolean;
+  monthlyChargeLimitCents: bigint;
+  monthlyUsedCents: bigint;
+  currency: string;
+}
+
+const FIXED_POINT_CENTS = 1_000_000n;
+
 export function getKimiUsageToken(): string | null {
   const credential = AuthStorage.create().get(PROVIDER_ID);
   if (credential?.type === "oauth" && credential.access) return credential.access;
@@ -62,6 +77,82 @@ function getResetTime(record: Record<string, unknown>): string | undefined {
   return undefined;
 }
 
+function toBigInt(value: unknown): bigint | null {
+  if (typeof value === "bigint") return value;
+  if (typeof value === "string" && /^-?\d+$/.test(value)) return BigInt(value);
+  if (typeof value === "number" && Number.isSafeInteger(value)) return BigInt(value);
+  return null;
+}
+
+function fixedPointToCents(value: bigint): bigint {
+  if (value > 0n && value < FIXED_POINT_CENTS) return 1n;
+  return (value + FIXED_POINT_CENTS / 2n) / FIXED_POINT_CENTS;
+}
+
+function parseMoney(value: unknown): { cents: bigint; currency: string } | null {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) return null;
+  const record = value as Record<string, unknown>;
+  const cents = toBigInt(record.priceInCents);
+  if (cents === null) return null;
+  return {
+    cents,
+    currency: typeof record.currency === "string" ? record.currency : "",
+  };
+}
+
+function parseBoosterWallet(value: unknown): BoosterWalletInfo | null {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) return null;
+  const record = value as Record<string, unknown>;
+  const balance = record.balance;
+  if (typeof balance !== "object" || balance === null || Array.isArray(balance)) return null;
+  const balanceRecord = balance as Record<string, unknown>;
+  if (balanceRecord.type !== "BOOSTER") return null;
+  const amount = toBigInt(balanceRecord.amount);
+  if (amount === null || amount < 0n) return null;
+
+  const amountLeft = toBigInt(balanceRecord.amountLeft);
+  const monthlyLimit = parseMoney(record.monthlyChargeLimit);
+  const monthlyUsed = parseMoney(record.monthlyUsed);
+  return {
+    balanceCents: amountLeft === null ? 0n : fixedPointToCents(amountLeft),
+    monthlyChargeLimitEnabled: record.monthlyChargeLimitEnabled === true,
+    monthlyChargeLimitCents: monthlyLimit?.cents ?? 0n,
+    monthlyUsedCents: monthlyUsed?.cents ?? 0n,
+    currency: monthlyLimit?.currency || monthlyUsed?.currency || "USD",
+  };
+}
+
+function formatCurrency(cents: bigint, currency: string): string {
+  const symbol =
+    currency.toUpperCase() === "USD" ? "$" : currency.toUpperCase() === "CNY" ? "¥" : "";
+  const sign = cents < 0n ? "-" : "";
+  const absolute = cents < 0n ? -cents : cents;
+  const amount = `${sign}${absolute / 100n}.${String(absolute % 100n).padStart(2, "0")}`;
+  return symbol ? `${symbol}${amount}` : `${amount} ${currency}`.trim();
+}
+
+function formatExtraUsage(info: BoosterWalletInfo): string[] {
+  const hasMonthlyLimit = info.monthlyChargeLimitEnabled && info.monthlyChargeLimitCents > 0n;
+  const lines = ["Extra Usage"];
+  if (hasMonthlyLimit) {
+    const used =
+      info.monthlyUsedCents < info.monthlyChargeLimitCents
+        ? info.monthlyUsedCents
+        : info.monthlyChargeLimitCents;
+    const percent = Number(
+      (used * 100n + info.monthlyChargeLimitCents / 2n) / info.monthlyChargeLimitCents,
+    );
+    const barUnits = Number((used * 400n) / info.monthlyChargeLimitCents);
+    lines.push(`${quotaBar(barUnits, 400)} ${percent}% used`);
+  }
+  lines.push(`Used this month: ${formatCurrency(info.monthlyUsedCents, info.currency)}`);
+  lines.push(
+    `Monthly limit: ${hasMonthlyLimit ? formatCurrency(info.monthlyChargeLimitCents, info.currency) : "Unlimited"}`,
+  );
+  lines.push(`Balance: ${formatCurrency(info.balanceCents, info.currency)}`);
+  return lines;
+}
+
 export function parseUsageRow(value: unknown, fallbackLabel: string): UsageRow | null {
   if (typeof value !== "object" || value === null || Array.isArray(value)) return null;
   const record = value as Record<string, unknown>;
@@ -115,6 +206,12 @@ export function parseUsageSummary(payload: unknown, options: UsageFormatOptions
     }
   }
 
+  const extraUsage = parseBoosterWallet(record.boosterWallet);
+  if (extraUsage) {
+    if (lines.length > 0) lines.push("");
+    lines.push(...formatExtraUsage(extraUsage));
+  }
+
   while (lines.at(-1) === "") lines.pop();
   return lines.length === 0 ? "Usage: no usage data" : lines.join("\n");
 }
@@ -140,7 +237,7 @@ function formatWindowLabel(value: unknown, fallbackLabel: string): string {
   return `Current ${minutes}m window`;
 }
 
-export function parseMembership(record: Record<string, unknown>): string | null {
+export function parseMembershipLevel(record: Record<string, unknown>): string | null {
   const user = record.user;
   if (typeof user !== "object" || user === null || Array.isArray(user)) return null;
   const membership = (user as Record<string, unknown>).membership;
@@ -148,7 +245,12 @@ export function parseMembership(record: Record<string, unknown>): string | null
     return null;
   }
   const level = (membership as Record<string, unknown>).level;
-  if (typeof level !== "string" || !level) return null;
+  return typeof level === "string" && level ? level : null;
+}
+
+export function parseMembership(record: Record<string, unknown>): string | null {
+  const level = parseMembershipLevel(record);
+  if (!level) return null;
   const name = MEMBERSHIP_LEVEL_NAMES[level];
   return name ? `Membership: ${name} (${level})` : `Membership: ${level}`;
 }
@@ -244,31 +346,49 @@ function fetchKimiUsage(token: string, signal: AbortSignal): Promise<Response> {
   return fetch(buildKimiUsageUrl(), {
     method: "GET",
     headers: {
-      ...getCommonHeaders(),
+      ...getKimiProviderHeaders(),
       Authorization: `Bearer ${token}`,
     },
     signal,
   });
 }
 
-export async function fetchKimiUsageSummary(): Promise<string> {
-  const token = getKimiUsageToken();
-  if (!token) return "Usage: missing credentials. Run /login kimi-coding.";
+export async function fetchKimiUsageSnapshot(
+  options: { timeoutMs?: number; token?: string; refreshOnUnauthorized?: boolean } = {},
+): Promise<KimiUsageSnapshot> {
+  const token = options.token ?? getKimiUsageToken();
+  if (!token) {
+    return {
+      summary: "Usage: missing credentials. Run /login kimi-coding.",
+      membershipLevel: null,
+    };
+  }
 
   const controller = new AbortController();
-  const timeout = setTimeout(() => controller.abort(), 15_000);
+  const timeout = setTimeout(() => controller.abort(), options.timeoutMs ?? 15_000);
   try {
     let response = await fetchKimiUsage(token, controller.signal);
-    if (response.status === 401) {
-      const refreshed = await refreshKimiAuthToken(token);
+    if (response.status === 401 && options.refreshOnUnauthorized !== false) {
+      const refreshed = await refreshKimiAuthToken(token, { signal: controller.signal });
       if (refreshed) response = await fetchKimiUsage(refreshed, controller.signal);
     }
-    if (!response.ok) return `Usage: fetch failed (${response.status})`;
-    return parseUsageSummary(await response.json());
+    if (!response.ok) {
+      return { summary: `Usage: fetch failed (${response.status})`, membershipLevel: null };
+    }
+    const payload = (await response.json()) as unknown;
+    const membershipLevel =
+      typeof payload === "object" && payload !== null && !Array.isArray(payload)
+        ? parseMembershipLevel(payload as Record<string, unknown>)
+        : null;
+    return { summary: parseUsageSummary(payload), membershipLevel };
   } catch (error) {
     const message = error instanceof Error ? error.message : String(error);
-    return `Usage: fetch failed (${message})`;
+    return { summary: `Usage: fetch failed (${message})`, membershipLevel: null };
   } finally {
     clearTimeout(timeout);
   }
 }
+
+export async function fetchKimiUsageSummary(): Promise<string> {
+  return (await fetchKimiUsageSnapshot()).summary;
+}