diff --git a/internal/httpapi/openai/leaked_output_sanitize_test.go b/internal/httpapi/openai/leaked_output_sanitize_test.go index 939f73f..20ace0e 100644 --- a/internal/httpapi/openai/leaked_output_sanitize_test.go +++ b/internal/httpapi/openai/leaked_output_sanitize_test.go @@ -35,6 +35,24 @@ func TestSanitizeLeakedOutputRemovesFullwidthDelimitedMetaMarkers(t *testing.T) } } +func TestSanitizeLeakedOutputRemovesAssistantEndOfToolCallsMarkers(t *testing.T) { + fw := "\uff5c" + raw := "A<|Assistant_END_OF_TOOL_CALLS|>B<" + fw + "Assistant▁END▁OF▁TOOL_CALLS" + fw + ">C<|end_of_tool_calls|>D" + got := sanitizeLeakedOutput(raw) + if got != "ABCD" { + t.Fatalf("unexpected sanitize result for assistant end-of-tool-calls markers: %q", got) + } +} + +func TestSanitizeLeakedOutputRemovesFullToolResultSection(t *testing.T) { + fw := "\uff5c" + raw := "开始<" + fw + "Tool" + fw + ">[{\"content\":\"openjdk version 21\"}]<" + fw + "end▁of▁toolresults" + fw + ">结束" + got := sanitizeLeakedOutput(raw) + if got != "开始结束" { + t.Fatalf("unexpected sanitize result for leaked tool result section: %q", got) + } +} + func TestSanitizeLeakedOutputRemovesThinkAndBosMarkers(t *testing.T) { raw := "ABC<|begin▁of▁sentence|>D<| begin_of_sentence |>E<|begin_of_sentence|>F" got := sanitizeLeakedOutput(raw) diff --git a/internal/httpapi/openai/shared/leaked_output_sanitize.go b/internal/httpapi/openai/shared/leaked_output_sanitize.go index 9293e78..64a6059 100644 --- a/internal/httpapi/openai/shared/leaked_output_sanitize.go +++ b/internal/httpapi/openai/shared/leaked_output_sanitize.go @@ -10,6 +10,9 @@ import ( var emptyJSONFencePattern = regexp.MustCompile("(?is)```json\\s*```") var leakedToolCallArrayPattern = regexp.MustCompile(`(?is)\[\{\s*"function"\s*:\s*\{[\s\S]*?\}\s*,\s*"id"\s*:\s*"call[^"]*"\s*,\s*"type"\s*:\s*"function"\s*}\]`) var leakedToolResultBlobPattern = regexp.MustCompile(`(?is)<\s*\|\s*tool\s*\|\s*>\s*\{[\s\S]*?"tool_call_id"\s*:\s*"call[^"]*"\s*}`) +var leakedToolResultOpenMarkerPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*tool\s*[\|\x{ff5c}]>`) +var leakedToolResultCloseMarkerPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*end[_▁]of[_▁]tool[_▁]?results\s*[\|\x{ff5c}]>`) +var leakedToolResultSectionPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*tool\s*[\|\x{ff5c}]>[\s\S]*?<[\|\x{ff5c}]\s*end[_▁]of[_▁]tool[_▁]?results\s*[\|\x{ff5c}]>`) var leakedThinkTagPattern = regexp.MustCompile(`(?is)`) @@ -29,7 +32,8 @@ var leakedThoughtMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:begi // halfwidth or legacy U+FF5C fullwidth delimiters: // - ASCII underscore: <|end_of_sentence|>, <|end_of_toolresults|>, <|end_of_instructions|> // - U+2581 variant: <|end▁of▁sentence|>, <|end▁of▁toolresults|>, <|end▁of▁instructions|> -var leakedMetaMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:assistant|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]toolresults|end[_▁]of[_▁]instructions)\s*[\|\x{ff5c}]>`) +// - compound assistant markers: <|Assistant_END_OF_TOOL_CALLS|> +var leakedMetaMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:assistant(?:[_▁]end[_▁]of[_▁]tool[_▁]?calls)?|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]tool[_▁]?results|end[_▁]of[_▁]tool[_▁]?calls|end[_▁]of[_▁]instructions)\s*[\|\x{ff5c}]>`) // leakedAgentXMLBlockPatterns catch agent-style XML blocks that leak through // when the sieve fails to capture them. These are applied only to complete @@ -52,6 +56,7 @@ func sanitizeLeakedOutput(text string) string { } out := emptyJSONFencePattern.ReplaceAllString(text, "") out = leakedToolCallArrayPattern.ReplaceAllString(out, "") + out = leakedToolResultSectionPattern.ReplaceAllString(out, "") out = leakedToolResultBlobPattern.ReplaceAllString(out, "") out = stripDanglingThinkSuffix(out) out = leakedThinkTagPattern.ReplaceAllString(out, "") @@ -63,6 +68,40 @@ func sanitizeLeakedOutput(text string) string { return out } +func stripLeakedToolResultSectionsDelta(text string, inside *bool) string { + if text == "" || inside == nil { + return text + } + var b strings.Builder + pos := 0 + for pos < len(text) { + if *inside { + loc := leakedToolResultCloseMarkerPattern.FindStringIndex(text[pos:]) + if loc == nil { + return b.String() + } + *inside = false + pos += loc[1] + continue + } + loc := leakedToolResultOpenMarkerPattern.FindStringIndex(text[pos:]) + if loc == nil { + b.WriteString(text[pos:]) + break + } + start := pos + loc[0] + openEnd := pos + loc[1] + b.WriteString(text[pos:start]) + closeLoc := leakedToolResultCloseMarkerPattern.FindStringIndex(text[openEnd:]) + if closeLoc == nil { + *inside = true + break + } + pos = openEnd + closeLoc[1] + } + return b.String() +} + func stripLeakedToolCallWrapperBlocks(text string) string { if text == "" { return text diff --git a/internal/httpapi/openai/shared/stream_accumulator.go b/internal/httpapi/openai/shared/stream_accumulator.go index 472748e..c84de87 100644 --- a/internal/httpapi/openai/shared/stream_accumulator.go +++ b/internal/httpapi/openai/shared/stream_accumulator.go @@ -16,6 +16,9 @@ type StreamAccumulator struct { ToolDetectionThinking strings.Builder RawText strings.Builder Text strings.Builder + + thinkingToolResultSectionOpen bool + textToolResultSectionOpen bool } type StreamPartDelta struct { @@ -69,7 +72,8 @@ func (a *StreamAccumulator) applyThinkingPart(text string) StreamPartDelta { if !a.ThinkingEnabled || rawTrimmed == "" { return delta } - cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers) + visibleCandidate := stripLeakedToolResultSectionsDelta(rawTrimmed, &a.thinkingToolResultSectionOpen) + cleanedText := CleanVisibleOutput(visibleCandidate, a.StripReferenceMarkers) if cleanedText == "" { return delta } @@ -89,11 +93,15 @@ func (a *StreamAccumulator) applyTextPart(text string) StreamPartDelta { } a.RawText.WriteString(rawTrimmed) delta := StreamPartDelta{Type: "text", RawText: rawTrimmed} - if a.SearchEnabled && sse.IsCitation(rawTrimmed) { + visibleCandidate := stripLeakedToolResultSectionsDelta(rawTrimmed, &a.textToolResultSectionOpen) + if visibleCandidate == "" { + return delta + } + if a.SearchEnabled && sse.IsCitation(visibleCandidate) { delta.CitationOnly = true return delta } - cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers) + cleanedText := CleanVisibleOutput(visibleCandidate, a.StripReferenceMarkers) trimmed := sse.TrimContinuationOverlapFromBuilder(&a.Text, cleanedText) if trimmed == "" { return delta diff --git a/internal/httpapi/openai/shared/stream_accumulator_test.go b/internal/httpapi/openai/shared/stream_accumulator_test.go index 1f4fe93..20292ae 100644 --- a/internal/httpapi/openai/shared/stream_accumulator_test.go +++ b/internal/httpapi/openai/shared/stream_accumulator_test.go @@ -96,6 +96,87 @@ func TestStreamAccumulatorSuppressesCitationTextWhenSearchEnabled(t *testing.T) } } +func TestStreamAccumulatorStripsToolResultSectionAcrossTextChunks(t *testing.T) { + acc := StreamAccumulator{StripReferenceMarkers: true} + first := acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: "visible:<|Tool|>"}}, + }) + second := acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: `[{"content":"secret","tool_call_id":"call_123"}]`}}, + }) + third := acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: "<|end_of_toolresults|> after"}}, + }) + + if got := acc.RawText.String(); got != `visible:<|Tool|>[{"content":"secret","tool_call_id":"call_123"}]<|end_of_toolresults|> after` { + t.Fatalf("raw text = %q", got) + } + if got := acc.Text.String(); got != "visible: after" { + t.Fatalf("visible text = %q", got) + } + if !first.ContentSeen || !second.ContentSeen || !third.ContentSeen { + t.Fatalf("expected all chunks to mark upstream content") + } + if got := first.Parts[0].VisibleText; got != "visible:" { + t.Fatalf("first visible delta = %q", got) + } + if got := second.Parts[0].VisibleText; got != "" { + t.Fatalf("payload visible delta = %q", got) + } + if got := third.Parts[0].VisibleText; got != " after" { + t.Fatalf("closing visible delta = %q", got) + } +} + +func TestStreamAccumulatorStripsFullwidthToolResultSectionAcrossTextChunks(t *testing.T) { + acc := StreamAccumulator{StripReferenceMarkers: true} + acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: "x<|Tool|>"}}, + }) + acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: `{"content":"secret"}`}}, + }) + acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "text", Text: "<|end▁of▁toolresults|>y"}}, + }) + + if got := acc.Text.String(); got != "xy" { + t.Fatalf("visible text = %q", got) + } +} + +func TestStreamAccumulatorStripsToolResultSectionAcrossThinkingChunks(t *testing.T) { + acc := StreamAccumulator{ThinkingEnabled: true, StripReferenceMarkers: true} + acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "thinking", Text: "thought <|Tool|>"}}, + }) + payload := acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "thinking", Text: `[{"content":"secret"}]`}}, + }) + acc.Apply(sse.LineResult{ + Parsed: true, + Parts: []sse.ContentPart{{Type: "thinking", Text: "<|end_of_toolresults|>resumes"}}, + }) + + if got := acc.RawThinking.String(); got != `thought <|Tool|>[{"content":"secret"}]<|end_of_toolresults|>resumes` { + t.Fatalf("raw thinking = %q", got) + } + if got := acc.Thinking.String(); got != "thought resumes" { + t.Fatalf("visible thinking = %q", got) + } + if got := payload.Parts[0].VisibleText; got != "" { + t.Fatalf("payload visible delta = %q", got) + } +} + func TestStreamAccumulatorStripsInlineCitationAndReferenceMarkers(t *testing.T) { acc := StreamAccumulator{SearchEnabled: true, StripReferenceMarkers: true} result := acc.Apply(sse.LineResult{