parsers: robust xml tool repair (#14961)

Previous xml repair for glm was a good start, but we need to go further and repair any incorrect open or closing tags Co-authored-by: Dongluo Chen <dongluo.chen@gmail.com>
2026-03-27 02:58:43 +07:00 · 2026-03-19 11:24:48 -07:00
parent 3f3a24b418
commit 126d8db7f3
2 changed files with 214 additions and 43 deletions
--- a/model/parsers/glm46.go
+++ b/model/parsers/glm46.go
@@ -345,44 +345,163 @@ func escapeGLM46Content(s string) string {
 	return result.String()
 }

-// repairUnclosedArgValues inserts missing </arg_value> closing tags.
-// GLM models sometimes omit the closing tag, producing XML like:
+// repairPhase represents the expected next tag in the repair cycle.
+type repairPhase int
+
+const (
+	phaseArgKeyOpen  repairPhase = iota // expecting <arg_key>
+	phaseArgKeyClose                    // expecting </arg_key>
+	phaseArgValOpen                     // expecting <arg_value>
+	phaseArgValClose                    // expecting </arg_value>
+	phaseCount                          // number of phases
+)
+
+// repairGLM46XML reconstructs well-formed XML from GLM model output that may
+// have missing or mismatched tags. The expected structure is:
 //
-//	<arg_value>value</tool_call>
+//	func_name
+//	<arg_key>key</arg_key>
+//	<arg_value>value</arg_value>
+//	...
 //
-// instead of:
-//
-//	<arg_value>value</arg_value></tool_call>
-func repairUnclosedArgValues(s string) string {
+// GLM models frequently omit opening or closing tags. This function follows
+// the expected tag cycle, scanning forward for each expected tag in sequence.
+// When a tag is missing, it inserts the tag and consumes any text in between.
+func repairGLM46XML(s string) string {
+	// tagCycle is the repeating sequence of tags after the function name.
+	tagCycle := [phaseCount]string{"<arg_key>", "</arg_key>", "<arg_value>", "</arg_value>"}
+
+	// findNextTag returns the index and identity of the earliest known tag in s.
+	findNextTag := func(s string) (int, string) {
+		bestIdx := -1
+		bestTag := ""
+		for _, tag := range tagCycle {
+			if idx := strings.Index(s, tag); idx != -1 && (bestIdx == -1 || idx < bestIdx) {
+				bestIdx = idx
+				bestTag = tag
+			}
+		}
+		return bestIdx, bestTag
+	}
+
+	// tagIndex returns the phase corresponding to the given tag.
+	tagIndex := func(tag string) repairPhase {
+		for i, t := range tagCycle {
+			if t == tag {
+				return repairPhase(i)
+			}
+		}
+		return -1
+	}
+
 	var result strings.Builder
-	for {
-		openIdx := strings.Index(s, "<arg_value>")
-		if openIdx == -1 {
+
+	idx, firstTag := findNextTag(s)
+	if idx == -1 {
+		return s
+	}
+	prefix := s[:idx]
+	s = s[idx:]
+
+	// If the first tag is not <arg_key>, the text before it may contain both
+	// the function name and key content (e.g. "weather city</arg_key>").
+	// Function names cannot contain space, so split at the first space.
+	phase := phaseArgKeyOpen
+	if firstTag != "<arg_key>" {
+		if spIdx := strings.IndexFunc(prefix, unicode.IsSpace); spIdx != -1 {
+			result.WriteString(prefix[:spIdx])
+			keyContent := strings.TrimLeftFunc(prefix[spIdx:], unicode.IsSpace)
+			result.WriteString("<arg_key>")
+			result.WriteString(keyContent)
+			phase = phaseArgKeyClose
+		} else {
+			result.WriteString(prefix)
+		}
+	} else {
+		result.WriteString(prefix)
+	}
+
+	// Walk through the expected tag cycle. At each step, look for the
+	// expected tag. If a different tag appears first, emit the missing
+	// tags to catch up, then continue.
+	for len(s) > 0 {
+		idx, found := findNextTag(s)
+		expected := tagCycle[phase]
+		isOpen := phase%2 == 0 // even phases are opening tags
+
+		if idx == -1 {
+			// No more tags — emit remaining text with fixups
+			if isOpen {
+				// Expecting an opening tag but nothing left — we're done
+				break
+			}
+			// Expecting a closing tag — emit text then close
 			result.WriteString(s)
+			result.WriteString(expected)
+			phase = (phase + 1) % phaseCount
 			break
 		}
-		afterOpen := openIdx + len("<arg_value>")
-		closeIdx := strings.Index(s[afterOpen:], "</arg_value>")
-		nextKeyIdx := strings.Index(s[afterOpen:], "<arg_key>")
-		// Check if properly closed before the next <arg_key> (or no next key)
-		if closeIdx != -1 && (nextKeyIdx == -1 || closeIdx < nextKeyIdx) {
-			end := afterOpen + closeIdx + len("</arg_value>")
-			result.WriteString(s[:end])
-			s = s[end:]
+
+		if found == expected {
+			// Found the expected tag — emit any text before it, then the tag
+			result.WriteString(s[:idx])
+			result.WriteString(expected)
+			s = s[idx+len(expected):]
+			phase = (phase + 1) % phaseCount
 			continue
 		}
-		// Unclosed — insert </arg_value> before the next <arg_key> or at end
-		if nextKeyIdx != -1 {
-			insertAt := afterOpen + nextKeyIdx
-			result.WriteString(s[:insertAt])
-			result.WriteString("</arg_value>")
-			s = s[insertAt:]
-		} else {
-			result.WriteString(s)
-			result.WriteString("</arg_value>")
-			break
+
+		// Found a different tag. Insert missing tags to catch up.
+		foundIdx := tagIndex(found)
+
+		if isOpen && idx > 0 {
+			// Text before the found tag while expecting an opening tag —
+			// the opening tag was omitted. Emit it before the text.
+			result.WriteString(expected)
+			// Advance to the next phase (text content) and then look
+			// for the closing tag — but the found tag might be that
+			// closing tag or something further ahead. Emit text up to
+			// the found tag and insert any missing tags between.
+			result.WriteString(s[:idx])
+			phase = (phase + 1) % phaseCount // now expecting closing
+			s = s[idx:]
+			// Fall through to re-evaluate with the closing tag expected
+			continue
 		}
+
+		// Emit missing tags to advance from current phase to the found tag's phase
+		for phase != foundIdx {
+			tag := tagCycle[phase]
+			if phase%2 == 0 {
+				result.WriteString(tag)
+			} else {
+				// Closing tag — emit any text before the found tag first,
+				// but only if we're one step before the found tag
+				if (phase+1)%phaseCount == foundIdx && idx > 0 {
+					result.WriteString(s[:idx])
+					s = s[idx:]
+					idx = 0
+				}
+				result.WriteString(tag)
+			}
+			phase = (phase + 1) % phaseCount
+		}
+		// Now phase == foundIdx, re-process without advancing s
 	}
+
+	// If we stopped mid-pair (after an opening tag), close it
+	switch phase {
+	case phaseArgKeyClose: // after <arg_key>, expecting text/</arg_key>
+		result.WriteString("</arg_key>")
+		result.WriteString("<arg_value>")
+		result.WriteString("</arg_value>")
+	case phaseArgValOpen: // after </arg_key>, expecting <arg_value>
+		result.WriteString("<arg_value>")
+		result.WriteString("</arg_value>")
+	case phaseArgValClose: // after <arg_value>, expecting text/</arg_value>
+		result.WriteString("</arg_value>")
+	}
+
 	return result.String()
 }

@@ -398,7 +517,7 @@ func parseGLM46ToolCall(raw glm46EventRawToolCall, tools []api.Tool) (api.ToolCa
 	var parsed GLMToolCallXML
 	if err := xml.Unmarshal([]byte(xmlString), &parsed); err != nil {
 		parsed = GLMToolCallXML{}
-		repaired := "<tool_call>" + repairUnclosedArgValues(escaped) + "</tool_call>"
+		repaired := "<tool_call>" + repairGLM46XML(escaped) + "</tool_call>"
 		if err2 := xml.Unmarshal([]byte(repaired), &parsed); err2 != nil {
 			return api.ToolCall{}, fmt.Errorf("failed to parse XML: %w", err)
 		}
--- a/model/parsers/glm46_test.go
+++ b/model/parsers/glm46_test.go
@@ -887,6 +887,28 @@ line3</arg_value>`,
 				},
 			},
 		},
+		{
+			name:        "unopened arg_value after arg_key",
+			tools:       []api.Tool{},
+			rawToolCall: "get-weather\n<arg_key>city</arg_key>\nNew York</arg_value>\n<arg_key>unit</arg_key>\ncelsius</arg_value>",
+			wantToolCall: api.ToolCall{
+				Function: api.ToolCallFunction{
+					Name:      "get-weather",
+					Arguments: args(`{"city": "New York", "unit": "celsius"}`),
+				},
+			},
+		},
+		{
+			name:        "mixed unopened and valid arg_values",
+			tools:       []api.Tool{},
+			rawToolCall: "get-weather\n<arg_key>city</arg_key>\n<arg_value>Paris</arg_value>\n<arg_key>unit</arg_key>\ncelsius</arg_value>",
+			wantToolCall: api.ToolCall{
+				Function: api.ToolCallFunction{
+					Name:      "get-weather",
+					Arguments: args(`{"city": "Paris", "unit": "celsius"}`),
+				},
+			},
+		},
 	}

 	for i, tc := range cases {
@@ -902,7 +924,7 @@ line3</arg_value>`,
 	}
 }

-func TestRepairUnclosedArgValues(t *testing.T) {
+func TestRepairGLM46XML(t *testing.T) {
 	cases := []struct {
 		name  string
 		input string
@@ -910,33 +932,63 @@ func TestRepairUnclosedArgValues(t *testing.T) {
 	}{
 		{
 			name:  "already valid",
-			input: `<arg_key>k</arg_key><arg_value>v</arg_value>`,
-			want:  `<arg_key>k</arg_key><arg_value>v</arg_value>`,
+			input: `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
 		},
 		{
-			name:  "unclosed at end",
-			input: `<arg_key>k</arg_key><arg_value>v`,
-			want:  `<arg_key>k</arg_key><arg_value>v</arg_value>`,
+			name:  "missing </arg_value> at end",
+			input: `func<arg_key>k</arg_key><arg_value>v`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
 		},
 		{
-			name:  "unclosed before next arg_key",
-			input: `<arg_key>a</arg_key><arg_value>1<arg_key>b</arg_key><arg_value>2</arg_value>`,
-			want:  `<arg_key>a</arg_key><arg_value>1</arg_value><arg_key>b</arg_key><arg_value>2</arg_value>`,
+			name:  "missing </arg_value> before next arg_key",
+			input: `func<arg_key>a</arg_key><arg_value>1<arg_key>b</arg_key><arg_value>2</arg_value>`,
+			want:  `func<arg_key>a</arg_key><arg_value>1</arg_value><arg_key>b</arg_key><arg_value>2</arg_value>`,
 		},
 		{
-			name:  "no arg_value tags",
+			name:  "no tags at all",
 			input: `just plain text`,
 			want:  `just plain text`,
 		},
 		{
-			name:  "multiple unclosed",
-			input: `<arg_key>a</arg_key><arg_value>1<arg_key>b</arg_key><arg_value>2`,
-			want:  `<arg_key>a</arg_key><arg_value>1</arg_value><arg_key>b</arg_key><arg_value>2</arg_value>`,
+			name:  "missing <arg_value> open tag",
+			input: `func<arg_key>k</arg_key>v</arg_value>`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
+		},
+		{
+			name:  "missing </arg_key> close tag",
+			input: `func<arg_key>k<arg_value>v</arg_value>`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
+		},
+		{
+			name:  "missing <arg_key> open tag",
+			input: `func k</arg_key><arg_value>v</arg_value>`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
+		},
+		{
+			name:  "all closing tags missing",
+			input: `func<arg_key>k<arg_value>v`,
+			want:  `func<arg_key>k</arg_key><arg_value>v</arg_value>`,
+		},
+		{
+			name:  "all opening tags missing",
+			input: "func k</arg_key>v</arg_value>",
+			want:  "func<arg_key>k</arg_key><arg_value>v</arg_value>",
+		},
+		{
+			name:  "multiple pairs with mixed missing tags",
+			input: `func<arg_key>a</arg_key>1</arg_value><arg_key>b<arg_value>2</arg_value>`,
+			want:  `func<arg_key>a</arg_key><arg_value>1</arg_value><arg_key>b</arg_key><arg_value>2</arg_value>`,
+		},
+		{
+			name:  "newlines preserved",
+			input: "func\n<arg_key>city</arg_key>\nNew York</arg_value>",
+			want:  "func\n<arg_key>city</arg_key><arg_value>\nNew York</arg_value>",
 		},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got := repairUnclosedArgValues(tc.input)
+			got := repairGLM46XML(tc.input)
 			if got != tc.want {
 				t.Errorf("got %q, want %q", got, tc.want)
 			}