pi-tool-repair 0.2.5 → 0.3.0

package/README.md

@@ -195,11 +195,13 @@ With Pi Fabric full code mode, `pi-tool-repair` sees `fabric_exec` as the active
 
 Phase 0 schema sanitization activates for models matching these patterns:
 
-| Pattern      | Models           |
-| ------------ | ---------------- |
-| `/kimi-k2/i` | Kimi K2 variants |
-| `/minimax/i` | MiniMax variants |
-| `/glm/i`     | GLM variants     |
+| Pattern           | Models                              |
+| ----------------- | ----------------------------------- |
+| `/kimi-k2/i`      | Kimi K2 variants                    |
+| `/moonshotai\/kimi/i` | Moonshot Kimi (direct API)      |
+| `/@cf\/.*kimi/i`  | @cf Kimi (e.g. `@cf/moonshotai/kimi-k2.6`) |
+| `/minimax/i`      | MiniMax variants                    |
+| `/glm/i`          | GLM variants                        |
 
 To add more models, edit `anchorBleedModels` in [`src/index.ts`](./src/index.ts).
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-tool-repair",
-  "version": "0.2.5",
+  "version": "0.3.0",
   "description": "Validate-then-repair extension for pi — fixes common LLM tool-call mistakes (null fields, stringified arrays, wrong field names, anchor bleed) before tools execute",
   "type": "module",
   "author": "Tom X Nguyen",

package/src/grammar-repair.ts

@@ -116,6 +116,7 @@ export function defaultConfigPath(): string {
   return join(getAgentDir(), "extensions", "pi-tool-repair.json");
 }
 
