pi-provider-kimi-code 0.6.0 → 0.6.2

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,19 +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
 // =============================================================================
@@ -163,8 +179,21 @@ export function streamSimpleKimi(
   )?.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
@@ -172,13 +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 supportsThinkingType = (
-    model as Model<Api> & { supportsThinkingType?: "only" | "no" | "both" }
-  ).supportsThinkingType;
-
   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.
@@ -192,14 +218,12 @@ export function streamSimpleKimi(
 
         if (isRecord(nextPayload)) {
           await applyKimiPayloadMutations(nextPayload, {
-            api: PROTOCOL,
+            api: apiProtocol,
             upload,
             cacheKey,
             cacheRetention,
             reasoning: options?.reasoning,
-            thinkingKeep,
-            envOverrides,
-            supportsThinkingType,
+            modelConfig,
           });
           if (
             typeof wireModelId === "string" &&
@@ -230,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/datasource.ts

@@ -104,6 +104,45 @@ function extractResponseText(data: DatasourceToolResponse): string {
   return JSON.stringify(data, null, 2);
 }
 
+function extractResponseFiles(data: DatasourceToolResponse): string[] {
+  const raw = data as Record<string, unknown>;
+  if (!Array.isArray(raw.files)) return [];
+  const names: string[] = [];
+  for (const file of raw.files) {
+    if (
+      typeof file === "object" &&
+      file !== null &&
+      typeof (file as Record<string, unknown>).name === "string"
+    ) {
+      names.push((file as Record<string, unknown>).name as string);
+    }
+  }
+  return names;
+}
+
+const REQUEST_ID_HEADERS = [
+  "x-request-id",
+  "x-trace-id",
+  "x-msh-request-id",
+  "x-msh-trace-id",
+  "request-id",
+];
+
+function extractRequestId(headers: Headers): string | undefined {
+  for (const key of REQUEST_ID_HEADERS) {
+    const value = headers.get(key);
+    if (typeof value === "string" && value.trim()) return value.trim();
+  }
+  return undefined;
+}
+
+function buildTraceLine(toolCallId: string, requestId: string | undefined): string {
+  const parts: string[] = [];
+  if (requestId) parts.push(`request-id: ${requestId}`);
+  parts.push(`tool-call-id: ${toolCallId}`);
+  return `\n\n[kimi-datasource] ${parts.join(" · ")}`;
+}
+
 export function buildKimiDatasourceTool(options: BuildKimiToolOptions = {}) {
   const deps = buildKimiToolDeps(options);
   const defaultCollapsed = options.defaultCollapsed ?? true;
@@ -138,7 +177,7 @@ export function buildKimiDatasourceTool(options: BuildKimiToolOptions = {}) {
             }
           : {
               method: "get_data_source_desc",
-              params: { data_source_name: params.data_source_name },
+              params: { name: params.data_source_name },
             };
 
         const response = await fetchWithAuthRetry(deps, accessToken, (token) =>
@@ -150,10 +189,13 @@ export function buildKimiDatasourceTool(options: BuildKimiToolOptions = {}) {
           }),
         );
 
+        const requestId = extractRequestId(response.headers);
+        const trace = buildTraceLine(toolCallId, requestId);
+
         if (!response.ok) {
           const bodyText = await readErrorBody(response);
           return errorResult<string>(
-            `kimi_datasource failed: ${response.status}${bodyText ? ` ${bodyText}` : ""}`,
+            `kimi_datasource failed: ${response.status}${bodyText ? ` ${bodyText}` : ""}${trace}`,
           );
         }
 
@@ -163,10 +205,16 @@ export function buildKimiDatasourceTool(options: BuildKimiToolOptions = {}) {
             typeof data.error === "string"
               ? data.error
               : JSON.stringify(data.error ?? data, null, 2);
-          return errorResult<string>(`kimi_datasource error: ${errorText}`);
+          return errorResult<string>(`kimi_datasource error: ${errorText}${trace}`);
+        }
+
+        let text = extractResponseText(data);
+        const files = extractResponseFiles(data);
+        if (files.length > 0) {
+          text += `\n\nNote: the backend returned ${files.length} file(s) that cannot be written in this context: ${files.join(", ")}`;
         }
+        text += trace;
 
-        const text = extractResponseText(data);
         return {
           content: [{ type: "text", text }],
           details: text,