diff --git a/internal/adapter/openai/leaked_output_sanitize.go b/internal/adapter/openai/leaked_output_sanitize.go index 481492d..cb6e7c4 100644 --- a/internal/adapter/openai/leaked_output_sanitize.go +++ b/internal/adapter/openai/leaked_output_sanitize.go @@ -24,6 +24,8 @@ var leakedAgentXMLBlockPatterns = []*regexp.Regexp{ } var leakedAgentWrapperTagPattern = regexp.MustCompile(`(?is)]*>`) +var leakedAgentWrapperPlusResultOpenPattern = regexp.MustCompile(`(?is)<(?:attempt_completion|ask_followup_question|new_task)\b[^>]*>\s*`) +var leakedAgentResultPlusWrapperClosePattern = regexp.MustCompile(`(?is)\s*]*>`) var leakedAgentResultTagPattern = regexp.MustCompile(`(?is)`) func sanitizeLeakedOutput(text string) string { @@ -53,11 +55,16 @@ func sanitizeLeakedAgentXMLBlocks(text string) string { } // 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. + // wrapper tags, strip only adjacent tags to avoid exposing agent + // markup without altering unrelated user-visible examples. if leakedAgentWrapperTagPattern.MatchString(out) { + out = leakedAgentWrapperPlusResultOpenPattern.ReplaceAllStringFunc(out, func(match string) string { + return leakedAgentResultTagPattern.ReplaceAllString(match, "") + }) + out = leakedAgentResultPlusWrapperClosePattern.ReplaceAllStringFunc(out, func(match string) string { + return leakedAgentResultTagPattern.ReplaceAllString(match, "") + }) 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 1898289..6548d39 100644 --- a/internal/adapter/openai/leaked_output_sanitize_test.go +++ b/internal/adapter/openai/leaked_output_sanitize_test.go @@ -57,3 +57,12 @@ func TestSanitizeLeakedOutputRemovesDanglingAgentXMLClosingTags(t *testing.T) { t.Fatalf("unexpected sanitize result for dangling closing tags: %q", got) } } + +func TestSanitizeLeakedOutputPreservesUnrelatedResultTagsWhenWrapperLeaks(t *testing.T) { + raw := "Done.Some final answer\nExample XML: value" + got := sanitizeLeakedOutput(raw) + want := "Done.Some final answer\nExample XML: value" + if got != want { + t.Fatalf("unexpected sanitize result for mixed leaked wrapper + xml example: %q", got) + } +}