openclaw/extensions/kimi-coding/kimi-coding.live.test.ts

import type { StreamFn } from "openclaw/plugin-sdk/agent-core";
import {
  streamSimple,
  type AssistantMessage,
  type Context,
  type Model,
  type Tool,
} from "openclaw/plugin-sdk/llm";
import { registerSingleProviderPlugin } from "openclaw/plugin-sdk/plugin-test-runtime";
import { extractNonEmptyAssistantText, isLiveTestEnabled } from "openclaw/plugin-sdk/test-live";
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import plugin from "./index.js";
import { buildKimiCodingProvider, normalizeKimiCodingModelId } from "./provider-catalog.js";

const describeLive =
  isLiveTestEnabled() && process.env.KIMI_API_KEY?.trim() ? describe : describe.skip;

async function collectDoneMessage(stream: ReturnType<StreamFn>): Promise<AssistantMessage> {
  let doneMessage: AssistantMessage | undefined;
  for await (const event of await stream) {
    if (event.type === "error") {
      throw new Error(event.error?.errorMessage || "Kimi live request failed");
    }
    if (event.type === "done") {
      doneMessage = event.message;
    }
  }
  if (!doneMessage) {
    throw new Error("Kimi live stream ended without a done message");
  }
  return doneMessage;
}

function resolveModel(modelId: "k3" | "k3-256k"): Model<"anthropic-messages"> {
  const provider = buildKimiCodingProvider();
  const normalizedModelId = normalizeKimiCodingModelId(modelId);
  const definition = provider.models.find((model) => model.id === normalizedModelId);
  if (!definition) {
    throw new Error(`Missing model ${modelId}`);
  }
  return {
    provider: "kimi",
    baseUrl: provider.baseUrl,
    headers: provider.headers,
    ...definition,
    api: "anthropic-messages",
  } as Model<"anthropic-messages">;
}

function countContentChars(message: AssistantMessage, type: "text" | "thinking"): number {
  return message.content.reduce((total, block) => {
    if (type === "text" && block.type === "text") {
      return total + block.text.length;
    }
    if (type === "thinking" && block.type === "thinking") {
      return total + block.thinking.length;
    }
    return total;
  }, 0);
}

async function runReasoningScenario(params: {
  modelId: "k3" | "k3-256k";
  thinkingLevel: "off" | "low" | "adaptive" | "max";
}): Promise<AssistantMessage> {
  const registered = await registerSingleProviderPlugin(plugin);
  const wrapped = registered.wrapStreamFn?.({
    provider: "kimi",
    modelId: params.modelId,
    thinkingLevel: params.thinkingLevel,
    extraParams: { thinking: params.thinkingLevel === "off" ? "off" : "enabled" },
    streamFn: streamSimple,
  } as never);
  if (!wrapped) {
    throw new Error("Missing Kimi stream wrapper");
  }

  const context: Context = {
    messages: [
      {
        role: "user",
        content: "Reply with exactly LIVE_OK and no punctuation.",
        timestamp: Date.now(),
      },
    ],
  };
  return collectDoneMessage(
    wrapped(resolveModel(params.modelId), context, {
      apiKey: process.env.KIMI_API_KEY?.trim() ?? "",
      maxTokens: 4096,
    }),
  );
}

describeLive("Kimi Code K3 reasoning live", () => {
  it.each(["k3", "k3-256k"] as const)(
    "%s honors off and max reasoning",
    async (modelId) => {
      const off = await runReasoningScenario({ modelId, thinkingLevel: "off" });
      expect(countContentChars(off, "thinking")).toBe(0);
      expect(countContentChars(off, "text")).toBeGreaterThan(0);

      const max = await runReasoningScenario({ modelId, thinkingLevel: "max" });
      expect(countContentChars(max, "thinking")).toBeGreaterThan(0);
      expect(countContentChars(max, "text")).toBeGreaterThan(0);
    },
    180_000,
  );

  it.each(["low", "adaptive"] as const)(
    "k3 accepts %s reasoning",
    async (thinkingLevel) => {
      const message = await runReasoningScenario({ modelId: "k3", thinkingLevel });
      expect(countContentChars(message, "thinking")).toBeGreaterThan(0);
      expect(countContentChars(message, "text")).toBeGreaterThan(0);
    },
    180_000,
  );

  it("preserves reasoning across a K3 tool-result replay", async () => {
    const registered = await registerSingleProviderPlugin(plugin);
    const wrapped = registered.wrapStreamFn?.({
      provider: "kimi",
      modelId: "k3",
      thinkingLevel: "max",
      streamFn: streamSimple,
    } as never);
    if (!wrapped) {
      throw new Error("Missing Kimi stream wrapper");
    }

    const sourcePath = "src/retry-delay.ts";
    const source =
      "export function retryDelayMs(attempt: number): number {\n" +
      "  return Math.min(250 * 2 ** attempt, 4000);\n" +
      "}\n";
    const tool: Tool = {
      name: "read_file",
      description: "Read a source file by its relative path.",
      parameters: Type.Object({ path: Type.String() }, { additionalProperties: false }),
    };
    const firstUser = {
      role: "user" as const,
      content:
        `Read ${sourcePath} with read_file and determine what retryDelayMs(3) returns. ` +
        "Read the implementation before answering, then return only the numeric result.",
      timestamp: Date.now(),
    };
    const model = resolveModel("k3");
    const options = { apiKey: process.env.KIMI_API_KEY?.trim() ?? "", maxTokens: 4096 };
    const first = await collectDoneMessage(
      wrapped(model, { messages: [firstUser], tools: [tool] }, options),
    );
    expect(countContentChars(first, "thinking")).toBeGreaterThan(0);
    const toolCall = first.content.find((block) => block.type === "toolCall");
    if (!toolCall || toolCall.type !== "toolCall") {
      throw new Error(`Kimi K3 did not read the source file: ${first.stopReason}`);
    }
    expect(toolCall.name).toBe("read_file");
    expect(toolCall.arguments).toEqual({ path: sourcePath });

    const second = await collectDoneMessage(
      wrapped(
        model,
        {
          messages: [
            firstUser,
            first,
            {
              role: "toolResult",
              toolCallId: toolCall.id,
              toolName: toolCall.name,
              content: [{ type: "text", text: source }],
              isError: false,
              timestamp: Date.now(),
            },
          ],
          tools: [tool],
        },
        options,
      ),
    );
    expect(extractNonEmptyAssistantText(second.content)).toBe("2000");
  }, 180_000);
});