+
 export function normalizeGrammarRepairConfig(raw: Partial<GrammarRepairConfig> = {}): GrammarRepairConfig {
   const grammarSet = new Set<GrammarName>(ALL_GRAMMARS);
   const grammars = Array.isArray(raw.grammars)
@@ -168,6 +169,27 @@ export function resolveGrammarRepairForModel(
     : config;
 }
 
+// Open models occasionally append index tokens to tool names (for example
+// `bash_1_234456789` for `bash`). Recover the original name only when the
+// trailing part is made up of separators and digits, and prefer the longest
+// known tool so `read_file_1_2` never collapses to `read`. Names that are
+// already known, or whose suffix contains real characters
+// (`read_multiple_files`), are returned untouched.
+const MANGLED_TOOL_SUFFIX = /^[_-\d]+$/;
+
+function recoverMangledToolName(name: string, knownTools: Set<string>): string {
+  if (knownTools.has(name)) return name;
+
+  let best: string | undefined;
+  for (const tool of knownTools) {
+    if (tool.length >= name.length || !name.startsWith(tool)) continue;
+    if (!MANGLED_TOOL_SUFFIX.test(name.slice(tool.length))) continue;
+    if (best === undefined || tool.length > best.length) best = tool;
+  }
+
+  return best ?? name;
+}
+
 export function repairAssistantMessageGrammarLeaks(
   message: MinimalAssistantMessage,
   config: GrammarRepairConfig,
@@ -188,6 +210,10 @@ export function repairAssistantMessageGrammarLeaks(
     if (text === undefined) return part;
 
     const candidates = selectCandidates(parseToolGrammarCandidates(text, enabled))
+      .map((candidate) => {
+        const name = recoverMangledToolName(candidate.name, knownTools);
+        return name === candidate.name ? candidate : { ...candidate, name };
+      })
       .filter((candidate) => candidate.stripOnly || isAllowedTool(candidate.name, config, knownTools));
 
     if (candidates.length === 0) return part;
@@ -274,6 +300,9 @@ function parseToolGrammarCandidates(text: string, enabled: Set<GrammarName>): Ca
     candidates.push(...parseDsmlDanglingMarkers(text));
   }
   if (enabled.has("kimi")) candidates.push(...parseKimi(text));
+  if (enabled.has("qwen") || enabled.has("granite") || enabled.has("kimi")) {
+    candidates.push(...parseOrphanToolCallClosers(text));
+  }
   if (enabled.has("mistral")) {
     candidates.push(...parseMistral(text));
     candidates.push(...parseBareJsonToolCalls(text, "mistral"));
@@ -420,6 +449,138 @@ function parseKimi(text: string): Candidate[] {
     }
   }
 
+  candidates.push(...parseKimiLooseToolCallXml(text));
+  candidates.push(...parseKimiDanglingMarkers(text));
+  return candidates;
+}
+
+function parseKimiLooseToolCallXml(text: string): Candidate[] {
+  const candidates: Candidate[] = [];
+  const openRe = /<tool_call>/gi;
+
+  for (const match of text.matchAll(openRe)) {
+    if (match.index === undefined || isInsideCodeFence(text, match.index)) continue;
+
+    const start = match.index;
+    const bodyStart = start + match[0].length;
+    if (findPattern(text, /<\/tool_call>/gi, bodyStart)) continue;
+
+    const looseClose = findPattern(
+      text,
+      /<\/(?:redacted_thinking|thinking)>/gi,
+      bodyStart,
+    );
+    if (!looseClose) continue;
+
+    const body = text.slice(bodyStart, looseClose.start);
+    for (const call of parseToolCallJsonBody(body, "kimi")) {
+      candidates.push({ ...call, range: { start, end: looseClose.end } });
+    }
+  }
+
+  return candidates;
+}
+
+function parseKimiDanglingMarkers(text: string): Candidate[] {
+  const candidates: Candidate[] = [];
+  const sectionOpenRe = /<\|(?:redacted_)?tool_calls?_section_begin\|>/gi;
+
+  for (const match of text.matchAll(sectionOpenRe)) {
+    if (match.index === undefined || isInsideCodeFence(text, match.index)) continue;
+    const sectionStart = match.index;
+    if (findPattern(text, /<\|(?:redacted_)?tool_calls?_section_end\|>/gi, sectionStart)) continue;
+
+    const end = findBestKimiUnclosedEnd(text, sectionStart + match[0].length) ?? sectionStart + match[0].length;
+    candidates.push({
+      name: "",
+      arguments: {},
+      grammar: "kimi",
+      range: { start: sectionStart, end },
+      stripOnly: true,
+    });
+  }
+
+  const openRe = /<tool_call>/gi;
+  for (const match of text.matchAll(openRe)) {
+    if (match.index === undefined || isInsideCodeFence(text, match.index)) continue;
+    const start = match.index;
+    const bodyStart = start + match[0].length;
+    const strictClose = findPattern(text, /<\/tool_call>/gi, bodyStart);
+    const looseClose = findPattern(text, /<\/(?:redacted_thinking|thinking)>/gi, bodyStart);
+    const close = strictClose ?? looseClose;
+    if (close && parseToolCallJsonBody(text.slice(bodyStart, close.start), "kimi").length > 0) continue;
+
+    candidates.push({
+      name: "",
+      arguments: {},
+      grammar: "kimi",
+      range: { start, end: close?.end ?? bodyStart },
+      stripOnly: true,
+    });
+  }
+
+  const markerRe = /<\|(?:redacted_)?tool_call(?:s)?(?:_section)?_(?:begin(?:_kimi)?|argument_begin|end(?:_kimi)?)\|>/gi;
+  for (const match of text.matchAll(markerRe)) {
+    if (match.index === undefined || isInsideCodeFence(text, match.index)) continue;
+    const range = { start: match.index, end: match.index + match[0].length };
+    if (candidates.some((candidate) => rangesOverlap(candidate.range, range))) continue;
+    candidates.push({
+      name: "",
+      arguments: {},
+      grammar: "kimi",
+      range,
+      stripOnly: true,
+    });
+  }
+
+  return candidates;
+}
+
+function findBestKimiUnclosedEnd(text: string, from: number): number | undefined {
+  const markerRe = /<\|(?:redacted_)?tool_call(?:s)?(?:_section)?_(?:begin(?:_kimi)?|argument_begin|end(?:_kimi)?)\|>/gi;
+  markerRe.lastIndex = from;
+  let end: number | undefined;
+  for (;;) {
+    const match = markerRe.exec(text);
+    if (!match) break;
+    end = match.index + match[0].length;
+  }
+  if (end === undefined) return undefined;
+
+  const tail = text.slice(end);
+  const idMatch = /^\s*functions\.[A-Za-z_][\w.-]*:\d*/.exec(tail);
+  if (idMatch) end += idMatch[0].length;
+
+  const argMarker = findPattern(text, /<\|(?:redacted_)?tool_call_argument_begin\|>/gi, end);
+  if (argMarker) {
+    end = argMarker.end;
+    const json = extractFirstBalancedJson(text.slice(end));
+    if (json) end += json.end;
+  }
+
+  return end;
+}
+
+function parseOrphanToolCallClosers(text: string): Candidate[] {
+  const candidates: Candidate[] = [];
+  const closeRe = /<\/tool_call>\s*/gi;
+
+  for (const match of text.matchAll(closeRe)) {
+    if (match.index === undefined || isInsideCodeFence(text, match.index)) continue;
+    const before = text.slice(0, match.index);
+    const opens = before.match(/<tool_call>/gi)?.length ?? 0;
+    const closes = before.match(/<\/tool_call>/gi)?.length ?? 0;
+    if (opens > closes) continue;
+
+    candidates.push({
+      name: "",
+      arguments: {},
+      grammar: "kimi",
+      range: { start: match.index, end: match.index + match[0].length },
+      stripOnly: true,
+    });
+  }
+
   return candidates;
 }
 

package/src/index.ts

@@ -24,7 +24,9 @@ export interface RepairConfig {
 export const DEFAULT_CONFIG: RepairConfig = {
   debug: Boolean(process.env.PI_TOOL_REPAIR_DEBUG),
   anchorBleedModels: [
-    /kimi-k2/i,
+    /kimi[-_]?k2/i,
+    /moonshotai\/kimi/i,
+    /@cf\/.*kimi/i,
     /minimax/i,
     /glm/i,
   ],