From d5ea5f27ac2b784a23d86420649e56e103d4afa7 Mon Sep 17 00:00:00 2001 From: Ayaan Zaidi Date: Fri, 3 Apr 2026 09:39:19 +0530 Subject: [PATCH] fix: parse kimi tagged tool calls (#60051) * fix: parse kimi tagged tool calls * fix: parse kimi tagged tool calls (#60051) * fix: parse kimi tagged tool calls (#60051) --- CHANGELOG.md | 1 + extensions/kimi-coding/index.ts | 2 + extensions/kimi-coding/stream.test.ts | 211 ++++++++++++++++++++++++++ extensions/kimi-coding/stream.ts | 181 ++++++++++++++++++++++ 4 files changed, 395 insertions(+) create mode 100644 extensions/kimi-coding/stream.test.ts create mode 100644 extensions/kimi-coding/stream.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 6bb2e5c5ac9..51e6323aa33 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,7 @@ Docs: https://docs.openclaw.ai - MS Teams/threading: preserve channel reply threading when proactive fallback reconstructs the conversation ID, so thread replies land in the originating thread instead of the channel root. (#55198) Thanks @hyojin. - Gateway/lock: detect PID recycling in gateway lock files on Windows and macOS so stale locks from crashed processes no longer block new gateway invocations. (#59843) Thanks @TonyDerek-dot. - Plugins/Ollama: prefer real cloud auth credentials over the local-only marker so Ollama instances behind authenticated proxies use the configured API key instead of skipping auth. +- Plugins/Kimi Coding: parse tagged Kimi tool-call text into structured tool calls on the provider stream path so tools execute instead of echoing raw markup. (#60051) Thanks @obviyus. ## 2026.4.2 diff --git a/extensions/kimi-coding/index.ts b/extensions/kimi-coding/index.ts index 6aa4c17605d..fd717037f37 100644 --- a/extensions/kimi-coding/index.ts +++ b/extensions/kimi-coding/index.ts @@ -2,6 +2,7 @@ import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry"; import { createProviderApiKeyAuthMethod } from "openclaw/plugin-sdk/provider-auth-api-key"; import { applyKimiCodeConfig, KIMI_CODING_MODEL_REF } from "./onboard.js"; import { buildKimiCodingProvider } from "./provider-catalog.js"; +import { createKimiToolCallMarkupWrapper } from "./stream.js"; const PLUGIN_ID = "kimi"; const PROVIDER_ID = "kimi"; @@ -85,6 +86,7 @@ export default definePluginEntry({ openAiPayloadNormalizationMode: "moonshot-thinking", preserveAnthropicThinkingSignatures: false, }, + wrapStreamFn: (ctx) => createKimiToolCallMarkupWrapper(ctx.streamFn), }); }, }); diff --git a/extensions/kimi-coding/stream.test.ts b/extensions/kimi-coding/stream.test.ts new file mode 100644 index 00000000000..d627f8f528e --- /dev/null +++ b/extensions/kimi-coding/stream.test.ts @@ -0,0 +1,211 @@ +import type { StreamFn } from "@mariozechner/pi-agent-core"; +import type { Context, Model } from "@mariozechner/pi-ai"; +import { describe, expect, it } from "vitest"; +import { createKimiToolCallMarkupWrapper } from "./stream.js"; + +type FakeStream = { + result: () => Promise; + [Symbol.asyncIterator]: () => AsyncIterator; +}; + +function createFakeStream(params: { events: unknown[]; resultMessage: unknown }): FakeStream { + return { + async result() { + return params.resultMessage; + }, + [Symbol.asyncIterator]() { + return (async function* () { + for (const event of params.events) { + yield event; + } + })(); + }, + }; +} + +const KIMI_TOOL_TEXT = + ' <|tool_calls_section_begin|> <|tool_call_begin|> functions.read:0 <|tool_call_argument_begin|> {"file_path":"./package.json"} <|tool_call_end|> <|tool_calls_section_end|>'; +const KIMI_MULTI_TOOL_TEXT = + ' <|tool_calls_section_begin|> <|tool_call_begin|> functions.read:0 <|tool_call_argument_begin|> {"file_path":"./package.json"} <|tool_call_end|> <|tool_call_begin|> functions.write:1 <|tool_call_argument_begin|> {"file_path":"./out.txt","content":"done"} <|tool_call_end|> <|tool_calls_section_end|>'; + +describe("kimi tool-call markup wrapper", () => { + it("converts tagged Kimi tool-call text into structured tool calls", async () => { + const partial = { + role: "assistant", + content: [{ type: "text", text: KIMI_TOOL_TEXT }], + stopReason: "stop", + }; + const message = { + role: "assistant", + content: [{ type: "text", text: KIMI_TOOL_TEXT }], + stopReason: "stop", + }; + const finalMessage = { + role: "assistant", + content: [ + { type: "thinking", thinking: "Need to read the file first." }, + { type: "text", text: KIMI_TOOL_TEXT }, + ], + stopReason: "stop", + }; + + const baseStreamFn: StreamFn = () => + createFakeStream({ + events: [{ type: "message_end", partial, message }], + resultMessage: finalMessage, + }) as ReturnType; + + const wrapped = createKimiToolCallMarkupWrapper(baseStreamFn); + const stream = wrapped( + { api: "anthropic-messages", provider: "kimi", id: "k2p5" } as Model<"anthropic-messages">, + { messages: [] } as Context, + {}, + ) as FakeStream; + + const events: unknown[] = []; + for await (const event of stream) { + events.push(event); + } + const result = (await stream.result()) as { + content: unknown[]; + stopReason: string; + }; + + expect(events).toEqual([ + { + type: "message_end", + partial: { + role: "assistant", + content: [ + { + type: "toolCall", + id: "functions.read:0", + name: "functions.read", + arguments: { file_path: "./package.json" }, + }, + ], + stopReason: "toolUse", + }, + message: { + role: "assistant", + content: [ + { + type: "toolCall", + id: "functions.read:0", + name: "functions.read", + arguments: { file_path: "./package.json" }, + }, + ], + stopReason: "toolUse", + }, + }, + ]); + expect(result).toEqual({ + role: "assistant", + content: [ + { type: "thinking", thinking: "Need to read the file first." }, + { + type: "toolCall", + id: "functions.read:0", + name: "functions.read", + arguments: { file_path: "./package.json" }, + }, + ], + stopReason: "toolUse", + }); + }); + + it("leaves normal assistant text unchanged", async () => { + const finalMessage = { + role: "assistant", + content: [{ type: "text", text: "normal response" }], + stopReason: "stop", + }; + const baseStreamFn: StreamFn = () => + createFakeStream({ + events: [], + resultMessage: finalMessage, + }) as ReturnType; + + const wrapped = createKimiToolCallMarkupWrapper(baseStreamFn); + const stream = wrapped( + { api: "anthropic-messages", provider: "kimi", id: "k2p5" } as Model<"anthropic-messages">, + { messages: [] } as Context, + {}, + ) as FakeStream; + + await expect(stream.result()).resolves.toBe(finalMessage); + }); + + it("supports async stream functions", async () => { + const finalMessage = { + role: "assistant", + content: [{ type: "text", text: KIMI_TOOL_TEXT }], + stopReason: "stop", + }; + const baseStreamFn: StreamFn = async () => + createFakeStream({ + events: [], + resultMessage: finalMessage, + }) as ReturnType; + + const wrapped = createKimiToolCallMarkupWrapper(baseStreamFn); + const stream = (await wrapped( + { api: "anthropic-messages", provider: "kimi", id: "k2p5" } as Model<"anthropic-messages">, + { messages: [] } as Context, + {}, + )) as FakeStream; + + await expect(stream.result()).resolves.toEqual({ + role: "assistant", + content: [ + { + type: "toolCall", + id: "functions.read:0", + name: "functions.read", + arguments: { file_path: "./package.json" }, + }, + ], + stopReason: "toolUse", + }); + }); + + it("parses multiple tagged tool calls in one section", async () => { + const finalMessage = { + role: "assistant", + content: [{ type: "text", text: KIMI_MULTI_TOOL_TEXT }], + stopReason: "stop", + }; + const baseStreamFn: StreamFn = () => + createFakeStream({ + events: [], + resultMessage: finalMessage, + }) as ReturnType; + + const wrapped = createKimiToolCallMarkupWrapper(baseStreamFn); + const stream = wrapped( + { api: "anthropic-messages", provider: "kimi", id: "k2p5" } as Model<"anthropic-messages">, + { messages: [] } as Context, + {}, + ) as FakeStream; + + await expect(stream.result()).resolves.toEqual({ + role: "assistant", + content: [ + { + type: "toolCall", + id: "functions.read:0", + name: "functions.read", + arguments: { file_path: "./package.json" }, + }, + { + type: "toolCall", + id: "functions.write:1", + name: "functions.write", + arguments: { file_path: "./out.txt", content: "done" }, + }, + ], + stopReason: "toolUse", + }); + }); +}); diff --git a/extensions/kimi-coding/stream.ts b/extensions/kimi-coding/stream.ts new file mode 100644 index 00000000000..d497aed55d0 --- /dev/null +++ b/extensions/kimi-coding/stream.ts @@ -0,0 +1,181 @@ +import type { StreamFn } from "@mariozechner/pi-agent-core"; +import { streamSimple } from "@mariozechner/pi-ai"; + +const TOOL_CALLS_SECTION_BEGIN = "<|tool_calls_section_begin|>"; +const TOOL_CALLS_SECTION_END = "<|tool_calls_section_end|>"; +const TOOL_CALL_BEGIN = "<|tool_call_begin|>"; +const TOOL_CALL_ARGUMENT_BEGIN = "<|tool_call_argument_begin|>"; +const TOOL_CALL_END = "<|tool_call_end|>"; + +type KimiToolCallBlock = { + type: "toolCall"; + id: string; + name: string; + arguments: Record; +}; + +function stripTaggedToolCallCounter(value: string): string { + return value.trim().replace(/:\d+$/, ""); +} + +function parseKimiTaggedToolCalls(text: string): KimiToolCallBlock[] | null { + const trimmed = text.trim(); + // Kimi emits tagged tool-call sections as standalone text blocks on this path. + if (!trimmed.startsWith(TOOL_CALLS_SECTION_BEGIN) || !trimmed.endsWith(TOOL_CALLS_SECTION_END)) { + return null; + } + + let cursor = TOOL_CALLS_SECTION_BEGIN.length; + const sectionEndIndex = trimmed.length - TOOL_CALLS_SECTION_END.length; + const toolCalls: KimiToolCallBlock[] = []; + + while (cursor < sectionEndIndex) { + while (cursor < sectionEndIndex && /\s/.test(trimmed[cursor] ?? "")) { + cursor += 1; + } + if (cursor >= sectionEndIndex) { + break; + } + if (!trimmed.startsWith(TOOL_CALL_BEGIN, cursor)) { + return null; + } + + const nameStart = cursor + TOOL_CALL_BEGIN.length; + const argMarkerIndex = trimmed.indexOf(TOOL_CALL_ARGUMENT_BEGIN, nameStart); + if (argMarkerIndex < 0 || argMarkerIndex >= sectionEndIndex) { + return null; + } + + const rawId = trimmed.slice(nameStart, argMarkerIndex).trim(); + if (!rawId) { + return null; + } + + const argsStart = argMarkerIndex + TOOL_CALL_ARGUMENT_BEGIN.length; + const callEndIndex = trimmed.indexOf(TOOL_CALL_END, argsStart); + if (callEndIndex < 0 || callEndIndex > sectionEndIndex) { + return null; + } + + const rawArgs = trimmed.slice(argsStart, callEndIndex).trim(); + let parsedArgs: unknown; + try { + parsedArgs = JSON.parse(rawArgs); + } catch { + return null; + } + if (!parsedArgs || typeof parsedArgs !== "object" || Array.isArray(parsedArgs)) { + return null; + } + + const name = stripTaggedToolCallCounter(rawId); + if (!name) { + return null; + } + + toolCalls.push({ + type: "toolCall", + id: rawId, + name, + arguments: parsedArgs as Record, + }); + + cursor = callEndIndex + TOOL_CALL_END.length; + } + + return toolCalls.length > 0 ? toolCalls : null; +} + +function rewriteKimiTaggedToolCallsInMessage(message: unknown): void { + if (!message || typeof message !== "object") { + return; + } + + const content = (message as { content?: unknown }).content; + if (!Array.isArray(content)) { + return; + } + + let changed = false; + const nextContent: unknown[] = []; + for (const block of content) { + if (!block || typeof block !== "object") { + nextContent.push(block); + continue; + } + const typedBlock = block as { type?: unknown; text?: unknown }; + if (typedBlock.type !== "text" || typeof typedBlock.text !== "string") { + nextContent.push(block); + continue; + } + + const parsed = parseKimiTaggedToolCalls(typedBlock.text); + if (!parsed) { + nextContent.push(block); + continue; + } + + nextContent.push(...parsed); + changed = true; + } + + if (!changed) { + return; + } + + (message as { content: unknown[] }).content = nextContent; + const typedMessage = message as { stopReason?: unknown }; + if (typedMessage.stopReason === "stop") { + typedMessage.stopReason = "toolUse"; + } +} + +function wrapKimiTaggedToolCalls( + stream: ReturnType, +): ReturnType { + const originalResult = stream.result.bind(stream); + stream.result = async () => { + const message = await originalResult(); + rewriteKimiTaggedToolCallsInMessage(message); + return message; + }; + + const originalAsyncIterator = stream[Symbol.asyncIterator].bind(stream); + (stream as { [Symbol.asyncIterator]: typeof originalAsyncIterator })[Symbol.asyncIterator] = + function () { + const iterator = originalAsyncIterator(); + return { + async next() { + const result = await iterator.next(); + if (!result.done && result.value && typeof result.value === "object") { + const event = result.value as { + partial?: unknown; + message?: unknown; + }; + rewriteKimiTaggedToolCallsInMessage(event.partial); + rewriteKimiTaggedToolCallsInMessage(event.message); + } + return result; + }, + async return(value?: unknown) { + return iterator.return?.(value) ?? { done: true as const, value: undefined }; + }, + async throw(error?: unknown) { + return iterator.throw?.(error) ?? { done: true as const, value: undefined }; + }, + }; + }; + + return stream; +} + +export function createKimiToolCallMarkupWrapper(baseStreamFn: StreamFn | undefined): StreamFn { + const underlying = baseStreamFn ?? streamSimple; + return (model, context, options) => { + const maybeStream = underlying(model, context, options); + if (maybeStream && typeof maybeStream === "object" && "then" in maybeStream) { + return Promise.resolve(maybeStream).then((stream) => wrapKimiTaggedToolCalls(stream)); + } + return wrapKimiTaggedToolCalls(maybeStream); + }; +}