Skip to content

Commit 6a819e2

Browse files
Filter XML tool call tags from SDK stream chunks
Implement stateful XML parser to remove <codebuff_tool_call> tags from handleStreamChunk output while preserving text before/after tool calls and handling tags split across chunk boundaries. 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent e1da277 commit 6a819e2

File tree

1 file changed

+32
-2
lines changed

1 file changed

+32
-2
lines changed

sdk/src/run.ts

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,10 @@ export async function run({
112112
})
113113

114114
// TODO: bad pattern, switch to using SSE and move off of websockets
115+
let insideToolCall = false
116+
let buffer = ''
117+
const BUFFER_SIZE = 100
118+
115119
const websocketHandler = new WebSocketHandler({
116120
apiKey,
117121
onWebsocketError: (error) => {
@@ -144,8 +148,34 @@ export async function run({
144148
onResponseChunk: async (action) => {
145149
const { userInputId, chunk } = action
146150
if (typeof chunk === 'string') {
147-
await handleStreamChunk?.(chunk)
148-
} else {
151+
buffer += chunk
152+
153+
if (!insideToolCall && buffer.includes('<codebuff_tool_call>')) {
154+
const openTagIndex = buffer.indexOf('<codebuff_tool_call>')
155+
const beforeTag = buffer.substring(0, openTagIndex)
156+
if (beforeTag) {
157+
await handleStreamChunk?.(beforeTag)
158+
}
159+
insideToolCall = true
160+
buffer = buffer.substring(openTagIndex)
161+
} else if (insideToolCall && buffer.includes('</codebuff_tool_call>')) {
162+
const closeTagIndex = buffer.indexOf('</codebuff_tool_call>')
163+
insideToolCall = false
164+
buffer = buffer.substring(
165+
closeTagIndex + '</codebuff_tool_call>'.length,
166+
)
167+
} else if (!insideToolCall) {
168+
if (buffer.length > 50) {
169+
const safeToOutput = buffer.substring(0, buffer.length - 50)
170+
await handleStreamChunk?.(safeToOutput)
171+
buffer = buffer.substring(buffer.length - 50)
172+
}
173+
}
174+
175+
if (insideToolCall && buffer.length > BUFFER_SIZE * 10) {
176+
buffer = buffer.slice(-BUFFER_SIZE * 10)
177+
}
178+
} else if (chunk.type !== 'tool_call') {
149179
await handleEvent?.(chunk)
150180
}
151181
},

0 commit comments

Comments
 (0)