Fix dangling agent XML cleanup and escape tool call prompt XML

This commit is contained in:
CJACK.
2026-03-30 11:23:16 +08:00
parent c95bf7b667
commit b790545d82
4 changed files with 53 additions and 2 deletions

View File

@@ -23,6 +23,7 @@ var leakedAgentXMLBlockPatterns = []*regexp.Regexp{
regexp.MustCompile(`(?is)<new_task\b[^>]*>(.*?)</new_task>`),
}
var leakedAgentWrapperTagPattern = regexp.MustCompile(`(?is)</?(?:attempt_completion|ask_followup_question|new_task)\b[^>]*>`)
var leakedAgentResultTagPattern = regexp.MustCompile(`(?is)</?result>`)
func sanitizeLeakedOutput(text string) string {
@@ -50,5 +51,13 @@ func sanitizeLeakedAgentXMLBlocks(text string) string {
return leakedAgentResultTagPattern.ReplaceAllString(submatches[1], "")
})
}
// Fallback for truncated output streams: strip any dangling wrapper tags
// that were not part of a complete block replacement. If we detect leaked
// wrapper tags, also strip sibling <result> tags to avoid exposing agent
// markup in user-visible text.
if leakedAgentWrapperTagPattern.MatchString(out) {
out = leakedAgentWrapperTagPattern.ReplaceAllString(out, "")
out = leakedAgentResultTagPattern.ReplaceAllString(out, "")
}
return out
}

View File

@@ -41,3 +41,19 @@ func TestSanitizeLeakedOutputPreservesStandaloneResultTags(t *testing.T) {
t.Fatalf("unexpected sanitize result for standalone result tag: %q", got)
}
}
func TestSanitizeLeakedOutputRemovesDanglingAgentXMLOpeningTags(t *testing.T) {
raw := "Done.<attempt_completion><result>Some final answer"
got := sanitizeLeakedOutput(raw)
if got != "Done.Some final answer" {
t.Fatalf("unexpected sanitize result for dangling opening tags: %q", got)
}
}
func TestSanitizeLeakedOutputRemovesDanglingAgentXMLClosingTags(t *testing.T) {
raw := "Done.Some final answer</result></attempt_completion>"
got := sanitizeLeakedOutput(raw)
if got != "Done.Some final answer" {
t.Fatalf("unexpected sanitize result for dangling closing tags: %q", got)
}
}

View File

@@ -5,6 +5,12 @@ import (
"strings"
)
var promptXMLTextEscaper = strings.NewReplacer(
"&", "&amp;",
"<", "&lt;",
">", "&gt;",
)
// FormatToolCallsForPrompt renders a tool_calls slice into the canonical
// prompt-visible history block used across adapters.
func FormatToolCallsForPrompt(raw any) string {
@@ -82,8 +88,8 @@ func formatToolCallForPrompt(call map[string]any) string {
}
return " <tool_call>\n" +
" <tool_name>" + name + "</tool_name>\n" +
" <parameters>" + StringifyToolCallArguments(argsRaw) + "</parameters>\n" +
" <tool_name>" + escapeXMLText(name) + "</tool_name>\n" +
" <parameters>" + escapeXMLText(StringifyToolCallArguments(argsRaw)) + "</parameters>\n" +
" </tool_call>"
}
@@ -122,3 +128,10 @@ func asString(v any) string {
}
return ""
}
func escapeXMLText(v string) string {
if v == "" {
return ""
}
return promptXMLTextEscaper.Replace(v)
}

View File

@@ -26,3 +26,16 @@ func TestFormatToolCallsForPromptXML(t *testing.T) {
t.Fatalf("unexpected formatted tool call XML: %q", got)
}
}
func TestFormatToolCallsForPromptEscapesXMLEntities(t *testing.T) {
got := FormatToolCallsForPrompt([]any{
map[string]any{
"name": "search<&>",
"arguments": `{"q":"a < b && c > d"}`,
},
})
want := "<tool_calls>\n <tool_call>\n <tool_name>search&lt;&amp;&gt;</tool_name>\n <parameters>{\"q\":\"a &lt; b &amp;&amp; c &gt; d\"}</parameters>\n </tool_call>\n</tool_calls>"
if got != want {
t.Fatalf("unexpected escaped tool call XML: %q", got)
}
}