From 0652d87aca341ae8b5ccacdad787d92b046d787f Mon Sep 17 00:00:00 2001
From: kami <kamixtrash@gmail.com>
Date: Tue, 30 Jun 2026 20:29:28 +0400
Subject: [PATCH] fix(inference): salvage truncated/malformed tool-call JSON
 from content

Small local models emit tool calls into message content but truncate them
(unterminated arguments string, missing closing braces), so strict decode
fails and the blob is treated as a failing artifact, exhausting the stage's
retry budget. Add a lenient fallback that recovers the function name and the
first brace-balanced arguments object, gated on a tool-call marker so genuine
artifact JSON is never misread as a call.
---
 .../llama/cpp/LlamaCppInferenceProvider.kt    | 48 ++++++++++++++++++-
 .../llama/cpp/SalvageToolCallsTest.kt         | 10 ++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt
index 7518b70d..68d40cac 100644
--- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt
+++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt
@@ -11,6 +11,7 @@ import com.correx.core.inference.ProviderHealth
 import com.correx.core.inference.ResponseFormat
 import com.correx.core.inference.TokenUsage
 import com.correx.core.inference.Tokenizer
+import com.correx.core.inference.ToolCallFunction
 import com.correx.core.inference.ToolCallRequest
 import com.correx.infrastructure.inference.commons.ModelDescriptor
 import io.ktor.client.HttpClient
@@ -69,10 +70,55 @@ private val log = LoggerFactory.getLogger(LlamaCppInferenceProvider::class.java)
 internal fun salvageToolCalls(content: String?): List<ToolCallRequest> {
     val s = content?.trim().orEmpty()
     if (s.isEmpty() || (s.first() != '{' && s.first() != '[')) return emptyList()
-    return runCatching { json.decodeFromString<List<ToolCallRequest>>(s) }
+    val strict = runCatching { json.decodeFromString<List<ToolCallRequest>>(s) }
         .recoverCatching { listOf(json.decodeFromString<ToolCallRequest>(s)) }
         .getOrDefault(emptyList())
         .filter { it.function.name.isNotBlank() }
+    return strict.ifEmpty { salvageLenient(s) }
+}
+
+private val NAME_RE = Regex("\"name\"\\s*:\\s*\"([^\"]+)\"")
+
+// Small local models often emit a tool call into `content` but truncate it — an unterminated
+// arguments string and missing closing braces — so strict JSON decode fails. Only the `function`
+// shape (name + arguments) matters to the orchestrator; recover those by hand. Gated on a tool-call
+// marker so genuine artifact JSON (no `function`/`arguments` field) is never misread as a call.
+private fun salvageLenient(s: String): List<ToolCallRequest> {
+    if (!s.contains("\"function\"") && !s.contains("\"arguments\"")) return emptyList()
+    val name = NAME_RE.find(s)?.groupValues?.get(1) ?: return emptyList()
+    val argsIdx = s.indexOf("\"arguments\"")
+    val args = if (argsIdx >= 0) {
+        // The arguments value is a JSON string holding escaped JSON; unescape, then take the
+        // first balanced object so trailing truncation garbage is dropped.
+        val unescaped = s.substring(argsIdx).replace("\\\"", "\"").replace("\\\\", "\\")
+        firstBalancedObject(unescaped) ?: "{}"
+    } else {
+        "{}"
+    }
+    return listOf(ToolCallRequest(function = ToolCallFunction(name = name, arguments = args)))
+}
+
+// Returns the first brace-balanced JSON object substring, or null if none closes (true truncation).
+private fun firstBalancedObject(s: String): String? {
+    val start = s.indexOf('{')
+    if (start < 0) return null
+    var depth = 0
+    var inStr = false
+    var esc = false
+    for (i in start until s.length) {
+        val c = s[i]
+        when {
+            esc -> esc = false
+            c == '\\' -> esc = true
+            c == '"' -> inStr = !inStr
+            !inStr && c == '{' -> depth++
+            !inStr && c == '}' -> {
+                depth--
+                if (depth == 0) return s.substring(start, i + 1)
+            }
+        }
+    }
+    return null
 }
 
 @Suppress("TooGenericExceptionCaught", "MagicNumber")
diff --git a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/SalvageToolCallsTest.kt b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/SalvageToolCallsTest.kt
index c227deb8..a323c937 100644
--- a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/SalvageToolCallsTest.kt
+++ b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/SalvageToolCallsTest.kt
@@ -23,6 +23,16 @@ class SalvageToolCallsTest {
         assertEquals(listOf("file_read", "ls"), calls.map { it.function.name })
     }
 
+    @Test
+    fun `recovers a truncated tool call with unterminated arguments string`() {
+        // Exact shape from session 390983ae: missing closing quote on arguments + trailing braces.
+        val content = """{"id":"rcu08","function":{"name":"file_read","arguments":"{\"path\":\"frontend/src/\"}}"""
+        val calls = salvageToolCalls(content)
+        assertEquals(1, calls.size)
+        assertEquals("file_read", calls.first().function.name)
+        assertEquals("""{"path":"frontend/src/"}""", calls.first().function.arguments)
+    }
+
     @Test
     fun `leaves a real artifact JSON alone`() {
         val artifact = """{"summary":"add a web ui","requirements":["one","two"]}"""