pi-tool-repair 0.1.3 → 0.1.5

package/README.md

@@ -168,7 +168,7 @@ Safety gates:
 
 - `requireKnownTool: true` only recovers calls whose name is in pi's active tool registry.
 - Markup inside fenced code blocks is ignored so syntax discussions and examples are preserved.
-- Incomplete or unparseable blocks are left alone.
+- Incomplete or unparseable blocks are not recovered as tool calls. For DSML, dangling or truncated marker tokens (e.g. a stream that died at `<｜DSML｜tool_calls` with no closing `>`) are still stripped from visible text so the raw marker doesn't persist in the transcript.
 - If the provider already emitted native `toolCall` blocks, leaked shadow text is stripped but duplicate calls are not added.
 
 Covered grammar families: DeepSeek DSML, MiniMax/Anthropic `<invoke>`, Qwen/Hermes `<tool_call>`, Kimi sentinels, Mistral `[TOOL_CALLS]`, Llama `<|python_tag|>`, GLM `arg_key`/`arg_value`, Granite JSON `<tool_call>`, MiniMax-Text-01 TypeScript calls, and OLMo3 `<function_calls>` pythonic calls. See [`docs/tool-call-grammar-leakage-survey.md`](./docs/tool-call-grammar-leakage-survey.md) for the survey.

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "pi-tool-repair",
-  "version": "0.1.3",
-  "description": "Validate-then-repair extension for pi \u2014 fixes common LLM tool-call mistakes (null fields, stringified arrays, wrong field names, anchor bleed) before tools execute",
+  "version": "0.1.5",
+  "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",
   "license": "MIT",
@@ -33,13 +33,6 @@
     "docs/",
     "README.md"
   ],
-  "scripts": {
-    "test": "vitest run",
-    "test:watch": "vitest",
-    "test:coverage": "vitest run --coverage",
-    "typecheck": "tsc --noEmit",
-    "lint:dead": "knip --no-gitignore"
-  },
   "devDependencies": {
     "@earendil-works/pi-coding-agent": "0.79.4",
     "@types/node": "25.9.1",
@@ -52,5 +45,12 @@
     "extensions": [
       "./tool-repair.ts"
     ]
+  },
+  "scripts": {
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "test:coverage": "vitest run --coverage",
+    "typecheck": "tsc --noEmit",
+    "lint:dead": "knip --no-gitignore"
   }
-}
+}

package/src/grammar-repair.ts

