Skip to content

Commit 279fb39

Browse files
elkaixelkaix
committed
fix: require closing tag before flush extraction and handle whitespace in tag prefixes
1 parent 20784b0 commit 279fb39

4 files changed

Lines changed: 96 additions & 38 deletions

File tree

‎packages/agent-core-v2/src/kosong/provider/bases/openai/dsml-tool-parser.ts‎

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,5 @@
11
import type { StreamedMessagePart, ToolCall } from '#/kosong/contract/message';
22

3-
const CANDIDATE_TARGETS = [
4-
'dsml｜tool_calls',
5-
'dsml|tool_calls',
6-
'dsmltool_calls',
7-
'dsml｜invoke',
8-
'dsml|invoke',
9-
'dsmlinvoke',
10-
'tool_calls',
11-
'tool_call',
12-
'invoke',
13-
];
14-
153
const CONTAINER_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i;
164
const CONTAINER_CLOSE_RE = /^<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i;
175
const INVOKE_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke(?:\s+[^>]*)?>/i;
@@ -110,9 +98,8 @@ function parseInvokeTag(invokeBlock: string): ToolCall | null {
11098
if (!toolName) return null;
11199

112100
const closeMatch = INVOKE_CLOSE_RE.exec(invokeBlock);
113-
const innerContent = closeMatch
114-
? invokeBlock.slice(openMatch[0].length, closeMatch.index)
115-
: invokeBlock.slice(openMatch[0].length);
101+
if (!closeMatch) return null;
102+
const innerContent = invokeBlock.slice(openMatch[0].length, closeMatch.index);
116103

117104
const args = parseInvokeBody(innerContent);
118105
return {
@@ -124,6 +111,7 @@ function parseInvokeTag(invokeBlock: string): ToolCall | null {
124111
}
125112

126113
function parseHermesToolCall(toolCallBlock: string): ToolCall | null {
114+
if (!HERMES_CLOSE_RE.test(toolCallBlock)) return null;
127115
const inner = toolCallBlock
128116
.replace(/^<tool_call>/i, '')
129117
.replace(/<\/tool_call>$/i, '')
@@ -150,11 +138,19 @@ function isPotentialTagPrefix(s: string): boolean {
150138
if (!s.startsWith('<')) return false;
151139
const lower = s.toLowerCase();
152140
let rest = lower.startsWith('</') ? lower.slice(2) : lower.slice(1);
141+
rest = rest.trimStart();
153142
if (rest.startsWith('｜') || rest.startsWith('|')) {
154-
rest = rest.slice(1);
143+
rest = rest.slice(1).trimStart();
144+
}
145+
if (rest.startsWith('dsml')) {
146+
rest = rest.slice(4).trimStart();
147+
if (rest.startsWith('｜') || rest.startsWith('|')) {
148+
rest = rest.slice(1).trimStart();
149+
}
155150
}
156151
if (rest.length === 0) return true;
157-
return CANDIDATE_TARGETS.some((target) => target.startsWith(rest) || rest.startsWith(target));
152+
const targets = ['tool_calls', 'tool_call', 'invoke', 'parameter'];
153+
return targets.some((t) => t.startsWith(rest) || rest.startsWith(t));
158154
}
159155

160156
export class DsmlStreamParser {
@@ -276,7 +272,7 @@ export class DsmlStreamParser {
276272
const parts: StreamedMessagePart[] = [];
277273
if (this._buffer.length > 0) {
278274
const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer);
279-
if (invokeOpen) {
275+
if (invokeOpen && INVOKE_CLOSE_RE.test(this._buffer)) {
280276
const toolCall = parseInvokeTag(this._buffer);
281277
if (toolCall) {
282278
this._hasExtractedToolCalls = true;
@@ -286,7 +282,7 @@ export class DsmlStreamParser {
286282
}
287283
}
288284
const hermesOpen = HERMES_OPEN_RE.exec(this._buffer);
289-
if (hermesOpen) {
285+
if (hermesOpen && HERMES_CLOSE_RE.test(this._buffer)) {
290286
const toolCall = parseHermesToolCall(this._buffer);
291287
if (toolCall) {
292288
this._hasExtractedToolCalls = true;

‎packages/agent-core-v2/test/kosong/provider/dsml-tool-parser.test.ts‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,5 +354,38 @@ describe('agent-core-v2: DsmlStreamParser and extractDsmlToolCalls', () => {
354354
expect(result.cleanText).toBe(input);
355355
expect(result.toolCalls).toHaveLength(0);
356356
});
357+
358+
it('handles stream split after whitespace in tag prefix', () => {
359+
const parser = new DsmlStreamParser();
360+
const chunks = [
361+
'< ',
362+
'| DSML ',
363+
'| invoke name="Read">\n< | DSML | parameter name="filePath">src/app.ts</ | DSML | parameter>\n</ | DSML | invoke>',
364+
];
365+
366+
const parts = [];
367+
for (const chunk of chunks) {
368+
parts.push(...parser.feed(chunk));
369+
}
370+
parts.push(...parser.flush());
371+
372+
expect(parser.hasExtractedToolCalls).toBe(true);
373+
const toolParts = parts.filter((p) => p.type === 'function');
374+
expect(toolParts).toHaveLength(1);
375+
expect(toolParts[0]?.name).toBe('Read');
376+
});
377+
378+
it('preserves unclosed Hermes block at flush as text', () => {
379+
const parser = new DsmlStreamParser();
380+
const parts = [
381+
...parser.feed('<tool_call>{"name": "Read"}'),
382+
...parser.flush(),
383+
];
384+
385+
expect(parser.hasExtractedToolCalls).toBe(false);
386+
expect(parts).toEqual([
387+
{ type: 'text', text: '<tool_call>{"name": "Read"}' },
388+
]);
389+
});
357390
});
358391
});

‎packages/kosong/src/providers/dsml-tool-parser.ts‎

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,5 @@
11
import type { StreamedMessagePart, ToolCall } from '#/message';
22

3-
const CANDIDATE_TARGETS = [
4-
'dsml｜tool_calls',
5-
'dsml|tool_calls',
6-
'dsmltool_calls',
7-
'dsml｜invoke',
8-
'dsml|invoke',
9-
'dsmlinvoke',
10-
'tool_calls',
11-
'tool_call',
12-
'invoke',
13-
];
14-
153
const CONTAINER_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i;
164
const CONTAINER_CLOSE_RE = /^<\/\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*tool_calls\s*>/i;
175
const INVOKE_OPEN_RE = /^<\s*[｜|]?\s*(?:DSML\s*[｜|]?)?\s*invoke(?:\s+[^>]*)?>/i;
@@ -110,9 +98,8 @@ function parseInvokeTag(invokeBlock: string): ToolCall | null {
11098
if (!toolName) return null;
11199

112100
const closeMatch = INVOKE_CLOSE_RE.exec(invokeBlock);
113-
const innerContent = closeMatch
114-
? invokeBlock.slice(openMatch[0].length, closeMatch.index)
115-
: invokeBlock.slice(openMatch[0].length);
101+
if (!closeMatch) return null;
102+
const innerContent = invokeBlock.slice(openMatch[0].length, closeMatch.index);
116103

117104
const args = parseInvokeBody(innerContent);
118105
return {
@@ -124,6 +111,7 @@ function parseInvokeTag(invokeBlock: string): ToolCall | null {
124111
}
125112

126113
function parseHermesToolCall(toolCallBlock: string): ToolCall | null {
114+
if (!HERMES_CLOSE_RE.test(toolCallBlock)) return null;
127115
const inner = toolCallBlock
128116
.replace(/^<tool_call>/i, '')
129117
.replace(/<\/tool_call>$/i, '')
@@ -150,11 +138,19 @@ function isPotentialTagPrefix(s: string): boolean {
150138
if (!s.startsWith('<')) return false;
151139
const lower = s.toLowerCase();
152140
let rest = lower.startsWith('</') ? lower.slice(2) : lower.slice(1);
141+
rest = rest.trimStart();
153142
if (rest.startsWith('｜') || rest.startsWith('|')) {
154-
rest = rest.slice(1);
143+
rest = rest.slice(1).trimStart();
144+
}
145+
if (rest.startsWith('dsml')) {
146+
rest = rest.slice(4).trimStart();
147+
if (rest.startsWith('｜') || rest.startsWith('|')) {
148+
rest = rest.slice(1).trimStart();
149+
}
155150
}
156151
if (rest.length === 0) return true;
157-
return CANDIDATE_TARGETS.some((target) => target.startsWith(rest) || rest.startsWith(target));
152+
const targets = ['tool_calls', 'tool_call', 'invoke', 'parameter'];
153+
return targets.some((t) => t.startsWith(rest) || rest.startsWith(t));
158154
}
159155

160156
export class DsmlStreamParser {
@@ -276,7 +272,7 @@ export class DsmlStreamParser {
276272
const parts: StreamedMessagePart[] = [];
277273
if (this._buffer.length > 0) {
278274
const invokeOpen = INVOKE_OPEN_RE.exec(this._buffer);
279-
if (invokeOpen) {
275+
if (invokeOpen && INVOKE_CLOSE_RE.test(this._buffer)) {
280276
const toolCall = parseInvokeTag(this._buffer);
281277
if (toolCall) {
282278
this._hasExtractedToolCalls = true;
@@ -286,7 +282,7 @@ export class DsmlStreamParser {
286282
}
287283
}
288284
const hermesOpen = HERMES_OPEN_RE.exec(this._buffer);
289-
if (hermesOpen) {
285+
if (hermesOpen && HERMES_CLOSE_RE.test(this._buffer)) {
290286
const toolCall = parseHermesToolCall(this._buffer);
291287
if (toolCall) {
292288
this._hasExtractedToolCalls = true;

‎packages/kosong/test/dsml-tool-parser.test.ts‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,5 +205,38 @@ describe('DsmlStreamParser and extractDsmlToolCalls', () => {
205205
expect(fullText).toBe('if (x < 5 && y > 2)');
206206
expect(parser.hasExtractedToolCalls).toBe(false);
207207
});
208+
209+
it('handles stream split after whitespace in tag prefix', () => {
210+
const parser = new DsmlStreamParser();
211+
const chunks = [
212+
'< ',
213+
'| DSML ',
214+
'| invoke name="Read">\n< | DSML | parameter name="filePath">src/app.ts</ | DSML | parameter>\n</ | DSML | invoke>',
215+
];
216+
217+
const parts = [];
218+
for (const chunk of chunks) {
219+
parts.push(...parser.feed(chunk));
220+
}
221+
parts.push(...parser.flush());
222+
223+
expect(parser.hasExtractedToolCalls).toBe(true);
224+
const toolParts = parts.filter((p) => p.type === 'function');
225+
expect(toolParts).toHaveLength(1);
226+
expect(toolParts[0]?.name).toBe('Read');
227+
});
228+
229+
it('preserves unclosed Hermes block at flush as text', () => {
230+
const parser = new DsmlStreamParser();
231+
const parts = [
232+
...parser.feed('<tool_call>{"name": "Read"}'),
233+
...parser.flush(),
234+
];
235+
236+
expect(parser.hasExtractedToolCalls).toBe(false);
237+
expect(parts).toEqual([
238+
{ type: 'text', text: '<tool_call>{"name": "Read"}' },
239+
]);
240+
});
208241
});
209242
});

0 commit comments

Comments
 (0)