diff --git a/src/extensions/orchestration/continuation-nudge.test.ts b/src/extensions/orchestration/continuation-nudge.test.ts index 34ee913171..4c807e93d0 100644 --- a/src/extensions/orchestration/continuation-nudge.test.ts +++ b/src/extensions/orchestration/continuation-nudge.test.ts @@ -258,6 +258,35 @@ describe("ContinuationNudge session-level tool tracking", () => { expect(guard.hasToolBeenCalledThisSession()).toBe(true) }) + it("resetForModelSwitch clears the session-level tool latch", () => { + const guard = new ContinuationNudge() + guard.resetForNewUserInput() + guard.recordToolCall() + guard.resetForNewUserInput() + expect(guard.hasToolBeenCalledThisSession()).toBe(true) + + guard.resetForModelSwitch() + + expect(guard.hasToolBeenCalledThisSession()).toBe(false) + expect(guard.hasToolBeenCalledThisCycle()).toBe(false) + expect(guard.hasToolBeenCalledThisRun()).toBe(false) + // A text-only turn after the model switch is treated like a fresh + // session and is not nudged. + expect(guard.evaluateTurn(textOnlyMessage)).toBe(false) + }) + + it("resetForModelSwitch clears pending nudge response state", () => { + const guard = new ContinuationNudge() + simulateSessionWithPriorToolCall(guard) + guard.evaluateTurn(textOnlyMessage) + expect(guard.isNudgeResponsePending()).toBe(true) + + guard.resetForModelSwitch() + + expect(guard.isNudgeResponsePending()).toBe(false) + expect(guard.isDoneSignalReceived()).toBe(false) + }) + it("does not consume a nudge slot while no tools have been called this session", () => { // Fresh session, no tools yet — repeated text-only turns must not burn // the per-cycle budget, so when a tool is eventually called the full @@ -396,6 +425,17 @@ describe("ContinuationNudge Agent-pending suppression", () => { expect(guard.evaluateTurn(textOnlyMessage)).toBe(false) }) + it("resetForModelSwitch does NOT clear pending delegation count", () => { + const guard = new ContinuationNudge() + simulateSessionWithPriorToolCall(guard) + guard.markDelegationCall() + // User switches models while an Agent result is still in flight. + guard.resetForModelSwitch() + // Delegated agents are independent of the orchestrator model, so the + // pending count must survive the switch to keep the nudge suppressed. + expect(guard.evaluateTurn(textOnlyMessage)).toBe(false) + }) + it("a regular tool call does NOT clear pending delegation count", () => { const guard = new ContinuationNudge() simulateSessionWithPriorToolCall(guard) @@ -547,6 +587,18 @@ describe("EmptyTurnNudge", () => { expect(guard.evaluateTurn(emptyMessage)).toBe(true) }) + it("re-arms after resetForModelSwitch", () => { + const guard = new EmptyTurnNudge() + expect(guard.evaluateTurn(emptyMessage)).toBe(true) + expect(guard.evaluateTurn(emptyMessage)).toBe(true) + expect(guard.evaluateTurn(emptyMessage)).toBe(false) + // A new model gets a fresh empty-turn budget even within the same cycle. + guard.resetForModelSwitch() + expect(guard.evaluateTurn(emptyMessage)).toBe(true) + expect(guard.evaluateTurn(emptyMessage)).toBe(true) + expect(guard.evaluateTurn(emptyMessage)).toBe(false) + }) + it("does not nudge when the user aborted the turn (stopReason: aborted)", () => { const guard = new EmptyTurnNudge() const aborted = { ...emptyMessage, stopReason: "aborted" as const } diff --git a/src/extensions/orchestration/continuation-nudge.ts b/src/extensions/orchestration/continuation-nudge.ts index ecdf0817d9..6484d8defb 100644 --- a/src/extensions/orchestration/continuation-nudge.ts +++ b/src/extensions/orchestration/continuation-nudge.ts @@ -109,6 +109,27 @@ export class ContinuationNudge { // would incorrectly allow the nudge to fire while Agents are still in flight. } + /** + * Reset the session-level tool latch when the active model changes. + * + * A new model has not yet demonstrated tool-calling behaviour in this + * session, so the fresh-session suppression should apply until it makes + * its first tool call. Without this reset, a model that legitimately + * opens with an orientation/clarification text turn (e.g. kimi-k2.7 in + * single-model mode) is immediately nudged because the previous model + * already called tools. + */ + resetForModelSwitch(): void { + this.toolsCalledThisSession = false + this.toolsCalledSinceLastUserInput = false + this.toolsCalledThisAgentRun = false + this.nudgeCountThisCycle = 0 + this.nudgeResponsePending = false + this.accumulatedResponseText = "" + // pendingDelegationCount is intentionally NOT reset here — delegated + // agents are independent of which orchestrator model is active. + } + recordToolCall(): void { this.toolsCalledSinceLastUserInput = true this.toolsCalledThisAgentRun = true @@ -232,6 +253,17 @@ export class EmptyTurnNudge { resetForNewUserInput(): void { this.nudgeCountThisCycle = 0 } + + /** + * Reset the per-cycle empty-turn budget when the active model changes. + * + * A new model has its own empty-turn behaviour, so it should get the + * full per-cycle budget rather than inheriting any budget consumed by + * the previous model. + */ + resetForModelSwitch(): void { + this.nudgeCountThisCycle = 0 + } } export const NUDGE_CUSTOM_TYPE = "nudge" diff --git a/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.test.ts b/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.test.ts new file mode 100644 index 0000000000..86fba3916d --- /dev/null +++ b/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.test.ts @@ -0,0 +1,113 @@ +import type { ContextEvent } from "@earendil-works/pi-coding-agent" +import { describe, expect, it } from "vitest" +import { isKimiK2Model, normalizeKimiToolCallIds } from "./normalize-kimi-tool-call-ids.js" + +type ContextMessage = ContextEvent["messages"][number] + +function assistantWithToolCall(id: string, name: string): ContextMessage { + return { + role: "assistant", + content: [{ type: "toolCall", id, name, arguments: {} }], + } as unknown as ContextMessage +} + +function toolResult(toolCallId: string, toolName: string): ContextMessage { + return { + role: "toolResult", + toolCallId, + toolName, + content: [{ type: "text", text: "ok" }], + } as unknown as ContextMessage +} + +function userMessage(): ContextMessage { + return { + role: "user", + content: [{ type: "text", text: "continue" }], + } as unknown as ContextMessage +} + +describe("isKimiK2Model", () => { + it.each([ + "kimi-k2", + "kimi-k2.5", + "kimi-k2.6", + "kimi-k2.7", + "kimi-k2.7-code", + "kimi-k2-thinking", + ])("matches %s", (id) => { + expect(isKimiK2Model(id)).toBe(true) + }) + + it.each(["kimi-k3", "glm-5.2-fp8", "minimax-m2.7", "deepseek-v4-flash", "gpt-5.4"])("does not match %s", (id) => { + expect(isKimiK2Model(id)).toBe(false) + }) + + it("does not match undefined", () => { + expect(isKimiK2Model(undefined)).toBe(false) + }) +}) + +describe("normalizeKimiToolCallIds", () => { + it("returns the same array reference when there are no tool calls", () => { + const messages = [userMessage()] + expect(normalizeKimiToolCallIds(messages)).toBe(messages) + }) + + it("rewrites a non-canonical tool call id and the paired toolResult id", () => { + const result = normalizeKimiToolCallIds([ + userMessage(), + assistantWithToolCall("call_abc123", "read"), + toolResult("call_abc123", "read"), + ]) + + const assistant = result[1] as { content: { type: string; id: string }[] } + const resultMsg = result[2] as { toolCallId: string } + expect(assistant.content[0].id).toBe("functions.read:0") + expect(resultMsg.toolCallId).toBe("functions.read:0") + }) + + it("increments the index globally across calls, keeping the tool name per call", () => { + const result = normalizeKimiToolCallIds([ + assistantWithToolCall("call_1", "read"), + toolResult("call_1", "read"), + assistantWithToolCall("call_2", "bash"), + toolResult("call_2", "bash"), + ]) + + expect((result[0] as { content: { id: string }[] }).content[0].id).toBe("functions.read:0") + expect((result[1] as { toolCallId: string }).toolCallId).toBe("functions.read:0") + expect((result[2] as { content: { id: string }[] }).content[0].id).toBe("functions.bash:1") + expect((result[3] as { toolCallId: string }).toolCallId).toBe("functions.bash:1") + }) + + it("renumbers already-canonical ids so the sequence stays globally unique", () => { + // k2.7's own history uses canonical ids; rest of history may not. + const result = normalizeKimiToolCallIds([ + assistantWithToolCall("functions.bash:7", "bash"), + toolResult("functions.bash:7", "bash"), + ]) + + expect((result[0] as { content: { id: string }[] }).content[0].id).toBe("functions.bash:0") + expect((result[1] as { toolCallId: string }).toolCallId).toBe("functions.bash:0") + }) + + it("leaves toolResults without a matching tool call untouched", () => { + const result = normalizeKimiToolCallIds([ + assistantWithToolCall("call_1", "read"), + toolResult("call_orphan", "read"), + ]) + + expect((result[1] as { toolCallId: string }).toolCallId).toBe("call_orphan") + }) + + it("passes non-matching messages through by reference and does not mutate the input", () => { + const user = userMessage() + const assistant = assistantWithToolCall("call_1", "read") + const before = JSON.stringify([user, assistant]) + const result = normalizeKimiToolCallIds([user, assistant]) + + expect(result[0]).toBe(user) + expect(JSON.stringify([user, assistant])).toBe(before) + }) +}) diff --git a/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.ts b/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.ts new file mode 100644 index 0000000000..47dd93ce25 --- /dev/null +++ b/src/extensions/prompt-construction/normalize-kimi-tool-call-ids.ts @@ -0,0 +1,86 @@ +/** + * kimi-k2.x historical tool_call ID normalization. + * + * Moonshot trains the K2 family with tool-call IDs in the canonical format + * `functions.{tool_name}:{idx}` (see Kimi-K2 docs/tool_call_guidance.md), and + * k2.x deployments abandon turns (immediate end-of-turn with no content) when + * conversation history carries non-canonical IDs, e.g. OpenAI-style `call_*` + * IDs written by another model during a mid-session model switch. + * Corroborated upstream: MoonshotAI/Kimi-K2#128, vllm-project/vllm#50782. + * Measured for kimchi (getkimchi/kimchi#1063): on a reconstructed 141-message + * / ~75k-token context, kimi-k2.7 acted in only 1/10 trials on canonical + * mismatch and 12/12 after rewriting IDs to the canonical format. + * + * Moonshot's hosted gateway rewrites IDs server-side; LiteLLM (the ai-enabler + * gateway our provider routes through) does not. Until it does, we rewrite + * here for kimi-k2.* targets. + * + * Every assistant toolCall ID is rewritten to `functions.{name}:{idx}` with a + * global incrementing index, and each `toolResult.toolCallId` is remapped with + * the same value so call/result pairing stays intact. Already-canonical IDs + * are renumbered as well: keeping them while rewriting the rest could yield + * duplicate IDs (e.g. a historical `functions.bash:0` colliding with an + * assigned `functions.bash:0`), and a single globally consistent sequence + * matches the training distribution. + * + * Copy-on-write: messages that need no rewriting are returned by reference; if + * the context contains no tool calls at all, the original array is returned. + */ + +import type { ContextEvent } from "@earendil-works/pi-coding-agent" + +type ContextMessages = ContextEvent["messages"] + +const KIMI_K2_ID = /kimi-k2(?:[.-]|$)/i + +/** True for kimi-k2.x model ids (kimi-k2.5/2.6/2.7, kimi-k2-thinking, ...). */ +export function isKimiK2Model(modelId: string | undefined): boolean { + return !!modelId && KIMI_K2_ID.test(modelId) +} + +type ToolCallBlockLite = { type?: string; id?: unknown; name?: unknown } +type MessageLite = { role?: string; content?: unknown; toolCallId?: unknown } + +export function normalizeKimiToolCallIds(messages: ContextMessages): ContextMessages { + // Pass 1: assign a canonical ID to every assistant toolCall block, in + // context order. + const renames = new Map() + let nextIdx = 0 + for (const msg of messages) { + const lite = msg as MessageLite + if (lite.role !== "assistant" || !Array.isArray(lite.content)) continue + for (const block of lite.content as ToolCallBlockLite[]) { + if (block?.type !== "toolCall" || typeof block.id !== "string" || block.id.length === 0) continue + const name = typeof block.name === "string" && block.name.length > 0 ? block.name : "tool" + renames.set(block.id, `functions.${name}:${nextIdx++}`) + } + } + if (renames.size === 0) return messages + + // Pass 2: rewrite assistant toolCall IDs and the paired toolResult + // toolCallIds with the same mapping. + return messages.map((msg) => { + const lite = msg as MessageLite + if (lite.role === "assistant" && Array.isArray(lite.content)) { + let changed = false + const content = (lite.content as ToolCallBlockLite[]).map((block) => { + if (block?.type === "toolCall" && typeof block.id === "string") { + const renamed = renames.get(block.id) + if (renamed && renamed !== block.id) { + changed = true + return { ...block, id: renamed } + } + } + return block + }) + return changed ? ({ ...(msg as object), content } as ContextMessages[number]) : msg + } + if (lite.role === "toolResult" && typeof lite.toolCallId === "string") { + const renamed = renames.get(lite.toolCallId) + if (renamed && renamed !== lite.toolCallId) { + return { ...(msg as object), toolCallId: renamed } as ContextMessages[number] + } + } + return msg + }) +} diff --git a/src/extensions/prompt-construction/prompt-enrichment.test.ts b/src/extensions/prompt-construction/prompt-enrichment.test.ts index 613f49e0ea..956a9bb68c 100644 --- a/src/extensions/prompt-construction/prompt-enrichment.test.ts +++ b/src/extensions/prompt-construction/prompt-enrichment.test.ts @@ -1019,4 +1019,83 @@ describe("continuation nudge turn_end handler", () => { expect(sendMessageCalls.length).toBe(1) expect((sendMessageCalls[0].message as { content?: string }).content).toContain("If you have finished") }) + + it("does not nudge after a model switch when the previous model called tools", async () => { + const { fire, sendMessageCalls } = buildNudgeHandlers() + + // Previous model called a tool during the session. + await fire("tool_execution_start", {}) + + // User switches models (e.g. via the UI model picker). + await fire("model_select", { + model: { id: "kimi-k2.7", provider: "kimchi-dev" }, + previousModel: undefined, + source: "set", + }) + + // New user input after the switch. + await fire("input", { source: "user" }) + + // New model responds with orientation text only, no tool calls. + await fire("turn_end", { + message: makeAssistantWithStop([{ type: "text", text: "I'll review the branch in detail." }]), + }) + + // No nudge should fire — the model switch reset the session-level + // tool latch, so the new model's orientation turn is treated like a + // fresh session. + expect(sendMessageCalls.length).toBe(0) + }) + + it("does not nudge after a model cycle when the previous model called tools", async () => { + const { fire, sendMessageCalls } = buildNudgeHandlers() + + // Previous model called a tool during the session. + await fire("tool_execution_start", {}) + + // User cycles models (e.g. via the keyboard shortcut). + await fire("model_select", { + model: { id: "kimi-k2.7", provider: "kimchi-dev" }, + previousModel: undefined, + source: "cycle", + }) + + // New user input after the cycle. + await fire("input", { source: "user" }) + + // New model responds with orientation text only, no tool calls. + await fire("turn_end", { + message: makeAssistantWithStop([{ type: "text", text: "I'll review the branch in detail." }]), + }) + + // Cycling is a user-initiated switch and must also reset the latch. + expect(sendMessageCalls.length).toBe(0) + }) + + it("still nudges after a model restore when the previous model called tools", async () => { + const { fire, sendMessageCalls } = buildNudgeHandlers() + + // Previous model called a tool during the session. + await fire("tool_execution_start", {}) + + // Session restore is not a user-initiated switch; the conversation + // continues and the session-level tool latch must stay true. + await fire("model_select", { + model: { id: "kimi-k2.7", provider: "kimchi-dev" }, + previousModel: undefined, + source: "restore", + }) + + // New user input after restore. + await fire("input", { source: "user" }) + + // Model responds with text only, no tool calls. + await fire("turn_end", { + message: makeAssistantWithStop([{ type: "text", text: "I'll review the branch in detail." }]), + }) + + // Nudge should fire because restore must not reset the latch. + expect(sendMessageCalls.length).toBe(1) + expect((sendMessageCalls[0].message as { customType?: string }).customType).toBe("nudge") + }) }) diff --git a/src/extensions/prompt-construction/prompt-enrichment.ts b/src/extensions/prompt-construction/prompt-enrichment.ts index 66708da594..e35055acf7 100644 --- a/src/extensions/prompt-construction/prompt-enrichment.ts +++ b/src/extensions/prompt-construction/prompt-enrichment.ts @@ -72,6 +72,7 @@ import { } from "../orchestration/model-roles.js" import { registerModelRolesCommand } from "../orchestration/model-roles-command.js" import { type ContextFile, loadGlobalContextFiles, loadProjectContextFiles } from "./context-files.js" +import { isKimiK2Model, normalizeKimiToolCallIds } from "./normalize-kimi-tool-call-ids.js" import { buildSystemPrompt, DELEGATION_TOOL_NAMES, @@ -290,8 +291,23 @@ export default function (skillPaths: string[]) { } }) - pi.on("model_select", async (_event, ctx) => { + pi.on("model_select", async (event, ctx) => { notifyIfDeprecated(ctx) + + // A user-initiated model switch (UI picker, /model, or cycling) + // is a fresh start for tool-calling behaviour from the new model's + // perspective. Reset the session-level latch so the first text-only + // turn after the switch is treated like the first turn of a new + // session (nudge suppressed until the new model calls a tool). + // Session restore is deliberately excluded: a restored session is + // continuing an existing conversation, not starting fresh. + if (event.source === "set" || event.source === "cycle") { + const sessionId = ctx.sessionManager.getSessionId() + const continuationNudge = getContinuationNudge(sessionId) + const emptyTurnNudge = getEmptyTurnNudge(sessionId) + continuationNudge.resetForModelSwitch() + emptyTurnNudge.resetForModelSwitch() + } }) // Detect the inverse of the context-event nudge below: the orchestrator reasons @@ -430,10 +446,16 @@ export default function (skillPaths: string[]) { ) }) - pi.on("context", async (event) => { + pi.on("context", async (event, ctx) => { let messages = stripStaleNudges(event.messages) messages = stripEmptyToolCalls(messages) messages = stripUiOnlyMessages(messages) + // kimi-k2.x stalls on historical tool calls whose IDs are not in + // Moonshot's canonical format (issue #1063) — normalize for those + // targets only. + if (isKimiK2Model(ctx.model?.id)) { + messages = normalizeKimiToolCallIds(messages) + } if (messages !== event.messages) return { messages } }) } @@ -444,8 +466,11 @@ export default function (skillPaths: string[]) { // and MiniMax M2.7) emit empty tool calls after a real write/edit call, // which the runtime rejects with a "Tool not found" result that would // otherwise accumulate in the subagent's context across turns. - pi.on("context", async (event) => { - const messages = stripEmptyToolCalls(event.messages) + pi.on("context", async (event, ctx) => { + let messages = stripEmptyToolCalls(event.messages) + if (isKimiK2Model(ctx.model?.id)) { + messages = normalizeKimiToolCallIds(messages) + } if (messages !== event.messages) return { messages } }) }