@@ -38,6 +38,7 @@ export interface RecoveredToolCall {
 
 interface Candidate extends RecoveredToolCall {
   range: Range;
+  stripOnly?: boolean;
 }
 
 interface Range {
@@ -147,13 +148,14 @@ export function repairAssistantMessageGrammarLeaks(
     if (text === undefined) return part;
 
     const candidates = selectCandidates(parseToolGrammarCandidates(text, enabled))
-      .filter((candidate) => isAllowedTool(candidate.name, config, knownTools));
+      .filter((candidate) => candidate.stripOnly || isAllowedTool(candidate.name, config, knownTools));
 
     if (candidates.length === 0) return part;
 
     strippedRanges += candidates.length;
     changed = true;
     for (const candidate of candidates) {
+      if (candidate.stripOnly) continue;
       recoveredCalls.push({
         name: candidate.name,
         arguments: candidate.arguments,
@@ -201,16 +203,21 @@ export function repairAssistantMessageGrammarLeaks(
 
 export function parseToolGrammarLeaks(text: string, grammars: Iterable<GrammarName> = ALL_GRAMMARS): RecoveredToolCall[] {
   const enabled = new Set(grammars);
-  return selectCandidates(parseToolGrammarCandidates(text, enabled)).map((candidate) => ({
-    name: candidate.name,
-    arguments: candidate.arguments,
-    grammar: candidate.grammar,
-  }));
+  return selectCandidates(parseToolGrammarCandidates(text, enabled))
+    .filter((candidate) => !candidate.stripOnly)
+    .map((candidate) => ({
+      name: candidate.name,
+      arguments: candidate.arguments,
+      grammar: candidate.grammar,
+    }));
 }
 
 function parseToolGrammarCandidates(text: string, enabled: Set<GrammarName>): Candidate[] {
   const candidates: Candidate[] = [];
-  if (enabled.has("dsml")) candidates.push(...parseDsml(text));
+  if (enabled.has("dsml")) {
+    candidates.push(...parseDsml(text));
+    candidates.push(...parseDsmlDanglingMarkers(text));
+  }
   if (enabled.has("kimi")) candidates.push(...parseKimi(text));
   if (enabled.has("mistral")) {
     candidates.push(...parseMistral(text));
@@ -252,6 +259,28 @@ function parseDsml(text: string): Candidate[] {
   return candidates;
 }
 
+function parseDsmlDanglingMarkers(text: string): Candidate[] {
+  if (!text.includes("DSML")) return [];
+  const prefix = "(?:｜{1,2}DSML｜{1,2}|DSML｜|\\s*\\|\\s*DSML\\s*\\|\\s*)";
+  const markerRe = new RegExp(
+    `</?${prefix}(?:tool_calls|function_calls|invoke|parameter)(?:\\s+[^>\\n]*)?>?`,
+    "giu",
+  );
+  const candidates: Candidate[] = [];
+  for (const match of text.matchAll(markerRe)) {
+    if (match.index === undefined) continue;
+    if (isInsideCodeFence(text, match.index)) continue;
+    candidates.push({
+      name: "",
+      arguments: {},
+      grammar: "dsml",
+      range: { start: match.index, end: match.index + match[0].length },
+      stripOnly: true,
+    });
+  }
+  return candidates;
+}
+
 function findDsmlClose(text: string, from: number, outerName: string): Range | undefined {
   const prefix = "(?:｜{1,2}DSML｜{1,2}|DSML｜|\\s*\\|\\s*DSML\\s*\\|\\s*)";
   const closeRe = new RegExp(`</${prefix}${outerName}>`, "giu");

package/src/index.ts

@@ -131,6 +131,67 @@ export function stripAnchorBleedInPlace(obj: Record<string, unknown>): void {
   }
 }
 
+// Leaked grammar markers from GLM/ChatGLM style tool-call grammars.
+// These can end up as literal prefixes/suffixes on parsed object keys
+// or string values instead of being interpreted as XML tags.
+const GRAMMAR_TOKEN_LEAKS = [
+  { tag: "<arg_key>", at: "start" as const },
+  { tag: "</arg_key>", at: "end" as const },
+  { tag: "<arg_value>", at: "start" as const },
+  { tag: "</arg_value>", at: "end" as const },
+];
+
+export function stripGrammarTokenLeaksInPlace(obj: Record<string, unknown>): void {
+  for (const key of Object.keys(obj)) {
+    const value = obj[key];
+    let newKey = key;
+    for (const { tag, at } of GRAMMAR_TOKEN_LEAKS) {
+      if (at === "start" && newKey.startsWith(tag)) {
+        newKey = newKey.slice(tag.length);
+      } else if (at === "end" && newKey.endsWith(tag)) {
+        newKey = newKey.slice(0, -tag.length);
+      }
+    }
+    newKey = newKey.trim();
+
+    if (newKey !== key) {
+      obj[newKey] = value;
+      delete obj[key];
+    }
+
+    if (typeof value === "string") {
+      let s = value;
+      for (const { tag, at } of GRAMMAR_TOKEN_LEAKS) {
+        if (at === "start" && s.startsWith(tag)) {
+          s = s.slice(tag.length);
+        } else if (at === "end" && s.endsWith(tag)) {
+          s = s.slice(0, -tag.length);
+        }
+      }
+      obj[newKey] = s.trim();
+    } else if (Array.isArray(value)) {
+      for (let i = 0; i < value.length; i++) {
+        const item = value[i];
+        if (typeof item === "string") {
+          let s = item;
+          for (const { tag, at } of GRAMMAR_TOKEN_LEAKS) {
+            if (at === "start" && s.startsWith(tag)) {
+              s = s.slice(tag.length);
+            } else if (at === "end" && s.endsWith(tag)) {
+              s = s.slice(0, -tag.length);
+            }
+          }
+          value[i] = s.trim();
+        } else if (item && typeof item === "object") {
+          stripGrammarTokenLeaksInPlace(item as Record<string, unknown>);
+        }
+      }
+    } else if (value && typeof value === "object") {
+      stripGrammarTokenLeaksInPlace(value as Record<string, unknown>);
+    }
+  }
+}
+
 // ─── Phase 2: Repair Rules ────────────────────────────────────────────────────
 
 export interface RepairResult {

package/tool-repair.ts

@@ -45,6 +45,7 @@ import {
   hasAnchorBleedBug,
   sanitizeSchemaAnchors,
   stripAnchorBleedInPlace,
+  stripGrammarTokenLeaksInPlace,
   wrapRootStringAsObject,
   repairToolInput,
   validateAgainstSchema,
@@ -213,6 +214,12 @@ export default function (pi: ExtensionAPI) {
       }
     }
 
+    // Defense-in-depth: strip leaked grammar tokens (e.g. GLM <arg_key>) from
+    // parsed tool-call keys/values before the schema sees them.
+    if (input && typeof input === "object") {
+      stripGrammarTokenLeaksInPlace(input as Record<string, unknown>);
+    }
+
     // Only repair built-in tools we have schemas for
     const schema = BUILTIN_SCHEMAS[toolName];
     if (!schema) return;