pi-provider-kimi-code 0.5.1 → 0.6.1

package/src/schema-dedup.ts

@@ -0,0 +1,256 @@
+// Tool schema deduplication: extract repeated sub-schemas into $defs/$ref
+// to stay under Moonshot's ~15 KB per-tool function.parameters limit.
+
+import { createHash } from "node:crypto";
+
+import type { JsonRecord } from "./payload.ts";
+
+const TOOL_SCHEMA_SIZE_THRESHOLD = 14_000;
+const MIN_FRAGMENT_SIZE = 50;
+const MAX_DEDUP_PASSES = 5;
+
+// ---------------------------------------------------------------------------
+// Internals
+// ---------------------------------------------------------------------------
+
+function isRecord(v: unknown): v is JsonRecord {
+  return typeof v === "object" && v !== null && !Array.isArray(v);
+}
+
+function jsonSize(v: unknown): number {
+  return Buffer.byteLength(JSON.stringify(v));
+}
+
+type PathSegment = string | number;
+type FragmentMap = Map<string, PathSegment[][]>;
+
+function collectFragments(
+  node: unknown,
+  path: PathSegment[],
+  map: FragmentMap,
+  minSize: number,
+): void {
+  if (node === null || typeof node !== "object") return;
+  if (Array.isArray(node)) {
+    node.forEach((item, i) => collectFragments(item, [...path, i], map, minSize));
+    return;
+  }
+  const serialized = JSON.stringify(node);
+  if (serialized.length >= minSize) {
+    let paths = map.get(serialized);
+    if (!paths) {
+      paths = [];
+      map.set(serialized, paths);
+    }
+    paths.push(path);
+  }
+  for (const [key, value] of Object.entries(node as JsonRecord)) {
+    collectFragments(value, [...path, key], map, minSize);
+  }
+}
+
+function setAtPath(root: JsonRecord, path: PathSegment[], value: unknown): void {
+  let current: unknown = root;
+  for (let i = 0; i < path.length - 1; i++) {
+    if (!isRecord(current) && !Array.isArray(current)) return;
+    current = (current as Record<string, unknown>)[path[i]];
+  }
+  if (isRecord(current) || Array.isArray(current)) {
+    (current as Record<string, unknown>)[path[path.length - 1]] = value;
+  }
+}
+
+function isDescendant(child: PathSegment[], ancestor: PathSegment[]): boolean {
+  if (child.length <= ancestor.length) return false;
+  for (let i = 0; i < ancestor.length; i++) {
+    if (child[i] !== ancestor[i]) return false;
+  }
+  return true;
+}
+
+function nextDefKey(existing: Set<string>, index: number): string {
+  let key = `d${index}`;
+  while (existing.has(key)) key = `d${++index}`;
+  return key;
+}
+
+// ---------------------------------------------------------------------------
+// Core dedup
+// ---------------------------------------------------------------------------
+
+function deduplicateSchema(schema: JsonRecord): JsonRecord {
+  const originalSize = jsonSize(schema);
+  const result: JsonRecord = JSON.parse(JSON.stringify(schema));
+  const defs: JsonRecord = isRecord(result.$defs) ? (result.$defs as JsonRecord) : {};
+  const existingKeys = new Set(Object.keys(defs));
+  let defIndex = 0;
+  const replaced: PathSegment[][] = [];
+
+  for (let pass = 0; pass < MAX_DEDUP_PASSES; pass++) {
+    const fragments: FragmentMap = new Map();
+    collectFragments(result, [], fragments, MIN_FRAGMENT_SIZE);
+
+    const candidates = [...fragments.entries()]
+      .filter(([, paths]) => paths.length >= 2)
+      .sort((a, b) => b[0].length - a[0].length);
+
+    let progress = false;
+    for (const [serialized, paths] of candidates) {
+      const active = paths.filter((p) => !replaced.some((r) => isDescendant(p, r)));
+      if (active.length < 2) continue;
+
+      const key = nextDefKey(existingKeys, defIndex++);
+      const refObject = JSON.stringify({ $ref: `#/$defs/${key}` });
+      const refSize = Buffer.byteLength(refObject);
+      const fragmentSize = Buffer.byteLength(serialized);
+      const defsEntryOverhead = Buffer.byteLength(JSON.stringify(key)) + 1;
+      const savings =
+        active.length * fragmentSize - (fragmentSize + defsEntryOverhead + active.length * refSize);
+      if (savings <= 0) continue;
+
+      existingKeys.add(key);
+      defs[key] = JSON.parse(serialized);
+
+      for (const p of active) {
+        setAtPath(result, p, { $ref: `#/$defs/${key}` });
+        replaced.push(p);
+      }
+      progress = true;
+    }
+    if (!progress) break;
+  }
+
+  if (Object.keys(defs).length > 0) {
+    result.$defs = defs;
+  }
+
+  if (jsonSize(result) >= originalSize) return schema;
+  return result;
+}
+
+// ---------------------------------------------------------------------------
+// Public API with caching
+// ---------------------------------------------------------------------------
+
+let cachedFingerprint: string | null = null;
+let cachedTools: unknown[] | null = null;
+
+const objectIds = new WeakMap<object, number>();
+const symbolIds = new Map<symbol, number>();
+let nextOpaqueId = 0;
+
+function identityId(t: object): number {
+  let id = objectIds.get(t);
+  if (id === undefined) {
+    id = nextOpaqueId++;
+    objectIds.set(t, id);
+  }
+  return id;
+}
+
+function symbolId(t: symbol): number {
+  let id = symbolIds.get(t);
+  if (id === undefined) {
+    id = nextOpaqueId++;
+    symbolIds.set(t, id);
+  }
+  return id;
+}
+
+function isPlainObject(v: object): boolean {
+  const proto = Object.getPrototypeOf(v);
+  return proto === null || proto === Object.prototype;
+}
+
+function hashValue(hash: ReturnType<typeof createHash>, value: unknown, seen: Set<unknown>): void {
+  if (value === null) {
+    hash.update("z");
+    return;
+  }
+  if (value === undefined) {
+    hash.update("u");
+    return;
+  }
+  const t = typeof value;
+  switch (t) {
+    case "boolean":
+      hash.update(value ? "T" : "F");
+      return;
+    case "string":
+      hash.update(`s${(value as string).length}:${value as string}`);
+      return;
+    case "number":
+      hash.update(`n:${Object.is(value, -0) ? "-0" : String(value)}`);
+      return;
+    case "bigint":
+      hash.update(`B:${value}`);
+      return;
+    case "symbol":
+      hash.update(`Y:${symbolId(value as symbol)}`);
+      return;
+    case "function":
+      hash.update(`f:${identityId(value as object)}`);
+      return;
+  }
+  if (seen.has(value)) {
+    hash.update(`R:${identityId(value as object)}`);
+    return;
+  }
+  seen.add(value);
+  if (Array.isArray(value)) {
+    hash.update("[");
+    for (let i = 0; i < value.length; i++) hashValue(hash, value[i], seen);
+    hash.update("]");
+    return;
+  }
+  if (!isPlainObject(value as object)) {
+    hash.update(`O:${identityId(value as object)}`);
+    return;
+  }
+  hash.update("{");
+  const entries = Object.entries(value as Record<string, unknown>);
+  for (const [k, v] of entries) {
+    hash.update(`k${k.length}:${k}`);
+    hashValue(hash, v, seen);
+  }
+  hash.update("}");
+}
+
+function toolsFingerprint(tools: unknown[]): string {
+  const hash = createHash("sha256");
+  const seen = new Set<unknown>();
+  for (const t of tools) {
+    hashValue(hash, t, seen);
+    hash.update("|");
+  }
+  return hash.digest("hex");
+}
+
+export function optimizeToolSchemas(tools: unknown[]): unknown[] {
+  const fp = toolsFingerprint(tools);
+  if (fp === cachedFingerprint && cachedTools) return cachedTools;
+
+  let changed = false;
+  const result = tools.map((tool) => {
+    if (!isRecord(tool) || !isRecord(tool.function)) return tool;
+    const fn = tool.function as JsonRecord;
+    const params = fn.parameters;
+    if (!isRecord(params)) return tool;
+
+    const size = jsonSize(params);
+    if (size <= TOOL_SCHEMA_SIZE_THRESHOLD) return tool;
+
+    const optimized = deduplicateSchema(params);
+    changed = true;
+    return { ...tool, function: { ...fn, parameters: optimized } };
+  });
+
+  cachedFingerprint = fp;
+  cachedTools = changed ? result : tools;
+  return cachedTools;
+}
+
+export function resetToolSchemaCache(): void {
+  cachedFingerprint = null;
+  cachedTools = null;
+}

