Conversation
Signed-off-by: lightseek-bot <243258330+lightseek-bot@users.noreply.github.com> Co-authored-by: Keyang Ru <rukeyang@gmail.com> Co-authored-by: Chen Hongtao <56470055+chenht2022@users.noreply.github.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthroughKimi-K3 support was added across native XTML rendering, tokenizer detection, reasoning and tool parsing, multimodal routing, and gateway handling. The change includes streaming state management, tool and response controls, model mappings, tool-call IDs, and comprehensive tests and fixtures. ChangesKimi-K3 XTML rendering
Kimi-K3 reasoning and tool parsing
Model routing and gateway integration
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
| let is_k3 = model | ||
| .as_bytes() | ||
| .windows(2) // "k3".len() | ||
| .any(|window| window.eq_ignore_ascii_case(b"k3")); |
There was a problem hiding this comment.
🟡 Nit: The 2-byte windows check for "k3" scans the entire model string, so any model whose name incidentally contains the ASCII substring k3 in a non-version position (e.g. a hypothetical kimi-7k3000 or kimi_k32) would be routed to the K3 id format.
A slightly tighter guard — anchoring on a word boundary or checking that the character after k3 is not a digit — would prevent those false positives while still matching kimi-k3, Kimi-K3-Chat, kimi_k3, etc. Not blocking because the existing K2 check has the same pattern and the false-positive space is tiny, but worth hardening if more Kimi generations land.
There was a problem hiding this comment.
Fair point, and agreed it's non-blocking. The windows(2) "k3" scan could match a substring in a hypothetical id like kimi_k32. I'm leaving it for now to keep this change focused: the K2 path (windows(4) "kimi") shares the same substring-scan shape, so a tightening should cover both symmetrically, and no such model id exists today. Noting it as a small follow-up hardening rather than addressing it here.
| self.buffer.push_str(chunk); | ||
| let current_text = self.buffer.clone(); |
There was a problem hiding this comment.
🟡 Nit: self.buffer.clone() allocates a full copy of the accumulated buffer on every streaming chunk. Because the buffer only grows, this is O(n²) total allocation over a response of length n. The clone is needed to work around the borrow conflict with extract_response_content(&mut self, &str), but you could eliminate it by splitting the mutable state (sent_content_idx, sent_tool_call_count) out of self or by inlining the content-extraction logic directly. Not urgent — the per-token cost is bounded by response length — but worth revisiting if K3 serves long tool-heavy responses.
There was a problem hiding this comment.
Same root cause as the O(n²) rescan flagged separately — the per-chunk buffer.clone() plus full re-scan is quadratic. I'm deferring both to the same follow-up. Note the clone isn't the dominant term (the call_re/arg_re/attr_re rescans are), so removing just the clone wouldn't change the complexity class; the proper fix is a scan cursor, which I'd rather do carefully in a dedicated change. Bounded by DEFAULT_MAX_BUFFER_SIZE meanwhile.
There was a problem hiding this comment.
Thorough review of the Kimi-K3 support across all five subsystems (multimodal, reasoning parser, tokenizer/XTML renderer, tool parser, gRPC gateway). The implementation is well-structured and well-tested.
Summary: 0 🔴 Important · 2 🟡 Nit · 0 🟣 Pre-existing
Both nits are minor optimization/robustness suggestions — nothing blocking.
Highlights:
- Reasoning parser streaming logic (overlap detection, partial marker holdback) is correctly implemented
- XTML renderer faithfully ports the Python reference with good documentation of known limitations
- Tool parser handles both streaming and non-streaming paths with proper marker-split protection
- Test coverage is excellent: golden fixture tests, unit tests for edge cases, factory resolution tests, and end-to-end tokenizer integration tests
- The
tool_choice/response_formatforwarding inbuild_chat_template_kwargsis a broader change (affects all models), but correctly gated — explicitchat_template_kwargsstill win via theextendordering
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 63dab72ee5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| .get("architectures") | ||
| .and_then(|v| v.as_array()) | ||
| .is_some_and(|a| { | ||
| a.iter() | ||
| .any(|v| v.as_str() == Some("KimiK25ForConditionalGeneration")) | ||
| }); | ||
| if is_kimi { | ||
| .map(|a| a.iter().filter_map(|v| v.as_str()).collect()) | ||
| .unwrap_or_default(); | ||
| if arch_strs.contains(&"KimiK3ForConditionalGeneration") { |
There was a problem hiding this comment.
Select XTML renderer for model_type-only K3 configs
When a K3 checkpoint's config.json identifies the model with "model_type": "kimi_k3" but does not include this exact architecture string, this branch falls through to Renderer::Jinja. That shape is now accepted elsewhere in this commit for tokenizer and vision detection, and K3 has no Jinja chat template, so chat rendering will fail or use the wrong format instead of the XTML renderer; please include model_type in this detection path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 851f389. detect_renderer_from_config now falls back to config.json model_type when architectures is absent or unrecognized: kimi_k3 → Renderer::KimiK3Xtml, kimi_k25 → Renderer::KimiK25Tools. This keeps renderer selection consistent with is_kimi_tokenizer, which already keys off model_type. Added a test covering the model_type-only config.
| registry.map_model("kimi-k3*", "kimi_k3"); | ||
| registry.map_model("Kimi-K3*", "kimi_k3"); | ||
| registry.map_model("moonshot*/Kimi-K3*", "kimi_k3"); |
There was a problem hiding this comment.
Map underscore Kimi_K3 ids to the K3 tool parser
For model ids that use the underscore form (kimi_k3 or moonshotai/Kimi_K3), auto-detection never reaches the new parser because these mappings only cover hyphenated names. The reasoning parser, vision registry, and tool-call id code all handle the underscore form, but resolve_model_to_parser will fall back to the default passthrough parser here, leaving XTML tool output in content instead of returning tool_calls; add the underscore K3 patterns as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 851f389 — added the underscore-spelled globs (kimi_k3*, Kimi_K3*, moonshot*/Kimi_K3*) alongside the hyphenated ones, matching the parity the reasoning-parser factory already has. Added a resolution test asserting all three resolve to the kimi_k3 parser.
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/reasoning_parser/src/parsers/kimi_k3.rs`:
- Around line 124-128: Update the response-open handling in the parser method
containing this block to remove all occurrences from the text, matching the
existing global replacement behavior of the other markers. Replace the
single-match `response_open_re.find` path while preserving the unchanged-text
behavior when no marker is present.
- Around line 122-146: The parser withholds marker-prefix suffixes in
content_ready_to_emit and reasoning_text_ready_to_emit without ever releasing
them at stream end. Add an explicit end-of-stream finalization path that emits
each residual held-back tail, invoke it when generation completes, and ensure
reset does not discard unflushed text; if the API cannot signal completion,
document the intentional loss instead.
In `@crates/tokenizer/src/encoders/kimi_k3_xtml.rs`:
- Around line 260-326: Update XTML serialization so user-controlled `<|...|>`
marker sequences are escaped in attribute values and string content before
output. Extend escape_attr_value and apply the corresponding content escaping in
push_content and render_assistant_segments, covering tool names, argument
keys/values, message names, and text. Preserve internally generated control
tokens such as OPEN_TOKEN, CLOSE_TOKEN, SEP_TOKEN, and IMAGE_PLACEHOLDER without
escaping.
In `@crates/tokenizer/src/tiktoken.rs`:
- Around line 565-591: Update think_in_prefill() for Renderer::KimiK3Xtml to
return the effective per-request thinking state resolved by apply_kimi_k3_xtml,
using template_kwargs["thinking"] before falling back to params.thinking.
Preserve the existing chat_template.think_in_prefill() behavior for other
renderers and ensure the result reflects whether the current request starts in
reasoning mode.
In `@crates/tool_parser/src/factory.rs`:
- Around line 421-423: Add underscore-form Kimi K3 model mappings alongside the
existing entries in the tool parser factory’s registry setup, including case and
provider-prefix variants as supported by the reasoning factory. Preserve the
existing hyphenated mappings so both model-id styles resolve to “kimi_k3”.
In `@crates/tool_parser/src/parsers/kimi_k3.rs`:
- Around line 365-383: The K3 streaming parsers reprocess their entire
accumulated output on every chunk, causing quadratic work. In
crates/tool_parser/src/parsers/kimi_k3.rs lines 365-383, update
parse_incremental to avoid cloning the full buffer and maintain a cursor past
decoded call blocks while scanning only new content; in
crates/reasoning_parser/src/parsers/kimi_k3.rs lines 217-311, replace whole-tail
current_safe recomputation and repeated replace_all passes with incremental
scanning of newly appended text plus a marker-length lookback, emitting only the
delta.
- Around line 392-402: Update the tool-call mapping around decoded_calls and
ToolCallItem so a model-provided decoded.tool_index is used only when it is
unique and follows the expected ordinal sequence. Otherwise, assign
self.sent_tool_call_count + i, preventing duplicate, sparse, or oversized
indices from colliding with emitted calls while preserving the existing
sent_tool_call_count update.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d071ddc3-04f2-4df9-8ac6-a81324099d4c
📒 Files selected for processing (18)
crates/multimodal/src/registry/kimi_k25.rscrates/multimodal/src/vision/processor.rscrates/reasoning_parser/Cargo.tomlcrates/reasoning_parser/src/factory.rscrates/reasoning_parser/src/lib.rscrates/reasoning_parser/src/parsers/kimi_k3.rscrates/reasoning_parser/src/parsers/mod.rscrates/tokenizer/src/encoders/kimi_k3_xtml.rscrates/tokenizer/src/encoders/mod.rscrates/tokenizer/src/kimi_k2_tokenizer.rscrates/tokenizer/src/tiktoken.rscrates/tokenizer/tests/fixtures/kimi_k3/k3_render_fixtures.jsoncrates/tokenizer/tests/kimi_k3_renderer.rscrates/tool_parser/src/factory.rscrates/tool_parser/src/lib.rscrates/tool_parser/src/parsers/kimi_k3.rscrates/tool_parser/src/parsers/mod.rsmodel_gateway/src/routers/grpc/utils/chat_utils.rs
…indexing - Detect the KimiK3Xtml / KimiK25Tools renderers from config.json `model_type` when `architectures` is absent or unrecognized, matching the existing model_type-based tokenizer detection (is_kimi_tokenizer). - Register underscore-spelled K3 tool-parser model globs (kimi_k3*, Kimi_K3*, moonshot*/Kimi_K3*) for parity with the reasoning-parser factory. - Assign streaming tool indices by emission order instead of trusting the model-supplied XTML `index`, so duplicate/sparse/out-of-order indices can no longer collide; this matches the non-streaming path. Remove the now-unused DecodedCall index plumbing. - Document the intentional response-open prefix-strip asymmetry and the non-flushed partial-marker tail in the reasoning parser, and clarify in the XTML encoder why HTML-escaping control tokens is not the correct fix. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 851f38927f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let mut arguments = serde_json::Map::new(); | ||
| for arg_match in self.arg_re.captures_iter(body) { |
There was a problem hiding this comment.
Decode XTML JSON blocks before emitting tool calls
When K3 emits a call using the supported <|open|>json type="object"<|sep|>... representation, this decoder only scans for argument blocks and therefore silently returns {} for the call arguments. The new renderer itself emits JSON blocks for raw or non-object argument strings in kimi_k3_xtml.rs, so this format can occur in K3 conversations; parse that block into the emitted arguments instead of discarding it.
Useful? React with 👍 / 👎.
| self.buffer.push_str(chunk); | ||
| let current_text = self.buffer.clone(); | ||
|
|
||
| let content = self.extract_response_content(¤t_text); |
There was a problem hiding this comment.
Stop retaining and rescanning the entire K3 stream
For long streamed responses, every delta appends to an ever-growing buffer, clones the complete generation, and then runs the response and tool regexes over it again. With token-sized chunks this makes parsing quadratic in total output length and retains two full copies of the response even after its content was emitted; drain already-consumed response bytes and retain only the incomplete marker/tool-call suffix needed for subsequent chunks.
Useful? React with 👍 / 👎.
Register the kimi_k3 parser with a structural-tag builder so guided
decoding (xgrammar) constrains Kimi-K3 XTML tool calls, mirroring the
existing kimi_k2 support. This closes the gap that let K3 emit invalid
or empty tool calls (schema-invalid args, empty {}, semantically empty).
KimiK3Parser::build_structural_tag emits a triggered_tags grammar for
K3 XTML: a single tools section whose content is a plus of call blocks
(native parallel calls). Each call pins the declared tool name and
index, and every argument is constrained from the JSON schema (Strict):
required properties forced present in declared order, optionals wrapped
in optional, and each type/value pinned per property type -- string
bodies verbatim, integer as number, enums as literal alternates, escaped
attribute values. Schemas without usable properties fall back to a
permissive skeleton so the grammar is never infeasible.
Like kimi_k2, stop_after_first is left unset so the trailing
message/end-of-msg tokens after the tools section remain valid.
Signed-off-by: Keyang Ru <rukeyang@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/tool_parser/src/parsers/kimi_k3.rs`:
- Around line 425-443: The argument grammar built by build_args_format currently
makes optional properties positional, while decode_call accepts them in any
order. Update the sequence construction near lines 425-443 so required
properties retain emitted map order and optional properties are represented by
an order-insensitive star/or region rather than positional optional slots. In
crates/tool_parser/src/parsers/kimi_k3.rs lines 138-140, update the
build_args_format documentation to describe the resulting layout, including
interleaved optionals, or adjust the implementation to match the documented
required-first ordering.
- Around line 425-443: The build_args_format sequence currently makes optional
arguments position-dependent, conflicting with decode_call’s parse-order
handling. Update build_args_format so optional properties are represented in an
order-insensitive structural grammar region while preserving required argument
handling and argument formats; only document a schema-order guarantee instead if
the encoder reliably enforces that order.
- Around line 524-558: Unify marker exclusions in permissive_argument_format and
permissive_call_format with the fixed attribute formats: exclude the complete
OPEN, CLOSE, and SEP markers rather than the partial "<|" prefix, while
retaining the quote exclusion. Ensure permissive attribute values cannot contain
partially formed protocol markers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 410733d3-7e7a-4d8b-bb8e-329c4c1191c1
📒 Files selected for processing (2)
crates/tool_parser/src/factory.rscrates/tool_parser/src/parsers/kimi_k3.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 28ab345dcd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let buffered_size = self.buffer.len() + text.len(); | ||
| if buffered_size > DEFAULT_MAX_BUFFER_SIZE { | ||
| return Err(ParseError::BufferOverflow(buffered_size)); | ||
| } |
There was a problem hiding this comment.
Drain emitted K3 reasoning stream data
When a streamed K3 completion exceeds DEFAULT_MAX_BUFFER_SIZE (4 MiB), this returns BufferOverflow even if nearly all of those bytes were already emitted to the client, because self.buffer is never drained after either reasoning or content deltas are produced. The same retention also makes each incremental call rescan and copy the full generation; retain only the incomplete control-marker suffix and other unconsumed bytes so long streams do not fail based on cumulative output size.
Useful? React with 👍 / 👎.
…r-exclusion intent Correct the build_structural_tag / build_args_format docs: the argument grammar emits one slot per property in the schema's declared order (preserve_order is on) with required properties mandatory and optionals individually skippable in place -- not "required first, then optional" as previously described. Note that this declared order is a generation constraint only; decode_call parses arguments order-insensitively. Also document why the permissive attribute regions exclude the `<|` marker prefix while free-text value regions exclude only the complete OPEN/CLOSE/SEP markers -- a deliberate distinction (quote-delimited identifiers vs marker-delimited free text), not an inconsistency. Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c4754790d2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| if params.add_generation_prompt { | ||
| push_open_tag(&mut out, "message", &[("role", "assistant".to_string())]); | ||
| push_open_tag(&mut out, if thinking { "think" } else { "response" }, &[]); |
There was a problem hiding this comment.
Keep continued assistant prefixes in the response channel
When continue_final_message=true and K3 thinking is left at its default-on setting, chat_utils.rs removes the final assistant message, invokes this renderer with add_generation_prompt=true, and appends that message's content afterward. This line opens the think channel, so the assistant prefix is placed inside reasoning rather than continuing its response, and subsequent generated tokens are classified as reasoning until a close marker appears; K3 continuations need to render the prefix inside response instead of using the normal thinking generation tail.
Useful? React with 👍 / 👎.
| let id = metadata.model_id.to_ascii_lowercase(); | ||
| id.contains("kimi") && id.contains("k2") | ||
| id.contains("kimi") && (id.contains("k2") || id.contains("k3")) | ||
| || metadata | ||
| .config_model_type() | ||
| .is_some_and(|mt| mt == "kimi_k25") | ||
| .is_some_and(|mt| mt == "kimi_k25" || mt == "kimi_k3") |
There was a problem hiding this comment.
Preserve authored media order for K3 prompts
For a K3 request whose content interleaves text and images, routing it through the K2.5 spec also inherits ModelProcessorSpec's MediaFirst ordering. The gRPC preparation path consequently collapses the content array into all <|media_pad|> anchors followed by all text before the K3 renderer sees it, even though push_content in the new XTML renderer processes parts positionally; prompts such as text-image-text therefore change meaning. K3 needs an authored-order override, likely through a distinct spec rather than sharing this one unchanged.
Useful? React with 👍 / 👎.
…ng test Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5eb832edb0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let m_open = self.response_open_re.find(current_text); | ||
| let body_start = m_open.map_or(0, |m| m.end()); |
There was a problem hiding this comment.
Wait for the response channel before streaming content
For an OpenAI Chat streaming request with tools, K3's default thinking enabled, and the default separate_reasoning=false, the model output begins with reasoning text because the generation prompt already consumed the think-open marker. Until the response-open marker arrives, m_open is None, so this sets body_start to zero and the no-marker branch emits each reasoning chunk as ordinary content; the non-streaming path instead extracts only the response channel. Buffer or skip the think channel until the response channel starts so streaming does not expose reasoning and diverge from non-streaming output.
Useful? React with 👍 / 👎.
| if properties.is_empty() { | ||
| return permissive_args_format(); |
There was a problem hiding this comment.
Honor closed no-argument schemas in structural constraints
When a no-argument tool declares {"type":"object","properties":{},"additionalProperties":false}, this converts its argument region into a permissive star of arbitrary argument blocks. Guided decoding can therefore produce parameters explicitly forbidden by the tool schema, and the parser forwards them to the caller. Treat an explicitly closed empty-object schema as an empty argument sequence, reserving the permissive fallback for schemas that actually allow unknown properties.
Useful? React with 👍 / 👎.
Description
Problem
SMG had no support for the Kimi-K3 model family. K3 uses a distinct XTML control-token format —
<|open|>/<|close|>/<|sep|>framing across thethink,response,message, andtoolschannels — for prompt rendering, reasoning, and tool calls. None of the existing renderers or parsers understood this format, so requests targeting K3 models could not be rendered, and K3 output could not be turned back into reasoning content or tool calls.Solution
Add end-to-end Kimi-K3 support across the tokenizer, reasoning parser, tool parser, and the multimodal / gRPC paths:
config.json(botharchitecturesand, as a fallback,model_type).thinkchannel into reasoning content and routes the remainder to normal content, honoring the model's thinking state.tools/responsechannels back intoToolCalls and content, plus a guided-decoding (xgrammar) structural tag that constrains K3 XTML tool calls to the declared tools.tool_choiceandresponse_formatto the backend and apply K3 tool-call id formatting.Changes
kimi_k3_xtmlencoder and register it in the encoder registry; extend renderer detection intiktoken.rsto map K3architectures/model_typeto the XTML renderer.kimi_k3reasoning parser and register it in the factory with its model-pattern mappings.kimi_k3tool parser (non-streaming + streaming) with a structural-tag builder for guided decoding, registered in the tool-parser factory with model globs.tool_choice/response_formatand format K3 tool-call ids in the gRPC chat utilities.Test Plan
All new and existing tests in the touched crates pass.
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
tool-choiceandresponse-format, and adjusted Kimi-K3 tool-call ID formatting.