A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 02:00 UTC
[Bugfix][Model] Prevent special token leakage in KimiK2ToolParser streaming mode (#28543)
Signed-off-by: Jscaldwell55 <jay.s.caldwell@gmail.com>
Jay Caldwell committed
Nov 16, 2025 at 23:54 UTC
6f374192442381b37a6a6ba29045c74a8ee2486d
2 files changed
+790
-5
tests/tool_use/test_kimi_k2_tool_parser.py
+593
@@ -209,3 +209,596 @@ def test_streaming_no_tool_calls(kimi_k2_tool_parser):
209
assert result is not None
210
assert hasattr(result, "content")
211
assert result.content == " without any tool calls."
212
+
213
+
214
+def test_token_leak_between_section_and_tool_begin(kimi_k2_tool_parser):
215
+ """
216
+ Test that text between <|tool_calls_section_begin|> and <|tool_call_begin|>
217
+ is suppressed and does not leak into reasoning_delta.
218
+ This is the main vulnerability being fixed.
219
+ """
220
+ kimi_k2_tool_parser.reset_streaming_state()
221
+
222
+ # Get token IDs for the markers
223
+ section_begin_token_id = kimi_k2_tool_parser.vocab.get(
224
+ "<|tool_calls_section_begin|>"
225
+ )
226
+ tool_call_begin_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
227
+
228
+ # Simulate streaming sequence:
229
+ # Delta 1: "I'll help you with that. "
230
+ result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
231
+ previous_text="",
232
+ current_text="I'll help you with that. ",
233
+ delta_text="I'll help you with that. ",
234
+ previous_token_ids=[],
235
+ current_token_ids=[1, 2, 3], # Regular tokens
236
+ delta_token_ids=[1, 2, 3],
237
+ request=None,
238
+ )
239
+ assert result1 is not None
240
+ assert result1.content == "I'll help you with that. "
241
+
242
+ # Delta 2: "<|tool_calls_section_begin|>"
243
+ prev_ids = [1, 2, 3]
244
+ curr_ids = prev_ids + [section_begin_token_id]
245
+ result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
246
+ previous_text="I'll help you with that. ",
247
+ current_text="I'll help you with that. <|tool_calls_section_begin|>",
248
+ delta_text="<|tool_calls_section_begin|>",
249
+ previous_token_ids=prev_ids,
250
+ current_token_ids=curr_ids,
251
+ delta_token_ids=[section_begin_token_id],
252
+ request=None,
253
+ )
254
+ # Section marker should be stripped and suppressed
255
+ assert result2 is None or (result2.content is None or result2.content == "")
256
+
257
+ # Delta 3: " spurious text or tokens " (THE LEAK SCENARIO)
258
+ prev_ids = curr_ids
259
+ curr_ids = curr_ids + [4, 5]
260
+ result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
261
+ previous_text="I'll help you with that. <|tool_calls_section_begin|>",
262
+ current_text="I'll help you with that. <|tool_calls_section_begin|> spurious text ",
263
+ delta_text=" spurious text ",
264
+ previous_token_ids=prev_ids,
265
+ current_token_ids=curr_ids,
266
+ delta_token_ids=[4, 5],
267
+ request=None,
268
+ )
269
+ # CRITICAL: This text should be suppressed, NOT returned as reasoning_delta
270
+ assert result3 is None or (result3.content is None or result3.content == "")
271
+
272
+ # Delta 4: "<|tool_call_begin|>..."
273
+ prev_ids = curr_ids
274
+ curr_ids = curr_ids + [tool_call_begin_token_id]
275
+ _result4 = kimi_k2_tool_parser.extract_tool_calls_streaming(
276
+ previous_text="I'll help you with that. <|tool_calls_section_begin|> spurious text ",
277
+ current_text="I'll help you with that. <|tool_calls_section_begin|> spurious text <|tool_call_begin|>",
278
+ delta_text="<|tool_call_begin|>",
279
+ previous_token_ids=prev_ids,
280
+ current_token_ids=curr_ids,
281
+ delta_token_ids=[tool_call_begin_token_id],
282
+ request=None,
283
+ )
284
+ # Now we're in tool call mode, result depends on internal state
285
+ # The key is that the spurious text from Delta 3 was not leaked
286
+
287
+
288
+def test_split_markers_across_deltas(kimi_k2_tool_parser):
289
+ """
290
+ Test that markers split across delta chunks are correctly detected
291
+ via the rolling buffer mechanism.
292
+ """
293
+ kimi_k2_tool_parser.reset_streaming_state()
294
+
295
+ section_begin_token_id = kimi_k2_tool_parser.vocab.get(
296
+ "<|tool_calls_section_begin|>"
297
+ )
298
+
299
+ # Delta 1: "...reasoning<|tool_calls_sec"
300
+ _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
301
+ previous_text="Some reasoning",
302
+ current_text="Some reasoning<|tool_calls_sec",
303
+ delta_text="<|tool_calls_sec",
304
+ previous_token_ids=[1, 2],
305
+ current_token_ids=[1, 2, 3], # Partial token
306
+ delta_token_ids=[3],
307
+ request=None,
308
+ )
309
+ # Partial token not recognized yet, might be buffered
310
+ # Should return as content or None (depends on implementation)
311
+
312
+ # Delta 2: "tion_begin|> " (completes the marker)
313
+ _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
314
+ previous_text="Some reasoning<|tool_calls_sec",
315
+ current_text="Some reasoning<|tool_calls_section_begin|> ",
316
+ delta_text="tion_begin|> ",
317
+ previous_token_ids=[1, 2, 3],
318
+ current_token_ids=[1, 2, section_begin_token_id, 4],
319
+ delta_token_ids=[section_begin_token_id, 4],
320
+ request=None,
321
+ )
322
+ # Now the complete marker should be detected via buffer
323
+ # The parser should enter tool section mode
324
+ assert kimi_k2_tool_parser.in_tool_section is True
325
+
326
+
327
+def test_marker_variants(kimi_k2_tool_parser):
328
+ """Test that both singular and plural marker variants are recognized."""
329
+ kimi_k2_tool_parser.reset_streaming_state()
330
+
331
+ # Test singular variant: <|tool_call_section_begin|> (note: singular "call")
332
+ singular_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_section_begin|>")
333
+
334
+ if singular_token_id is not None: # Only test if tokenizer supports it
335
+ _result = kimi_k2_tool_parser.extract_tool_calls_streaming(
336
+ previous_text="Reasoning ",
337
+ current_text="Reasoning <|tool_call_section_begin|>",
338
+ delta_text="<|tool_call_section_begin|>",
339
+ previous_token_ids=[1, 2],
340
+ current_token_ids=[1, 2, singular_token_id],
341
+ delta_token_ids=[singular_token_id],
342
+ request=None,
343
+ )
344
+ # Should enter tool section mode with singular variant too
345
+ assert kimi_k2_tool_parser.in_tool_section is True
346
+
347
+
348
+def test_reentry_to_reasoning_after_tool_section(kimi_k2_tool_parser):
349
+ """
350
+ Test that after exiting a tool section with <|tool_calls_section_end|>,
351
+ subsequent text is correctly returned as reasoning content.
352
+ """
353
+ kimi_k2_tool_parser.reset_streaming_state()
354
+
355
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
356
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
357
+
358
+ # Enter tool section
359
+ _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
360
+ previous_text="",
361
+ current_text="<|tool_calls_section_begin|>",
362
+ delta_text="<|tool_calls_section_begin|>",
363
+ previous_token_ids=[],
364
+ current_token_ids=[section_begin_id],
365
+ delta_token_ids=[section_begin_id],
366
+ request=None,
367
+ )
368
+ assert kimi_k2_tool_parser.in_tool_section is True
369
+
370
+ # Exit tool section
371
+ _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
372
+ previous_text="<|tool_calls_section_begin|>",
373
+ current_text="<|tool_calls_section_begin|><|tool_calls_section_end|>",
374
+ delta_text="<|tool_calls_section_end|>",
375
+ previous_token_ids=[section_begin_id],
376
+ current_token_ids=[section_begin_id, section_end_id],
377
+ delta_token_ids=[section_end_id],
378
+ request=None,
379
+ )
380
+ assert kimi_k2_tool_parser.in_tool_section is False
381
+
382
+ # Subsequent reasoning text should be returned normally
383
+ result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
384
+ previous_text="<|tool_calls_section_begin|><|tool_calls_section_end|>",
385
+ current_text="<|tool_calls_section_begin|><|tool_calls_section_end|> More reasoning",
386
+ delta_text=" More reasoning",
387
+ previous_token_ids=[section_begin_id, section_end_id],
388
+ current_token_ids=[section_begin_id, section_end_id, 10, 11],
389
+ delta_token_ids=[10, 11],
390
+ request=None,
391
+ )
392
+ assert result3 is not None
393
+ assert result3.content == " More reasoning"
394
+
395
+
396
+def test_empty_tool_section(kimi_k2_tool_parser):
397
+ """Test an empty tool section (begin immediately followed by end)."""
398
+ kimi_k2_tool_parser.reset_streaming_state()
399
+
400
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
401
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
402
+
403
+ # Section begin
404
+ _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
405
+ previous_text="Reasoning ",
406
+ current_text="Reasoning <|tool_calls_section_begin|>",
407
+ delta_text="<|tool_calls_section_begin|>",
408
+ previous_token_ids=[1],
409
+ current_token_ids=[1, section_begin_id],
410
+ delta_token_ids=[section_begin_id],
411
+ request=None,
412
+ )
413
+
414
+ # Immediate section end
415
+ _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
416
+ previous_text="Reasoning <|tool_calls_section_begin|>",
417
+ current_text="Reasoning <|tool_calls_section_begin|><|tool_calls_section_end|>",
418
+ delta_text="<|tool_calls_section_end|>",
419
+ previous_token_ids=[1, section_begin_id],
420
+ current_token_ids=[1, section_begin_id, section_end_id],
421
+ delta_token_ids=[section_end_id],
422
+ request=None,
423
+ )
424
+ # Should exit cleanly without errors
425
+ assert kimi_k2_tool_parser.in_tool_section is False
426
+
427
+
428
+def test_malformed_tool_section_recovery(kimi_k2_tool_parser):
429
+ """
430
+ Test that the parser recovers from a malformed tool section
431
+ that never closes properly.
432
+ """
433
+ kimi_k2_tool_parser.reset_streaming_state()
434
+
435
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
436
+
437
+ # Enter tool section
438
+ _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
439
+ previous_text="",
440
+ current_text="<|tool_calls_section_begin|>",
441
+ delta_text="<|tool_calls_section_begin|>",
442
+ previous_token_ids=[],
443
+ current_token_ids=[section_begin_id],
444
+ delta_token_ids=[section_begin_id],
445
+ request=None,
446
+ )
447
+ assert kimi_k2_tool_parser.in_tool_section is True
448
+
449
+ # Simulate a lot of text without proper tool calls or section end
450
+ # This should trigger the error recovery mechanism
451
+ large_text = "x" * 10000 # Exceeds max_section_chars
452
+
453
+ result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
454
+ previous_text="<|tool_calls_section_begin|>",
455
+ current_text="<|tool_calls_section_begin|>" + large_text,
456
+ delta_text=large_text,
457
+ previous_token_ids=[section_begin_id],
458
+ current_token_ids=[section_begin_id] + list(range(100, 100 + len(large_text))),
459
+ delta_token_ids=list(range(100, 100 + len(large_text))),
460
+ request=None,
461
+ )
462
+
463
+ # Parser should have force-exited the tool section
464
+ assert kimi_k2_tool_parser.in_tool_section is False
465
+ # And returned the content as reasoning
466
+ assert result2 is not None
467
+ assert result2.content == large_text
468
+
469
+
470
+def test_state_reset(kimi_k2_tool_parser):
471
+ """Test that reset_streaming_state() properly clears all state."""
472
+ # Put parser in a complex state
473
+ kimi_k2_tool_parser.in_tool_section = True
474
+ kimi_k2_tool_parser.token_buffer = "some buffer"
475
+ kimi_k2_tool_parser.current_tool_id = 5
476
+ kimi_k2_tool_parser.prev_tool_call_arr = [{"id": "test"}]
477
+ kimi_k2_tool_parser.section_char_count = 1000
478
+
479
+ # Reset
480
+ kimi_k2_tool_parser.reset_streaming_state()
481
+
482
+ # Verify all state is cleared
483
+ assert kimi_k2_tool_parser.in_tool_section is False
484
+ assert kimi_k2_tool_parser.token_buffer == ""
485
+ assert kimi_k2_tool_parser.current_tool_id == -1
486
+ assert kimi_k2_tool_parser.prev_tool_call_arr == []
487
+ assert kimi_k2_tool_parser.section_char_count == 0
488
+ assert kimi_k2_tool_parser.current_tool_name_sent is False
489
+ assert kimi_k2_tool_parser.streamed_args_for_tool == []
490
+
491
+
492
+def test_section_begin_noise_tool_begin_same_chunk(kimi_k2_tool_parser):
493
+ """
494
+ Test that begin→noise→tool_begin within the SAME chunk suppresses
495
+ the noise text correctly (not just across chunks).
496
+ """
497
+ kimi_k2_tool_parser.reset_streaming_state()
498
+
499
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
500
+ tool_call_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
501
+
502
+ # Single delta containing: section_begin + spurious text + tool_call_begin
503
+ combined_text = "<|tool_calls_section_begin|> noise text <|tool_call_begin|>"
504
+
505
+ result = kimi_k2_tool_parser.extract_tool_calls_streaming(
506
+ previous_text="Reasoning ",
507
+ current_text="Reasoning " + combined_text,
508
+ delta_text=combined_text,
509
+ previous_token_ids=[1, 2],
510
+ current_token_ids=[1, 2, section_begin_id, 3, 4, tool_call_begin_id],
511
+ delta_token_ids=[section_begin_id, 3, 4, tool_call_begin_id],
512
+ request=None,
513
+ )
514
+
515
+ # The noise text should NOT leak into content
516
+ # Result should either be None/empty or start tool call parsing
517
+ if result is not None and result.content is not None:
518
+ # If content is returned, it should not contain the noise
519
+ assert "noise text" not in result.content
520
+ assert result.content == "" or result.content.strip() == ""
521
+
522
+
523
+def test_stream_ends_without_section_end_marker(kimi_k2_tool_parser):
524
+ """
525
+ Test that if the stream ends (EOF) without a proper section end marker,
526
+ the parser doesn't leak text, doesn't crash, and resets state cleanly.
527
+ """
528
+ kimi_k2_tool_parser.reset_streaming_state()
529
+
530
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
531
+
532
+ # Enter tool section
533
+ _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
534
+ previous_text="",
535
+ current_text="<|tool_calls_section_begin|>",
536
+ delta_text="<|tool_calls_section_begin|>",
537
+ previous_token_ids=[],
538
+ current_token_ids=[section_begin_id],
539
+ delta_token_ids=[section_begin_id],
540
+ request=None,
541
+ )
542
+ assert kimi_k2_tool_parser.in_tool_section is True
543
+
544
+ # Some content in tool section
545
+ result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
546
+ previous_text="<|tool_calls_section_begin|>",
547
+ current_text="<|tool_calls_section_begin|> partial content",
548
+ delta_text=" partial content",
549
+ previous_token_ids=[section_begin_id],
550
+ current_token_ids=[section_begin_id, 10, 11],
551
+ delta_token_ids=[10, 11],
552
+ request=None,
553
+ )
554
+ # Content should be suppressed
555
+ assert result2.content == "" or result2.content is None
556
+
557
+ # Stream ends (EOF) - no more deltas, no section_end marker
558
+ # Simulate this by manually checking state and resetting
559
+ # (In real usage, the request handler would call reset_streaming_state)
560
+ assert kimi_k2_tool_parser.in_tool_section is True # Still in section
561
+
562
+ # Reset state (as would happen between requests)
563
+ kimi_k2_tool_parser.reset_streaming_state()
564
+
565
+ # Verify clean slate
566
+ assert kimi_k2_tool_parser.in_tool_section is False
567
+ assert kimi_k2_tool_parser.token_buffer == ""
568
+
569
+ # Next request should work normally
570
+ result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
571
+ previous_text="",
572
+ current_text="New reasoning",
573
+ delta_text="New reasoning",
574
+ previous_token_ids=[],
575
+ current_token_ids=[20, 21],
576
+ delta_token_ids=[20, 21],
577
+ request=None,
578
+ )
579
+ assert result3 is not None
580
+ assert result3.content == "New reasoning"
581
+
582
+
583
+def test_same_chunk_begin_and_end_markers(kimi_k2_tool_parser):
584
+ """
585
+ CRITICAL TEST: Verify that when both section_begin and section_end
586
+ markers appear in the SAME chunk, the parser correctly:
587
+ 1. Enters the tool section
588
+ 2. Immediately exits the tool section
589
+ 3. Does NOT get stuck in in_tool_section=True state
590
+
591
+ This tests the bug fix where elif was changed to if to handle
592
+ both state transitions in a single delta.
593
+ """
594
+ kimi_k2_tool_parser.reset_streaming_state()
595
+
596
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
597
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
598
+
599
+ # Single chunk with both markers (e.g., empty tool section)
600
+ combined_delta = "<|tool_calls_section_begin|><|tool_calls_section_end|>"
601
+
602
+ result = kimi_k2_tool_parser.extract_tool_calls_streaming(
603
+ previous_text="Some reasoning ",
604
+ current_text="Some reasoning " + combined_delta,
605
+ delta_text=combined_delta,
606
+ previous_token_ids=[1, 2],
607
+ current_token_ids=[1, 2, section_begin_id, section_end_id],
608
+ delta_token_ids=[section_begin_id, section_end_id],
609
+ request=None,
610
+ )
611
+
612
+ # CRITICAL: Parser should NOT be stuck in tool section
613
+ assert kimi_k2_tool_parser.in_tool_section is False, (
614
+ "Parser stuck in tool section after processing both begin/end in same chunk. "
615
+ "This indicates the elif bug was not fixed."
616
+ )
617
+
618
+ # Result should be empty or contain only stripped content
619
+ assert result is not None
620
+ assert result.content == "" or result.content is None
621
+
622
+ # Verify subsequent content streams correctly (not suppressed)
623
+ result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
624
+ previous_text="Some reasoning " + combined_delta,
625
+ current_text="Some reasoning " + combined_delta + " More reasoning",
626
+ delta_text=" More reasoning",
627
+ previous_token_ids=[1, 2, section_begin_id, section_end_id],
628
+ current_token_ids=[1, 2, section_begin_id, section_end_id, 10, 11],
629
+ delta_token_ids=[10, 11],
630
+ request=None,
631
+ )
632
+
633
+ # This content should NOT be suppressed (we're out of tool section)
634
+ assert result2 is not None
635
+ assert result2.content == " More reasoning"
636
+
637
+
638
+def test_same_chunk_begin_content_end_markers(kimi_k2_tool_parser):
639
+ """
640
+ Test the same-chunk scenario with actual content between markers.
641
+ Example: <|tool_calls_section_begin|> text <|tool_calls_section_end|>
642
+ all arriving in one delta. The key is that the state machine correctly
643
+ transitions in and out within the same chunk.
644
+ """
645
+ kimi_k2_tool_parser.reset_streaming_state()
646
+
647
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
648
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
649
+
650
+ # Chunk with begin, some whitespace/noise, and end all together
651
+ # This simulates a tool section that opens and closes in the same chunk
652
+ combined_delta = "<|tool_calls_section_begin|> <|tool_calls_section_end|>"
653
+
654
+ _result = kimi_k2_tool_parser.extract_tool_calls_streaming(
655
+ previous_text="Reasoning ",
656
+ current_text="Reasoning " + combined_delta,
657
+ delta_text=combined_delta,
658
+ previous_token_ids=[1],
659
+ current_token_ids=[1, section_begin_id, 100, section_end_id],
660
+ delta_token_ids=[section_begin_id, 100, section_end_id],
661
+ request=None,
662
+ )
663
+
664
+ # Parser should exit cleanly (not stuck in tool section)
665
+ assert kimi_k2_tool_parser.in_tool_section is False
666
+
667
+ # Verify the fix: next content should stream normally, not be suppressed
668
+ result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
669
+ previous_text="Reasoning " + combined_delta,
670
+ current_text="Reasoning " + combined_delta + " Done",
671
+ delta_text=" Done",
672
+ previous_token_ids=[1, section_begin_id, 100, section_end_id],
673
+ current_token_ids=[1, section_begin_id, 100, section_end_id, 200],
674
+ delta_token_ids=[200],
675
+ request=None,
676
+ )
677
+
678
+ # Content after section should be returned (not suppressed)
679
+ assert result2 is not None
680
+ assert result2.content == " Done"
681
+
682
+
683
+def test_tool_call_end_and_section_end_same_chunk(kimi_k2_tool_parser):
684
+ """
685
+ CRITICAL TEST (P1): Verify that when both <|tool_call_end|> and
686
+ <|tool_calls_section_end|> appear in the SAME chunk, the parser:
687
+ 1. Processes the tool_call_end first (emits final arguments)
688
+ 2. THEN exits the section
689
+ 3. Does NOT drop the final tool call update
690
+ 4. Does NOT leak special tokens into reasoning
691
+
692
+ This tests the deferred section exit fix.
693
+ """
694
+ kimi_k2_tool_parser.reset_streaming_state()
695
+
696
+ section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")
697
+ section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")
698
+ tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")
699
+ tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")
700
+
701
+ # Simulate a streaming sequence for a SHORT tool call (all in one chunk):
702
+ # 1. Reasoning text
703
+ result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(
704
+ previous_text="",
705
+ current_text="Let me help. ",
706
+ delta_text="Let me help. ",
707
+ previous_token_ids=[],
708
+ current_token_ids=[1, 2],
709
+ delta_token_ids=[1, 2],
710
+ request=None,
711
+ )
712
+ assert result1 is not None
713
+ assert result1.content == "Let me help. "
714
+
715
+ # 2. Section begin
716
+ _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(
717
+ previous_text="Let me help. ",
718
+ current_text="Let me help. <|tool_calls_section_begin|>",
719
+ delta_text="<|tool_calls_section_begin|>",
720
+ previous_token_ids=[1, 2],
721
+ current_token_ids=[1, 2, section_begin_id],
722
+ delta_token_ids=[section_begin_id],
723
+ request=None,
724
+ )
725
+ assert kimi_k2_tool_parser.in_tool_section is True
726
+
727
+ # 3. Tool call begin + full content + tool_end + section_end ALL IN ONE CHUNK
728
+ # This is the critical scenario for short tool calls
729
+ combined = (
730
+ '<|tool_call_begin|>get_weather:0 <|tool_call_argument_begin|> {"city": "Paris"} '
731
+ "<|tool_call_end|><|tool_calls_section_end|>"
732
+ )
733
+
734
+ # Build up the previous text gradually to simulate realistic streaming
735
+ prev_text = "Let me help. <|tool_calls_section_begin|>"
736
+ curr_text = prev_text + combined
737
+
738
+ result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(
739
+ previous_text=prev_text,
740
+ current_text=curr_text,
741
+ delta_text=combined,
742
+ previous_token_ids=[1, 2, section_begin_id],
743
+ current_token_ids=[
744
+ 1,
745
+ 2,
746
+ section_begin_id,
747
+ tool_begin_id,
748
+ 10,
749
+ 11,
750
+ 12,
751
+ tool_end_id,
752
+ section_end_id,
753
+ ],
754
+ delta_token_ids=[tool_begin_id, 10, 11, 12, tool_end_id, section_end_id],
755
+ request=None,
756
+ )
757
+
758
+ # CRITICAL: Parser should have exited section AFTER processing tool
759
+ assert kimi_k2_tool_parser.in_tool_section is False
760
+
761
+ # Tool call should have been emitted (not dropped)
762
+ # The result might be the tool name or None depending on state, but
763
+ # importantly, it shouldn't be returning the literal tokens as content
764
+
765
+ if result3 is not None and result3.content is not None:
766
+ # Verify no special tokens leaked into content
767
+ assert "<|tool_call_end|>" not in result3.content
768
+ assert "<|tool_calls_section_end|>" not in result3.content
769
+
770
+ # 4. Verify subsequent content streams normally
771
+ result4 = kimi_k2_tool_parser.extract_tool_calls_streaming(
772
+ previous_text=curr_text,
773
+ current_text=curr_text + " Done",
774
+ delta_text=" Done",
775
+ previous_token_ids=[
776
+ 1,
777
+ 2,
778
+ section_begin_id,
779
+ tool_begin_id,
780
+ 10,
781
+ 11,
782
+ 12,
783
+ tool_end_id,
784
+ section_end_id,
785
+ ],
786
+ current_token_ids=[
787
+ 1,
788
+ 2,
789
+ section_begin_id,
790
+ tool_begin_id,
791
+ 10,
792
+ 11,
793
+ 12,
794
+ tool_end_id,
795
+ section_end_id,
796
+ 20,
797
+ ],
798
+ delta_token_ids=[20],
799
+ request=None,
800
+ )
801
+
802
+ # Content after tool section should stream normally
803
+ assert result4 is not None
804
+ assert result4.content == " Done"
vllm/entrypoints/openai/tool_parsers/kimi_k2_tool_parser.py
+197
-5
@@ -34,8 +34,27 @@ class KimiK2ToolParser(ToolParser):
34
str
35
] = [] # map what has been streamed for each tool so far to a list
36
37
+ # Section-level state management to prevent token leakage
38
+ self.in_tool_section: bool = False
39
+ self.token_buffer: str = ""
40
+ # Buffer size: empirical worst-case for longest marker (~30 chars) * 2
41
+ # + safety margin for unicode + partial overlap. Prevents unbounded growth.
42
+ self.buffer_max_size: int = 1024
43
+ self.section_char_count: int = 0 # Track characters processed in tool section
44
+ self.max_section_chars: int = 8192 # Force exit if section exceeds this
45
+ self._buffer_overflow_logged: bool = False # Log overflow once per session
46
+
47
+ # Support both singular and plural variants
48
self.tool_calls_start_token: str = "<|tool_calls_section_begin|>"
49
self.tool_calls_end_token: str = "<|tool_calls_section_end|>"
50
+ self.tool_calls_start_token_variants: list[str] = [
51
+ "<|tool_calls_section_begin|>",
52
+ "<|tool_call_section_begin|>", # singular variant
53
+ ]
54
+ self.tool_calls_end_token_variants: list[str] = [
55
+ "<|tool_calls_section_end|>",
56
+ "<|tool_call_section_end|>", # singular variant
57
+ ]
58
59
self.tool_call_start_token: str = "<|tool_call_begin|>"
60
self.tool_call_end_token: str = "<|tool_call_end|>"
@@ -58,6 +77,18 @@ class KimiK2ToolParser(ToolParser):
77
self.tool_calls_start_token_id = self.vocab.get(self.tool_calls_start_token)
78
self.tool_calls_end_token_id = self.vocab.get(self.tool_calls_end_token)
79
80
+ # Get token IDs for all variants
81
+ self.tool_calls_start_token_ids: list[int] = [
82
+ tid
83
+ for variant in self.tool_calls_start_token_variants
84
+ if (tid := self.vocab.get(variant)) is not None
85
+ ]
86
+ self.tool_calls_end_token_ids: list[int] = [
87
+ tid
88
+ for variant in self.tool_calls_end_token_variants
89
+ if (tid := self.vocab.get(variant)) is not None
90
+ ]
91
+
92
self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
93
self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
94
@@ -70,6 +101,51 @@ class KimiK2ToolParser(ToolParser):
101
"tokens in the tokenizer!"
102
)
103
104
+ def _check_and_strip_markers(self, text: str) -> tuple[str, bool, bool]:
105
+ """
106
+ Check for section begin/end markers in text and strip them.
107
+ Returns: (cleaned_text, found_section_begin, found_section_end)
108
+ """
109
+ found_begin = False
110
+ found_end = False
111
+ cleaned = text
112
+
113
+ # Check for section begin markers (any variant)
114
+ for variant in self.tool_calls_start_token_variants:
115
+ if variant in cleaned:
116
+ cleaned = cleaned.replace(variant, "")
117
+ found_begin = True
118
+
119
+ # Check for section end markers (any variant)
120
+ for variant in self.tool_calls_end_token_variants:
121
+ if variant in cleaned:
122
+ cleaned = cleaned.replace(variant, "")
123
+ found_end = True
124
+
125
+ return cleaned, found_begin, found_end
126
+
127
+ def _reset_section_state(self) -> None:
128
+ """Reset state when exiting tool section."""
129
+ self.in_tool_section = False
130
+ self.token_buffer = ""
131
+ self.section_char_count = 0
132
+
133
+ def reset_streaming_state(self) -> None:
134
+ """
135
+ Reset all streaming state. Call this between requests to prevent
136
+ state leakage when parser instance is reused.
137
+ """
138
+ # Reset section state
139
+ self._reset_section_state()
140
+
141
+ # Reset parent class state
142
+ self.current_tool_name_sent = False
143
+ self.prev_tool_call_arr = []
144
+ self.current_tool_id = -1
145
+ self.streamed_args_for_tool = []
146
+
147
+ logger.debug("Streaming state reset")
148
+
149
def extract_tool_calls(
150
self,
151
model_output: str,
@@ -131,13 +207,94 @@ class KimiK2ToolParser(ToolParser):
207
) -> DeltaMessage | None:
208
logger.debug("delta_text: %s", delta_text)
209
logger.debug("delta_token_ids: %s", delta_token_ids)
134
- # check to see if we should be streaming a tool call - is there a
135
- if self.tool_calls_start_token_id not in current_token_ids:
210
+
211
+ # Flag to defer section exit until after tool parsing completes
212
+ deferred_section_exit = False
213
+
214
+ # Add delta to buffer for split marker detection
215
+ self.token_buffer += delta_text
216
+
217
+ # Enforce buffer size limit to prevent memory issues
218
+ if len(self.token_buffer) > self.buffer_max_size:
219
+ if not self._buffer_overflow_logged:
220
+ logger.warning(
221
+ "Token buffer exceeded max size (%d bytes), flushing excess. "
222
+ "This may indicate very long markers or unusual tokenization.",
223
+ self.buffer_max_size,
224
+ )
225
+ self._buffer_overflow_logged = True
226
+ # Keep only the most recent content that might contain partial markers
227
+ self.token_buffer = self.token_buffer[-self.buffer_max_size // 2 :]
228
+
229
+ # Check buffer for section markers (handles split tokens)
230
+ buffered_text, found_section_begin, found_section_end = (
231
+ self._check_and_strip_markers(self.token_buffer)
232
+ )
233
+
234
+ # Track section state transitions
235
+ if found_section_begin and not self.in_tool_section:
236
+ logger.debug("Entering tool section")
237
+ self.in_tool_section = True
238
+ self.token_buffer = buffered_text # Use cleaned buffer
239
+ self.section_char_count = 0 # Reset counter for new section
240
+ if found_section_end and self.in_tool_section:
241
+ logger.debug("Detected section end marker")
242
+ # CRITICAL: Don't exit early if tool_call_end is in this chunk.
243
+ # Tool parser must emit final arguments/close first to avoid dropping
244
+ # the final tool update and leaking tokens into reasoning channel.
245
+ has_tool_end = self.tool_call_end_token_id in delta_token_ids
246
+ if has_tool_end:
247
+ # Defer exit until after tool parsing completes
248
+ deferred_section_exit = True
249
+ logger.debug("Deferring section exit: tool_call_end in same chunk")
250
+ self.token_buffer = buffered_text
251
+ else:
252
+ # No tool call ending, safe to exit immediately
253
+ logger.debug("Exiting tool section")
254
+ remaining = buffered_text
255
+ self._reset_section_state()
256
+ # Return remaining text as reasoning content if non-empty
257
+ if remaining.strip():
258
+ return DeltaMessage(content=remaining)
259
+ # Return empty delta to maintain function contract
260
+ # (always returns DeltaMessage)
261
+ return DeltaMessage(content="")
262
+ else:
263
+ self.token_buffer = buffered_text
264
+
265
+ # Check if any variant of section start token is in current_token_ids
266
+ has_section_token = any(
267
+ tid in current_token_ids for tid in self.tool_calls_start_token_ids
268
+ )
269
+
270
+ # Early return: if no section token detected yet, return as reasoning content
271
+ if not has_section_token and not self.in_tool_section:
272
logger.debug("No tool call tokens found!")
273
+ # Don't clear buffer - it needs to accumulate partial markers across deltas
274
+ # Buffer overflow is already protected by lines 215-224
275
return DeltaMessage(content=delta_text)
138
- delta_text = delta_text.replace(self.tool_calls_start_token, "").replace(
139
- self.tool_calls_end_token, ""
140
- )
276
+
277
+ # Strip section markers from delta_text for subsequent processing
278
+ # NOTE: This preprocessing happens BEFORE the regex-based tool call
279
+ # parsing (from PR #24847) to ensure markers are removed cleanly
280
+ # before pattern matching. No double-stripping occurs because
281
+ # section markers and tool call markers are distinct.
282
+ delta_text, _, _ = self._check_and_strip_markers(delta_text)
283
+
284
+ # Error recovery: If in tool section for too long, force exit
285
+ if self.in_tool_section:
286
+ self.section_char_count += len(delta_text)
287
+ if self.section_char_count > self.max_section_chars:
288
+ logger.warning(
289
+ "Tool section exceeded max length (%d chars), forcing exit. "
290
+ "This may indicate malformed model output.",
291
+ self.max_section_chars,
292
+ )
293
+ self._reset_section_state()
294
+ # Deferred exit already handled by forced exit above
295
+ # Return remaining content as reasoning (or empty delta if no content)
296
+ return DeltaMessage(content=delta_text if delta_text.strip() else "")
297
+
298
try:
299
# figure out where we are in the parsing by counting tool call
300
# start & end tags
@@ -158,6 +315,16 @@ class KimiK2ToolParser(ToolParser):
315
and prev_tool_end_count == cur_tool_end_count
316
and self.tool_call_end_token not in delta_text
317
):
318
+ # CRITICAL FIX: Suppress content if in tool section but
319
+ # no tool calls started
320
+ if self.in_tool_section and cur_tool_start_count == 0:
321
+ logger.debug(
322
+ "In tool section but no tool calls started yet. "
323
+ "Suppressing: %s",
324
+ delta_text,
325
+ )
326
+ # Return empty delta to maintain iterator contract
327
+ return DeltaMessage(content="")
328
logger.debug("Generating text content! skipping tool parsing.")
329
return DeltaMessage(content=delta_text)
330
@@ -209,6 +376,9 @@ class KimiK2ToolParser(ToolParser):
376
):
377
if self.prev_tool_call_arr is None or len(self.prev_tool_call_arr) == 0:
378
logger.debug("attempting to close tool call, but no tool call")
379
+ # Handle deferred section exit before returning
380
+ if deferred_section_exit and self.in_tool_section:
381
+ self._reset_section_state()
382
return None
383
diff = self.prev_tool_call_arr[self.current_tool_id].get("arguments")
384
if diff:
@@ -218,6 +388,9 @@ class KimiK2ToolParser(ToolParser):
388
else diff
389
)
390
if '"}' not in delta_text:
391
+ # Handle deferred section exit before returning
392
+ if deferred_section_exit and self.in_tool_section:
393
+ self._reset_section_state()
394
return None
395
end_loc = delta_text.rindex('"}')
396
diff = delta_text[:end_loc] + '"}'
@@ -227,6 +400,10 @@ class KimiK2ToolParser(ToolParser):
400
diff,
401
)
402
self.streamed_args_for_tool[self.current_tool_id] += diff
403
+ # Handle deferred section exit before returning
404
+ if deferred_section_exit and self.in_tool_section:
405
+ logger.debug("Completing deferred section exit")
406
+ self._reset_section_state()
407
return DeltaMessage(
408
tool_calls=[
409
DeltaToolCall(
@@ -240,9 +417,19 @@ class KimiK2ToolParser(ToolParser):
417
418
# case -- otherwise we're just generating text
419
else:
420
+ # Check if we're in tool section - if so, suppress
421
+ if self.in_tool_section:
422
+ logger.debug("In tool section, suppressing text generation")
423
+ # Handle deferred section exit before returning
424
+ if deferred_section_exit:
425
+ self._reset_section_state()
426
+ return DeltaMessage(content="")
427
text = delta_text.replace(self.tool_call_start_token, "")
428
text = text.replace(self.tool_call_end_token, "")
429
delta = DeltaMessage(tool_calls=[], content=text)
430
+ # Handle deferred section exit before returning
431
+ if deferred_section_exit and self.in_tool_section:
432
+ self._reset_section_state()
433
return delta
434
435
current_tool_call = dict()
@@ -390,6 +577,11 @@ class KimiK2ToolParser(ToolParser):
577
else:
578
self.prev_tool_call_arr.append(current_tool_call)
579
580
+ # Handle deferred section exit after tool parsing completes
581
+ if deferred_section_exit and self.in_tool_section:
582
+ logger.debug("Completing deferred section exit")
583
+ self._reset_section_state()
584
+
585
return delta
586
587
except Exception: