Skip to content
This repository was archived by the owner on May 13, 2026. It is now read-only.

Commit 92b3093

Browse files
CJackHwangCJackHwang
authored
Merge pull request #483 from NgoQuocViet2001/ai/tool-result-leak-sanitize
fix(openai): strip leaked tool result markers
2 parents 3569ae1 + 3e935c0 commit 92b3093

4 files changed

Lines changed: 150 additions & 4 deletions

File tree

‎internal/httpapi/openai/leaked_output_sanitize_test.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,24 @@ func TestSanitizeLeakedOutputRemovesFullwidthDelimitedMetaMarkers(t *testing.T)
3535
}
3636
}
3737

38+
func TestSanitizeLeakedOutputRemovesAssistantEndOfToolCallsMarkers(t *testing.T) {
39+
fw := "\uff5c"
40+
raw := "A<|Assistant_END_OF_TOOL_CALLS|>B<" + fw + "Assistant▁END▁OF▁TOOL_CALLS" + fw + ">C<|end_of_tool_calls|>D"
41+
got := sanitizeLeakedOutput(raw)
42+
if got != "ABCD" {
43+
t.Fatalf("unexpected sanitize result for assistant end-of-tool-calls markers: %q", got)
44+
}
45+
}
46+
47+
func TestSanitizeLeakedOutputRemovesFullToolResultSection(t *testing.T) {
48+
fw := "\uff5c"
49+
raw := "开始<" + fw + "Tool" + fw + ">[{\"content\":\"openjdk version 21\"}]<" + fw + "end▁of▁toolresults" + fw + ">结束"
50+
got := sanitizeLeakedOutput(raw)
51+
if got != "开始结束" {
52+
t.Fatalf("unexpected sanitize result for leaked tool result section: %q", got)
53+
}
54+
}
55+
3856
func TestSanitizeLeakedOutputRemovesThinkAndBosMarkers(t *testing.T) {
3957
raw := "A<think>B</think>C<|begin▁of▁sentence|>D<| begin_of_sentence |>E<|begin_of_sentence|>F"
4058
got := sanitizeLeakedOutput(raw)

‎internal/httpapi/openai/shared/leaked_output_sanitize.go‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ import (
1010
var emptyJSONFencePattern = regexp.MustCompile("(?is)```json\\s*```")
1111
var leakedToolCallArrayPattern = regexp.MustCompile(`(?is)\[\{\s*"function"\s*:\s*\{[\s\S]*?\}\s*,\s*"id"\s*:\s*"call[^"]*"\s*,\s*"type"\s*:\s*"function"\s*}\]`)
1212
var leakedToolResultBlobPattern = regexp.MustCompile(`(?is)<\s*\|\s*tool\s*\|\s*>\s*\{[\s\S]*?"tool_call_id"\s*:\s*"call[^"]*"\s*}`)
13+
var leakedToolResultOpenMarkerPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*tool\s*[\|\x{ff5c}]>`)
14+
var leakedToolResultCloseMarkerPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*end[_▁]of[_▁]tool[_▁]?results\s*[\|\x{ff5c}]>`)
15+
var leakedToolResultSectionPattern = regexp.MustCompile(`(?is)<[\|\x{ff5c}]\s*tool\s*[\|\x{ff5c}]>[\s\S]*?<[\|\x{ff5c}]\s*end[_▁]of[_▁]tool[_▁]?results\s*[\|\x{ff5c}]>`)
1316

1417
var leakedThinkTagPattern = regexp.MustCompile(`(?is)</?\s*think\s*>`)
1518

@@ -29,7 +32,8 @@ var leakedThoughtMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:begi
2932
// halfwidth or legacy U+FF5C fullwidth delimiters:
3033
// - ASCII underscore: <|end_of_sentence|>, <|end_of_toolresults|>, <|end_of_instructions|>
3134
// - U+2581 variant: <|end▁of▁sentence|>, <|end▁of▁toolresults|>, <|end▁of▁instructions|>
32-
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}]>`)
35+
// - compound assistant markers: <|Assistant_END_OF_TOOL_CALLS|>
36+
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}]>`)
3337

3438
// leakedAgentXMLBlockPatterns catch agent-style XML blocks that leak through
3539
// when the sieve fails to capture them. These are applied only to complete
@@ -52,6 +56,7 @@ func sanitizeLeakedOutput(text string) string {
5256
}
5357
out := emptyJSONFencePattern.ReplaceAllString(text, "")
5458
out = leakedToolCallArrayPattern.ReplaceAllString(out, "")
59+
out = leakedToolResultSectionPattern.ReplaceAllString(out, "")
5560
out = leakedToolResultBlobPattern.ReplaceAllString(out, "")
5661
out = stripDanglingThinkSuffix(out)
5762
out = leakedThinkTagPattern.ReplaceAllString(out, "")
@@ -63,6 +68,40 @@ func sanitizeLeakedOutput(text string) string {
6368
return out
6469
}
6570

71+
func stripLeakedToolResultSectionsDelta(text string, inside *bool) string {
72+
if text == "" || inside == nil {
73+
return text
74+
}
75+
var b strings.Builder
76+
pos := 0
77+
for pos < len(text) {
78+
if *inside {
79+
loc := leakedToolResultCloseMarkerPattern.FindStringIndex(text[pos:])
80+
if loc == nil {
81+
return b.String()
82+
}
83+
*inside = false
84+
pos += loc[1]
85+
continue
86+
}
87+
loc := leakedToolResultOpenMarkerPattern.FindStringIndex(text[pos:])
88+
if loc == nil {
89+
b.WriteString(text[pos:])
90+
break
91+
}
92+
start := pos + loc[0]
93+
openEnd := pos + loc[1]
94+
b.WriteString(text[pos:start])
95+
closeLoc := leakedToolResultCloseMarkerPattern.FindStringIndex(text[openEnd:])
96+
if closeLoc == nil {
97+
*inside = true
98+
break
99+
}
100+
pos = openEnd + closeLoc[1]
101+
}
102+
return b.String()
103+
}
104+
66105
func stripLeakedToolCallWrapperBlocks(text string) string {
67106
if text == "" {
68107
return text

‎internal/httpapi/openai/shared/stream_accumulator.go‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ type StreamAccumulator struct {
1616
ToolDetectionThinking strings.Builder
1717
RawText strings.Builder
1818
Text strings.Builder
19+
20+
thinkingToolResultSectionOpen bool
21+
textToolResultSectionOpen bool
1922
}
2023

2124
type StreamPartDelta struct {
@@ -69,7 +72,8 @@ func (a *StreamAccumulator) applyThinkingPart(text string) StreamPartDelta {
6972
if !a.ThinkingEnabled || rawTrimmed == "" {
7073
return delta
7174
}
72-
cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers)
75+
visibleCandidate := stripLeakedToolResultSectionsDelta(rawTrimmed, &a.thinkingToolResultSectionOpen)
76+
cleanedText := CleanVisibleOutput(visibleCandidate, a.StripReferenceMarkers)
7377
if cleanedText == "" {
7478
return delta
7579
}
@@ -89,11 +93,15 @@ func (a *StreamAccumulator) applyTextPart(text string) StreamPartDelta {
8993
}
9094
a.RawText.WriteString(rawTrimmed)
9195
delta := StreamPartDelta{Type: "text", RawText: rawTrimmed}
92-
if a.SearchEnabled && sse.IsCitation(rawTrimmed) {
96+
visibleCandidate := stripLeakedToolResultSectionsDelta(rawTrimmed, &a.textToolResultSectionOpen)
97+
if visibleCandidate == "" {
98+
return delta
99+
}
100+
if a.SearchEnabled && sse.IsCitation(visibleCandidate) {
93101
delta.CitationOnly = true
94102
return delta
95103
}
96-
cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers)
104+
cleanedText := CleanVisibleOutput(visibleCandidate, a.StripReferenceMarkers)
97105
trimmed := sse.TrimContinuationOverlapFromBuilder(&a.Text, cleanedText)
98106
if trimmed == "" {
99107
return delta

‎internal/httpapi/openai/shared/stream_accumulator_test.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,87 @@ func TestStreamAccumulatorSuppressesCitationTextWhenSearchEnabled(t *testing.T)
9696
}
9797
}
9898

99+
func TestStreamAccumulatorStripsToolResultSectionAcrossTextChunks(t *testing.T) {
100+
acc := StreamAccumulator{StripReferenceMarkers: true}
101+
first := acc.Apply(sse.LineResult{
102+
Parsed: true,
103+
Parts: []sse.ContentPart{{Type: "text", Text: "visible:<|Tool|>"}},
104+
})
105+
second := acc.Apply(sse.LineResult{
106+
Parsed: true,
107+
Parts: []sse.ContentPart{{Type: "text", Text: `[{"content":"secret","tool_call_id":"call_123"}]`}},
108+
})
109+
third := acc.Apply(sse.LineResult{
110+
Parsed: true,
111+
Parts: []sse.ContentPart{{Type: "text", Text: "<|end_of_toolresults|> after"}},
112+
})
113+
114+
if got := acc.RawText.String(); got != `visible:<|Tool|>[{"content":"secret","tool_call_id":"call_123"}]<|end_of_toolresults|> after` {
115+
t.Fatalf("raw text = %q", got)
116+
}
117+
if got := acc.Text.String(); got != "visible: after" {
118+
t.Fatalf("visible text = %q", got)
119+
}
120+
if !first.ContentSeen || !second.ContentSeen || !third.ContentSeen {
121+
t.Fatalf("expected all chunks to mark upstream content")
122+
}
123+
if got := first.Parts[0].VisibleText; got != "visible:" {
124+
t.Fatalf("first visible delta = %q", got)
125+
}
126+
if got := second.Parts[0].VisibleText; got != "" {
127+
t.Fatalf("payload visible delta = %q", got)
128+
}
129+
if got := third.Parts[0].VisibleText; got != " after" {
130+
t.Fatalf("closing visible delta = %q", got)
131+
}
132+
}
133+
134+
func TestStreamAccumulatorStripsFullwidthToolResultSectionAcrossTextChunks(t *testing.T) {
135+
acc := StreamAccumulator{StripReferenceMarkers: true}
136+
acc.Apply(sse.LineResult{
137+
Parsed: true,
138+
Parts: []sse.ContentPart{{Type: "text", Text: "x<｜Tool｜>"}},
139+
})
140+
acc.Apply(sse.LineResult{
141+
Parsed: true,
142+
Parts: []sse.ContentPart{{Type: "text", Text: `{"content":"secret"}`}},
143+
})
144+
acc.Apply(sse.LineResult{
145+
Parsed: true,
146+
Parts: []sse.ContentPart{{Type: "text", Text: "<｜end▁of▁toolresults｜>y"}},
147+
})
148+
149+
if got := acc.Text.String(); got != "xy" {
150+
t.Fatalf("visible text = %q", got)
151+
}
152+
}
153+
154+
func TestStreamAccumulatorStripsToolResultSectionAcrossThinkingChunks(t *testing.T) {
155+
acc := StreamAccumulator{ThinkingEnabled: true, StripReferenceMarkers: true}
156+
acc.Apply(sse.LineResult{
157+
Parsed: true,
158+
Parts: []sse.ContentPart{{Type: "thinking", Text: "thought <|Tool|>"}},
159+
})
160+
payload := acc.Apply(sse.LineResult{
161+
Parsed: true,
162+
Parts: []sse.ContentPart{{Type: "thinking", Text: `[{"content":"secret"}]`}},
163+
})
164+
acc.Apply(sse.LineResult{
165+
Parsed: true,
166+
Parts: []sse.ContentPart{{Type: "thinking", Text: "<|end_of_toolresults|>resumes"}},
167+
})
168+
169+
if got := acc.RawThinking.String(); got != `thought <|Tool|>[{"content":"secret"}]<|end_of_toolresults|>resumes` {
170+
t.Fatalf("raw thinking = %q", got)
171+
}
172+
if got := acc.Thinking.String(); got != "thought resumes" {
173+
t.Fatalf("visible thinking = %q", got)
174+
}
175+
if got := payload.Parts[0].VisibleText; got != "" {
176+
t.Fatalf("payload visible delta = %q", got)
177+
}
178+
}
179+
99180
func TestStreamAccumulatorStripsInlineCitationAndReferenceMarkers(t *testing.T) {
100181
acc := StreamAccumulator{SearchEnabled: true, StripReferenceMarkers: true}
101182
result := acc.Apply(sse.LineResult{

0 commit comments

Comments
 (0)