fix(parsers): recover Kimi K2 tool calls when section_end is missing

[keivenchang]

Overview:

Kimi K2 tool calls were getting silently dropped whenever the model hit max_tokens before emitting the section_end marker. The parser just threw away everything if that closing tag was missing -- same bug exists in Moonshot's reference Python parser and got copied into vLLM and SGLang too.

Details:

• the root cause is in extract_tool_calls (kimi_k2_parser.rs) -- the old code discarded the entire tool call section when section_end wasnt there. now it treats missing section_end as "section goes to EOF" and pulls out whatever complete individual calls it can find
• the streaming jail (jail.rs) had a related issue: find_tool_call_end_position_kimi_k2 couldn't tell "found section_end" apart from "section_end is missing", so it would early-exit and swallow parallel calls. changed the return type to Option<usize> so None skips the early-exit path, and finalize() recovers all the calls at stream end
• 4 new parser unit tests, 3 new streaming integration tests that reproduce the customer-reported scenario (DIS-1765)

Where should the reviewer start?

docs/agents/kimi_k2.md, then kimi_k2_parser.rs

Related Issues:

Relates to DIS-1765

/coderabbit profile chill

[coderabbitai]

Walkthrough

Refactors the Kimi K2 tool-call parser to handle truncated output gracefully by making the section-end marker optional. Updates the parser API to return Option<usize> instead of usize, modifying underlying implementations and callers to treat missing section-end markers as incomplete rather than errors.

Changes

Cohort / File(s)	Summary
Documentation docs/agents/kimi_k2.md	New documentation page describing the tool-call truncation bug and documenting the Python regex fix and Rust-side implementation status, including TODO items to upstream changes.
Kimi K2 Parser lib/parsers/src/tool_calling/xml/kimi_k2_parser.rs	Changed find_tool_call_end_position_kimi_k2 return type from usize to Option<usize>; returns None when section_end marker is missing. Updated extract_tool_calls to extract complete tool calls from truncated output lacking section_end.
Parser API lib/parsers/src/tool_calling/parsers.rs	Updated find_tool_call_end_position return type from usize to Option<usize>. All parser branches now wrap results in Some(...). Added doc comment clarifying that None indicates unclosed/incomplete tool-call sections.
Stream Handling lib/llm/src/protocols/openai/chat_completions/jail.rs	Modified JailedStream::should_end_jail early-exit logic to treat find_tool_call_end_position as optional; returns (false, accumulated_content.len()) when None is returned instead of unconditionally splitting.
Test Coverage lib/llm/tests/test_streaming_tool_parsers.rs	Added make_chunk test helper and three integration-style tests: baseline case with complete markers, DIS-1765 reproduction with truncated output (no section_end), and multi-tool-call variant with truncation.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically summarizes the main change: fixing Kimi K2 tool call recovery when the section_end marker is missing.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	PR description follows the template with all required sections: Overview explains the bug clearly, Details describe specific changes and affected files, Where should the reviewer start identifies key files, and Related Issues references DIS-1765.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/llm/src/protocols/openai/chat_completions/jail.rs (1)

803-824: ⚠️ Potential issue | 🟠 Major

Preserve the terminal finish reason on the new finalize() recovery path.

These Kimi None cases now intentionally defer emission to finalize(), but the parsed tool-call branch in create_tool_call_choice() always builds the emitted chunk with finish_reason: None. For a truncated stream ending with FinishReason::Length, the recovered tool call comes out without any terminal reason, and fix_finish_reason() cannot restore it because it only rewrites existing Some(Stop) values. Please thread base_choice.finish_reason through the successful parsed-tool-call path so recovered Kimi responses still terminate with Length (or get rewritten from Stop to ToolCalls).

Suggested fix

-                    let choice = create_choice_stream(
-                        choice_index,
-                        Some(Role::Assistant),
-                        normal_text.as_deref().unwrap_or(""),
-                        Some(tool_call_chunks),
-                        None,
-                        None,
-                        None,
-                    );
+                    let choice = create_choice_stream(
+                        choice_index,
+                        Some(Role::Assistant),
+                        normal_text.as_deref().unwrap_or(""),
+                        Some(tool_call_chunks),
+                        base_choice.finish_reason,
+                        base_choice.stop_reason.clone(),
+                        base_choice.logprobs.clone(),
+                    );

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@lib/llm/src/protocols/openai/chat_completions/jail.rs` around lines 803 -
824, create_tool_call_choice() currently constructs the emitted tool-call chunk
with finish_reason: None causing recovered Kimi tool-call responses (from
finalize()) to lose terminal reasons like FinishReason::Length; update the
parsed-tool-call branch in create_tool_call_choice() to propagate the original
base_choice.finish_reason into the emitted chunk (use base_choice.finish_reason
instead of None) so finalize() recovers tool calls with the correct terminal
reason and fix_finish_reason() can continue to rewrite Stop→ToolCalls as before;
ensure any helper paths that build the emitted choice (the parsed tool-call
success branch where try_tool_call_parse_aggregate(...) and
find_tool_call_end_position(...) return) use base_choice.finish_reason and keep
existing behavior for non-parsed branches.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@lib/llm/src/protocols/openai/chat_completions/jail.rs`:
- Around line 803-824: create_tool_call_choice() currently constructs the
emitted tool-call chunk with finish_reason: None causing recovered Kimi
tool-call responses (from finalize()) to lose terminal reasons like
FinishReason::Length; update the parsed-tool-call branch in
create_tool_call_choice() to propagate the original base_choice.finish_reason
into the emitted chunk (use base_choice.finish_reason instead of None) so
finalize() recovers tool calls with the correct terminal reason and
fix_finish_reason() can continue to rewrite Stop→ToolCalls as before; ensure any
helper paths that build the emitted choice (the parsed tool-call success branch
where try_tool_call_parse_aggregate(...) and find_tool_call_end_position(...)
return) use base_choice.finish_reason and keep existing behavior for non-parsed
branches.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b62536c8-0684-494a-b37f-c0a1eaf514d8

📥 Commits

Reviewing files that changed from the base of the PR and between 45364b5 and c0f12a8.

📒 Files selected for processing (5)

• docs/agents/kimi_k2.md
• lib/llm/src/protocols/openai/chat_completions/jail.rs
• lib/llm/tests/test_streaming_tool_parsers.rs
• lib/parsers/src/tool_calling/parsers.rs
• lib/parsers/src/tool_calling/xml/kimi_k2_parser.rs