mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-04 00:15:28 +08:00
Fix dangling agent XML cleanup and escape tool call prompt XML
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,12 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
var promptXMLTextEscaper = strings.NewReplacer(
|
||||
"&", "&",
|
||||
"<", "<",
|
||||
">", ">",
|
||||
)
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
||||
@@ -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<&></tool_name>\n <parameters>{\"q\":\"a < b && c > d\"}</parameters>\n </tool_call>\n</tool_calls>"
|
||||
if got != want {
|
||||
t.Fatalf("unexpected escaped tool call XML: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user