From f16b4e0cda56cef74bb92e264f8561cf8f4c1385 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Thu, 16 Jul 2026 12:26:56 +0000 Subject: [PATCH] add deferred tools support for kimi in openai-completions api --- packages/ai/scripts/generate-models.ts | 3 +- packages/ai/src/api/openai-completions.ts | 59 +++++++++++- packages/ai/src/types.ts | 2 + packages/ai/test/deferred-tools.test.ts | 94 +++++++++++++++++++ ...penai-completions-thinking-as-text.test.ts | 3 +- ...nai-completions-tool-result-images.test.ts | 4 +- packages/coding-agent/docs/models.md | 1 + .../examples/extensions/README.md | 1 + .../extensions/kimi-deferred-tools.ts | 61 ++++++++++++ .../coding-agent/src/core/model-config.ts | 1 + 10 files changed, 223 insertions(+), 6 deletions(-) create mode 100644 packages/coding-agent/examples/extensions/kimi-deferred-tools.ts diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 5fad8aaf..641717bd 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -373,8 +373,9 @@ const OPENAI_COMPLETIONS_DEFAULT_COMPAT = { supportsStrictMode: true, sendSessionAffinityHeaders: false, supportsLongCacheRetention: true, -} satisfies Required> & { +} satisfies Required> & { cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; + deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"]; }; type OpenAICompletionsResolvedCompat = typeof OPENAI_COMPLETIONS_DEFAULT_COMPAT & { diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 88df3143..61001777 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -77,6 +77,26 @@ function hasToolHistory(messages: Message[]): boolean { return false; } +function getDeferredToolNames(messages: Message[]): Set { + const names = new Set(); + for (const message of messages) { + if (message.role === "toolResult") { + for (const name of message.addedToolNames ?? []) { + names.add(name); + } + } + } + return names; +} + +function getToolsByName(tools: Tool[] | undefined, names: Iterable): Tool[] { + if (!tools) return []; + const toolsByName = new Map(tools.map((tool) => [tool.name, tool])); + return Array.from(names) + .map((name) => toolsByName.get(name)) + .filter((tool): tool is Tool => tool !== undefined); +} + function isTextContentBlock(block: { type: string }): block is TextContent { return block.type === "text"; } @@ -117,14 +137,23 @@ interface OpenAICompatCacheControl { ttl?: string; } -type ResolvedOpenAICompletionsCompat = Omit, "cacheControlFormat"> & { +type ResolvedOpenAICompletionsCompat = Omit< + Required, + "cacheControlFormat" | "deferredToolsMode" +> & { cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; + deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"]; }; type ResolvedChatTemplateKwargValue = string | number | boolean | null; type ChatCompletionInstructionMessageParam = ChatCompletionDeveloperMessageParam | ChatCompletionSystemMessageParam; +type KimiToolSystemMessageParam = { + role: "system"; + tools: OpenAI.Chat.Completions.ChatCompletionTool[]; +}; + type OpenAIEncryptedReasoningDetail = { type: "reasoning.encrypted"; id: string; @@ -585,8 +614,11 @@ function buildParams( params.temperature = options.temperature; } - if (context.tools && context.tools.length > 0) { - params.tools = convertTools(context.tools, compat); + const deferredToolNames = + compat.deferredToolsMode === "kimi" ? getDeferredToolNames(context.messages) : new Set(); + const activeTools = context.tools?.filter((tool) => !deferredToolNames.has(tool.name)); + if (activeTools && activeTools.length > 0) { + params.tools = convertTools(activeTools, compat); if (compat.zaiToolStream) { (params as any).tool_stream = true; } @@ -1025,6 +1057,7 @@ export function convertMessages( params.push(assistantMsg); } else if (msg.role === "toolResult") { const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = []; + const deferredToolNames = new Set(); let j = i; for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) { @@ -1051,6 +1084,12 @@ export function convertMessages( } params.push(toolResultMsg); + if (compat.deferredToolsMode === "kimi") { + for (const name of toolMsg.addedToolNames ?? []) { + deferredToolNames.add(name); + } + } + if (hasImages && model.input.includes("image")) { for (const block of toolMsg.content) { if (isImageContentBlock(block)) { @@ -1089,6 +1128,18 @@ export function convertMessages( } else { lastRole = "toolResult"; } + + if (deferredToolNames.size > 0) { + const deferredTools = getToolsByName(context.tools, deferredToolNames); + if (deferredTools.length > 0) { + const kimiToolMessage: KimiToolSystemMessageParam = { + role: "system", + tools: convertTools(deferredTools, compat), + }; + // Kimi accepts a system message with tools but omits the standard content field. + params.push(kimiToolMessage as unknown as ChatCompletionMessageParam); + } + } continue; } @@ -1256,6 +1307,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, cacheControlFormat, sendSessionAffinityHeaders: false, + deferredToolsMode: undefined, sessionAffinityFormat: isOpenRouter ? "openrouter" : "openai", supportsLongCacheRetention: !( isTogether || @@ -1296,6 +1348,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode, cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat, sendSessionAffinityHeaders: model.compat.sendSessionAffinityHeaders ?? detected.sendSessionAffinityHeaders, + deferredToolsMode: model.compat.deferredToolsMode ?? detected.deferredToolsMode, sessionAffinityFormat: model.compat.sessionAffinityFormat ?? detected.sessionAffinityFormat, supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention, }; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 69356b22..c123fd89 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -524,6 +524,8 @@ export interface OpenAICompletionsCompat { cacheControlFormat?: "anthropic"; /** Whether to send session-affinity data from `options.sessionId`. Default: false. */ sendSessionAffinityHeaders?: boolean; + /** Provider-specific deferred tool serialization mode. */ + deferredToolsMode?: "kimi"; /** Session-affinity header format: `openai` sends `session_id`, `x-client-request-id`, and `x-session-affinity`; `openai-nosession` sends `x-client-request-id` and `x-session-affinity`; `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param, which is governed by cache retention. Default: auto-detected. */ sessionAffinityFormat?: SessionAffinityFormat; /** Whether the provider supports long prompt cache retention (`prompt_cache_retention: "24h"` or Anthropic-style `cache_control.ttl: "1h"`, depending on format). Default: true. */ diff --git a/packages/ai/test/deferred-tools.test.ts b/packages/ai/test/deferred-tools.test.ts index f2567032..095c2fb9 100644 --- a/packages/ai/test/deferred-tools.test.ts +++ b/packages/ai/test/deferred-tools.test.ts @@ -1,5 +1,6 @@ import { Type } from "typebox"; import { describe, expect, it } from "vitest"; +import { convertMessages } from "../src/api/openai-completions.ts"; import { getModel, streamSimple } from "../src/compat.ts"; import type { Api, AssistantMessage, Context, Model, Tool, ToolResultMessage, UserMessage } from "../src/types.ts"; import { estimateContextTokens } from "../src/utils/estimate.ts"; @@ -49,6 +50,26 @@ interface OpenAIPayload { input?: Array; } +interface KimiTool { + type: "function"; + function: { + name: string; + description?: string; + parameters?: Record; + }; +} + +interface KimiMessage { + role: string; + content?: unknown; + tools?: KimiTool[]; +} + +interface KimiPayload { + tools?: KimiTool[]; + messages: KimiMessage[]; +} + class PayloadCaptured extends Error {} function makeTool(name: string): Tool { @@ -102,6 +123,22 @@ function makeContext(tools: Tool[], addedToolNames = ["late_tool"]): Context { }; } +function makeKimiModel(deferredToolsMode?: "kimi"): Model<"openai-completions"> { + return { + id: "deferred-tools-model", + name: "Deferred Tools Model", + api: "openai-completions", + provider: "moonshotai", + baseUrl: "http://127.0.0.1:9/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 4096, + compat: deferredToolsMode ? { deferredToolsMode } : undefined, + }; +} + async function capturePayload(model: Model, context: Context, apiKey = "fake-key"): Promise { let captured: T | undefined; const stream = streamSimple({ ...model, baseUrl: "http://127.0.0.1:9" }, context, { @@ -297,6 +334,63 @@ describe("deferred tools", () => { expect(payload.tools?.find((tool) => tool.name === "late_tool")?.defer_loading).toBe(true); }); + it("serializes Kimi deferred tools as system tool definitions", async () => { + const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]); + const payload = await capturePayload(makeKimiModel("kimi"), context); + + expect(payload.tools?.map((tool) => tool.function.name)).toEqual(["base_tool"]); + const toolResultIndex = payload.messages.findIndex((message) => message.role === "tool"); + const systemToolIndex = payload.messages.findIndex((message) => message.tools !== undefined); + expect(toolResultIndex).toBeGreaterThanOrEqual(0); + expect(systemToolIndex).toBeGreaterThan(toolResultIndex); + expect(payload.messages[systemToolIndex]?.tools?.map((tool) => tool.function.name)).toEqual(["late_tool"]); + }); + + it("emits Kimi deferred schemas after all tool results in a batch", () => { + const context = makeContext([makeTool("base_tool"), makeTool("late_tool"), makeTool("later_tool")]); + context.messages.splice(3, 0, { + ...makeToolResult(["later_tool"]), + toolCallId: "call_2", + }); + + const messages = convertMessages(makeKimiModel("kimi"), context, { + supportsStore: false, + supportsDeveloperRole: false, + supportsReasoningEffort: false, + supportsUsageInStreaming: true, + maxTokensField: "max_tokens", + requiresToolResultName: false, + requiresAssistantAfterToolResult: false, + requiresThinkingAsText: false, + requiresReasoningContentOnAssistantMessages: false, + thinkingFormat: "openai", + openRouterRouting: {}, + vercelGatewayRouting: {}, + chatTemplateKwargs: {}, + zaiToolStream: false, + supportsStrictMode: false, + cacheControlFormat: undefined, + sendSessionAffinityHeaders: false, + deferredToolsMode: "kimi", + sessionAffinityFormat: "openai", + supportsLongCacheRetention: false, + }); + + expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "tool", "tool", "system", "user"]); + expect((messages[4] as { tools?: KimiTool[] }).tools?.map((tool) => tool.function.name)).toEqual([ + "late_tool", + "later_tool", + ]); + }); + + it("leaves OpenAI Completions tools unchanged without Kimi mode", async () => { + const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]); + const payload = await capturePayload(makeKimiModel(), context); + + expect(payload.tools?.map((tool) => tool.function.name)).toEqual(["base_tool", "late_tool"]); + expect(payload.messages.some((message) => message.tools !== undefined)).toBe(false); + }); + it("loads an OpenAI Responses tool through client tool search", async () => { const context = makeContext([makeTool("base_tool"), makeTool("late_tool")]); const payload = await capturePayload(getModel("openai", "gpt-5.4"), context); diff --git a/packages/ai/test/openai-completions-thinking-as-text.test.ts b/packages/ai/test/openai-completions-thinking-as-text.test.ts index 88af01ea..6a2b125c 100644 --- a/packages/ai/test/openai-completions-thinking-as-text.test.ts +++ b/packages/ai/test/openai-completions-thinking-as-text.test.ts @@ -41,8 +41,9 @@ const compat = { sendSessionAffinityHeaders: false, sessionAffinityFormat: "openai", supportsLongCacheRetention: true, -} satisfies Required> & { +} satisfies Omit, "cacheControlFormat" | "deferredToolsMode"> & { cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; + deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"]; }; function buildModel(baseUrl = "http://127.0.0.1:1"): Model<"openai-completions"> { diff --git a/packages/ai/test/openai-completions-tool-result-images.test.ts b/packages/ai/test/openai-completions-tool-result-images.test.ts index f9ed9d30..4c3f8a9f 100644 --- a/packages/ai/test/openai-completions-tool-result-images.test.ts +++ b/packages/ai/test/openai-completions-tool-result-images.test.ts @@ -19,7 +19,9 @@ const emptyUsage: Usage = { cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }; -const compat: Required = { +const compat: Omit, "deferredToolsMode"> & { + deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"]; +} = { supportsStore: true, supportsDeveloperRole: true, supportsReasoningEffort: true, diff --git a/packages/coding-agent/docs/models.md b/packages/coding-agent/docs/models.md index c0a98716..50efc369 100644 --- a/packages/coding-agent/docs/models.md +++ b/packages/coding-agent/docs/models.md @@ -449,6 +449,7 @@ For providers with partial OpenAI compatibility, use the `compat` field. | `sendSessionAffinityHeaders` | For `openai-completions`, send session-affinity headers from the session id when caching is enabled. Default: `false`. | | `sessionAffinityFormat` | For `openai-completions` and `openai-responses`, the session-affinity header format: `openai` sends `session_id`/`x-client-request-id` (completions also `x-session-affinity`), `openai-nosession` omits the underscore-containing `session_id` header, `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param. Default: auto-detected. | | `supportsStrictMode` | Include the `strict` field in tool definitions | +| `deferredToolsMode` | Use provider-specific deferred tool serialization. Currently only `"kimi"` is supported for Kimi's OpenAI-compatible Chat Completions format. | | `supportsLongCacheRetention` | Whether the provider accepts long cache retention when cache retention is `long`: `prompt_cache_retention: "24h"` for OpenAI prompt caching, or `cache_control.ttl: "1h"` when `cacheControlFormat` is `anthropic`. Default: `true`. | | `openRouterRouting` | OpenRouter provider routing preferences. This object is sent as-is in the `provider` field of the [OpenRouter API request](https://openrouter.ai/docs/guides/routing/provider-selection). | | `vercelGatewayRouting` | Vercel AI Gateway routing config for provider selection (`only`, `order`) | diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index 3dbc14d0..c34e9825 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -36,6 +36,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/ | `questionnaire.ts` | Multi-question input with tab bar navigation between questions | | `tool-override.ts` | Override built-in tools (e.g., add logging/access control to `read`) | | `dynamic-tools.ts` | Register tools after startup (`session_start`) and at runtime via command, with prompt snippets and tool-specific prompt guidelines | +| `kimi-deferred-tools.ts` | Search for and progressively activate tools for Kimi's deferred-tool loading protocol | | `structured-output.ts` | Final structured-output tool that returns `terminate: true` so the agent can end on the tool call | | `built-in-tool-renderer.ts` | Custom compact rendering for built-in tools (read, bash, edit, write) while keeping original behavior | | `minimal-mode.ts` | Override built-in tool rendering for minimal display (only tool calls, no output in collapsed mode) | diff --git a/packages/coding-agent/examples/extensions/kimi-deferred-tools.ts b/packages/coding-agent/examples/extensions/kimi-deferred-tools.ts new file mode 100644 index 00000000..603dc57d --- /dev/null +++ b/packages/coding-agent/examples/extensions/kimi-deferred-tools.ts @@ -0,0 +1,61 @@ +/** + * Minimal Kimi deferred-tool loading demo. + * + * pi -e ./kimi-deferred-tools.ts + * example prompt: Use the available tools to calculate 100 + 500. Do not calculate it yourself. + */ + +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { Type } from "typebox"; + +function calculate(_expr: string): string { + return "42"; +} + +export default function (pi: ExtensionAPI): void { + pi.registerTool({ + name: "Calculator", + label: "Calculator", + description: "Evaluate a simple arithmetic expression.", + parameters: Type.Object({ + expr: Type.String({ description: "An expression such as 100 + 500" }), + }), + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: calculate(params.expr) }], + details: {}, + }; + }, + }); + + pi.registerTool({ + name: "tool_search", + label: "Tool Search", + description: "Find and activate tools for a capability.", + promptSnippet: "Search for additional tools when the active tools cannot perform the task", + parameters: Type.Object({ + query: Type.String({ description: "Capability to search for" }), + }), + async execute(_toolCallId, params) { + if (!params.query.toLowerCase().includes("calc")) { + return { + content: [{ type: "text", text: "The relevant tools do not exist." }], + details: { matches: [], added: [] }, + }; + } + + const active = pi.getActiveTools(); + const added = active.includes("Calculator") ? [] : ["Calculator"]; + if (added.length > 0) pi.setActiveTools([...active, ...added]); + + return { + content: [{ type: "text", text: "Success. Found 1 matching tool(s)" }], + details: { matches: ["Calculator"], added }, + }; + }, + }); + + pi.on("session_start", () => { + pi.setActiveTools(["tool_search"]); + }); +} diff --git a/packages/coding-agent/src/core/model-config.ts b/packages/coding-agent/src/core/model-config.ts index 54a9b0a0..d21b9e0b 100644 --- a/packages/coding-agent/src/core/model-config.ts +++ b/packages/coding-agent/src/core/model-config.ts @@ -99,6 +99,7 @@ const OpenAICompletionsCompatSchema = Type.Object({ vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema), supportsStrictMode: Type.Optional(Type.Boolean()), sendSessionAffinityHeaders: Type.Optional(Type.Boolean()), + deferredToolsMode: Type.Optional(Type.Literal("kimi")), sessionAffinityFormat: Type.Optional( Type.Union([Type.Literal("openai"), Type.Literal("openai-nosession"), Type.Literal("openrouter")]), ),