package/src/stream.ts

@@ -17,18 +17,35 @@ import {
   streamSimpleAnthropic,
   streamSimpleOpenAICompletions,
 } from "@earendil-works/pi-ai";
+import {
+  DEFAULT_KIMI_CODE_CONFIG,
+  type KimiCodeConfig,
+  type KimiResolvedModelConfig,
+} from "./config.ts";
 
-import { IS_OPENAI_PROTOCOL, PROTOCOL } from "./constants.ts";
+import { ENV_KIMI_CODE_PROTOCOL, getApiProtocol, getBaseUrl } from "./constants.ts";
+import { getCommonHeaders } from "./device.ts";
 import { isKimiAuthErrorMessage, refreshKimiAuthToken } from "./oauth.ts";
 import {
   type Uploader,
   applyKimiPayloadMutations,
   isRecord,
-  readEnvOverrides,
   resolveCacheRetention,
   uploadKimiFile,
 } from "./payload.ts";
 
+interface KimiStreamRuntimeConfig {
+  model: KimiResolvedModelConfig;
+  protocol: KimiCodeConfig["protocol"];
+  uploads: KimiCodeConfig["uploads"];
+}
+
+let resolvedStore: KimiStreamRuntimeConfig | null = null;
+
+export function setStoreResolvedKimiConfig(config: KimiStreamRuntimeConfig): void {
+  resolvedStore = config;
+}
+
 // =============================================================================
 // Event stream filter: suppress Kimi "(Empty response: ...)" text blocks
 // =============================================================================
