diff --git a/internal/adapter/openai/leaked_output_sanitize.go b/internal/adapter/openai/leaked_output_sanitize.go index 2e7dcc1..481492d 100644 --- a/internal/adapter/openai/leaked_output_sanitize.go +++ b/internal/adapter/openai/leaked_output_sanitize.go @@ -23,6 +23,7 @@ var leakedAgentXMLBlockPatterns = []*regexp.Regexp{ regexp.MustCompile(`(?is)]*>(.*?)`), } +var leakedAgentWrapperTagPattern = regexp.MustCompile(`(?is)]*>`) var leakedAgentResultTagPattern = regexp.MustCompile(`(?is)`) 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 tags to avoid exposing agent + // markup in user-visible text. + if leakedAgentWrapperTagPattern.MatchString(out) { + out = leakedAgentWrapperTagPattern.ReplaceAllString(out, "") + out = leakedAgentResultTagPattern.ReplaceAllString(out, "") + } return out } diff --git a/internal/adapter/openai/leaked_output_sanitize_test.go b/internal/adapter/openai/leaked_output_sanitize_test.go index 90ce9d1..1898289 100644 --- a/internal/adapter/openai/leaked_output_sanitize_test.go +++ b/internal/adapter/openai/leaked_output_sanitize_test.go @@ -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.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" + got := sanitizeLeakedOutput(raw) + if got != "Done.Some final answer" { + t.Fatalf("unexpected sanitize result for dangling closing tags: %q", got) + } +} diff --git a/internal/prompt/tool_calls.go b/internal/prompt/tool_calls.go index 718c48b..d8a2df9 100644 --- a/internal/prompt/tool_calls.go +++ b/internal/prompt/tool_calls.go @@ -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 " \n" + - " " + name + "\n" + - " " + StringifyToolCallArguments(argsRaw) + "\n" + + " " + escapeXMLText(name) + "\n" + + " " + escapeXMLText(StringifyToolCallArguments(argsRaw)) + "\n" + " " } @@ -122,3 +128,10 @@ func asString(v any) string { } return "" } + +func escapeXMLText(v string) string { + if v == "" { + return "" + } + return promptXMLTextEscaper.Replace(v) +} diff --git a/internal/prompt/tool_calls_test.go b/internal/prompt/tool_calls_test.go index 8ad4407..3eb2c1e 100644 --- a/internal/prompt/tool_calls_test.go +++ b/internal/prompt/tool_calls_test.go @@ -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 := "\n \n search<&>\n {\"q\":\"a < b && c > d\"}\n \n" + if got != want { + t.Fatalf("unexpected escaped tool call XML: %q", got) + } +}