@@ -136,21 +153,47 @@ export async function* filterEmptyResponseStream(
 // Stream wrapper: orchestrates payload mutation + event filter
 // =============================================================================
 
+const KIMI_API_KEY_ENV_REFERENCES = new Set(["$KIMI_API_KEY", "${KIMI_API_KEY}"]);
+
+export function resolveKimiApiKey(apiKey: string | undefined): string {
+  if (apiKey !== undefined && KIMI_API_KEY_ENV_REFERENCES.has(apiKey)) {
+    return process.env.KIMI_API_KEY?.trim() || "";
+  }
+  return apiKey || process.env.KIMI_API_KEY || "";
+}
+
+export function mergeKimiRequestHeaders(headers?: Record<string, string>): Record<string, string> {
+  return { ...getCommonHeaders(), ...headers };
+}
+
 export function streamSimpleKimi(
   model: Model<Api>,
   context: Context,
   options?: SimpleStreamOptions,
 ): AssistantMessageEventStream {
   const filtered = createAssistantMessageEventStream();
-  const initialKey = options?.apiKey || process.env.KIMI_API_KEY || "";
+  const initialKey = resolveKimiApiKey(options?.apiKey);
 
   const cacheKeyOverride = (
     options as (SimpleStreamOptions & { prompt_cache_key?: unknown }) | undefined
   )?.prompt_cache_key;
   const cacheKey = (typeof cacheKeyOverride === "string" && cacheKeyOverride) || options?.sessionId;
   const cacheRetention = resolveCacheRetention(options?.cacheRetention);
-  const envOverrides = readEnvOverrides();
-  const thinkingKeep = process.env.KIMI_MODEL_THINKING_KEEP;
+  const streamConfig = resolvedStore ?? {
+    model: {
+      contextWindow: model.contextWindow,
+      maxTokens: model.maxTokens,
+      input: ["text"],
+      reasoning: model.reasoning,
+      reasoningMap: {},
+      thinkingKeep: null,
+      generation: {},
+    },
+    protocol: ENV_KIMI_CODE_PROTOCOL,
+    uploads: DEFAULT_KIMI_CODE_CONFIG.uploads,
+  };
+  const modelConfig = streamConfig.model;
+  const apiProtocol = getApiProtocol(streamConfig.protocol);
   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
@@ -158,10 +201,10 @@ export function streamSimpleKimi(
   // the model object via modifyModels and rewritten into the request payload
   // here so /v1/chat/completions and /v1/messages see the real wire id.
   const wireModelId = (model as Model<Api> & { wireModelId?: unknown }).wireModelId;
-
   const buildPatchedOptions = (apiKey: string): SimpleStreamOptions => {
     const upload: Uploader | undefined = apiKey
-      ? (mimeType, data) => uploadKimiFile(apiKey, mimeType, data)
+      ? (mimeType, data) =>
+          uploadKimiFile(apiKey, mimeType, data, streamConfig.uploads.thresholdBytes)
       : undefined;
     // Only forward apiKey if we actually have one — never override the
     // caller's credential (e.g. Claude Code OAuth token) with an empty string.
@@ -169,18 +212,18 @@ export function streamSimpleKimi(
     return {
       ...options,
       ...apiKeyOverride,
+      headers: mergeKimiRequestHeaders(options?.headers),
       onPayload: async (payload, modelData) => {
         let nextPayload: unknown = payload;
 
         if (isRecord(nextPayload)) {
           await applyKimiPayloadMutations(nextPayload, {
-            api: PROTOCOL,
+            api: apiProtocol,
             upload,
             cacheKey,
             cacheRetention,
             reasoning: options?.reasoning,
-            thinkingKeep,
-            envOverrides,
+            modelConfig,
           });
           if (
             typeof wireModelId === "string" &&
@@ -211,13 +254,23 @@ export function streamSimpleKimi(
       // register with a custom api type (kimi-openai-completions /
       // kimi-anthropic-messages) to avoid overriding the built-in
       // Anthropic/OpenAI stream handlers.
-      const upstream = IS_OPENAI_PROTOCOL
-        ? streamSimpleOpenAICompletions(
-            model as Model<"openai-completions">,
-            context,
-            patchedOptions,
-          )
-        : streamSimpleAnthropic(model as Model<"anthropic-messages">, context, patchedOptions);
+      const runtimeModel = {
+        ...model,
+        api: apiProtocol,
+        baseUrl: getBaseUrl(streamConfig.protocol),
+      } as Model<Api>;
+      const upstream =
+        streamConfig.protocol === "openai"
+          ? streamSimpleOpenAICompletions(
+              runtimeModel as Model<"openai-completions">,
+              context,
+              patchedOptions,
+            )
+          : streamSimpleAnthropic(
+              runtimeModel as Model<"anthropic-messages">,
+              context,
+              patchedOptions,
+            );
 
       let shouldRetry = false;
       let prefixBuffer: AssistantMessageEvent[] = [];

package/src/tools/common.ts

@@ -0,0 +1,115 @@
+import { AuthStorage, type AgentToolResult } from "@earendil-works/pi-coding-agent";
+import { truncateToWidth } from "@earendil-works/pi-tui";
+
+import { PROVIDER_ID, getBaseUrl } from "../constants.ts";
+import { getCommonHeaders } from "../device.ts";
+import { refreshKimiAuthToken } from "../oauth.ts";
+
+export const KIMI_TOOL_TIMEOUT_MS = 180_000;
+
+export interface KimiToolDeps {
+  fetch: typeof fetch;
+  getAccessToken: () => string | null;
+  refreshAccessToken: (currentToken: string) => Promise<string | null>;
+}
+
+export interface BuildKimiToolOptions {
+  deps?: Partial<KimiToolDeps>;
+  defaultCollapsed?: boolean;
+}
+
+function defaultGetAccessToken(): string | null {
+  const credential = AuthStorage.create().get(PROVIDER_ID);
+  if (!credential || credential.type !== "oauth") return null;
+  return credential.access;
+}
+
+export function buildKimiToolDeps(options: BuildKimiToolOptions = {}): KimiToolDeps {
+  return {
+    fetch: options.deps?.fetch ?? fetch,
+    getAccessToken: options.deps?.getAccessToken ?? defaultGetAccessToken,
+    refreshAccessToken: options.deps?.refreshAccessToken ?? refreshKimiAuthToken,
+  };
+}
+
+export function buildTimeoutSignal(signal: AbortSignal | undefined): {
+  signal: AbortSignal;
+  cleanup: () => void;
+} {
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), KIMI_TOOL_TIMEOUT_MS);
+  const abort = () => controller.abort();
+  if (signal?.aborted) {
+    controller.abort();
+  } else {
+    signal?.addEventListener("abort", abort, { once: true });
+  }
+
+  return {
+    signal: controller.signal,
+    cleanup: () => {
+      clearTimeout(timeout);
+      signal?.removeEventListener("abort", abort);
+    },
+  };
+}
+
+export function getKimiBaseV1(): string {
+  const base = getBaseUrl().replace(/\/+$/, "");
+  return base.endsWith("/v1") ? base : `${base}/v1`;
+}
+
+export function getKimiDatasourceUrl(): string {
+  const explicit = process.env.KIMI_DATASOURCE_API_URL?.trim();
+  if (explicit) return explicit;
+  return `${getKimiBaseV1()}/tools`;
+}
+
+export function buildHeaders(accessToken: string, toolCallId: string): Record<string, string> {
+  return {
+    ...getCommonHeaders(),
+    Authorization: `Bearer ${accessToken}`,
+    "Content-Type": "application/json",
+    "X-Msh-Tool-Call-Id": toolCallId,
+  };
+}
+
+export function errorResult<T>(message: string): AgentToolResult<T> {
+  return {
+    content: [{ type: "text", text: message }],
+    details: undefined as T,
+  };
+}
+
+export async function readErrorBody(response: Response): Promise<string> {
+  return response.text().catch(() => "");
+}
+
+export async function fetchWithAuthRetry(
+  deps: KimiToolDeps,
+  accessToken: string,
+  request: (token: string) => Promise<Response>,
+): Promise<Response> {
+  const response = await request(accessToken);
+  if (response.status !== 401) return response;
+
+  const refreshed = await deps.refreshAccessToken(accessToken);
+  if (!refreshed || refreshed === accessToken) return response;
+  return request(refreshed);
+}
+
+export function textComponent(text: string) {
+  return {
+    render: (width: number) => text.split("\n").map((line) => truncateToWidth(line, width, "…")),
+    invalidate: () => {},
+  };
+}
+
+export function firstText(result: AgentToolResult<unknown>): string {
+  const first = result.content[0];
+  return first?.type === "text" ? first.text : "";
+}
+
+export function shouldCollapse(defaultCollapsed: boolean, expanded: boolean | undefined): boolean {
+  return defaultCollapsed && expanded !== true;
+}