@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ab98034d4c

[Frontend][Bugfix] Use default tool call IDs for Kimi K3 for conversation-level uniqueness (#50420)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jul 31, 2026 at 10:02 UTC ab98034d4ccccd62e814ebcdc3f3831586b1c8d7

5 files changed +55 -53

rust/src/chat/src/renderer/kimi_k3/encoding.rs

+2 -2

@@ -126,7 +126,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -

126 resolved.sort_by_key(|item| (item.0, item.1));

127 }

128

129 - for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() {

129 + for (position, _, content, name) in resolved {

130 let tool_name = name.as_deref().ok_or_else(|| {

131 Error::ChatTemplate(

132 "Kimi K3 tool messages need a resolvable tool name: \

@@ -135,7 +135,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -

135 .to_string(),

136 )

137 })?;

138 - write_tool_message(out, tool_name, xtml_index + 1, &content)?;

138 + write_tool_message(out, tool_name, position, &content)?;

139 }

140 Ok(())

141 };

rust/src/chat/src/renderer/kimi_k3/tests.rs

+30 -1

@@ -14,11 +14,11 @@ use vllm_tokenizer::Tokenizer;

14 use vllm_tokenizer::test_utils::TestTokenizer;

16 use super::KimiK3ChatRenderer;

17 -use crate::AssistantContentBlock;

17 use crate::ChatRenderer;

18 use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP};

19 use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request};

20 use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort};

21 +use crate::{AssistantContentBlock, AssistantToolCall};

23 const OPEN_ID: u32 = 256;

24 const CLOSE_ID: u32 = 257;

@@ -178,6 +178,35 @@ fn non_thinking_history_omits_reasoning_channel() {

178 ));

179 }

180

181 +#[test]

182 +fn partial_tool_results_keep_assistant_call_position() {

183 + let mut request = crate::request::ChatRequest::for_test();

184 + request.messages = vec![

185 + ChatMessage::assistant_blocks(vec![

186 + AssistantContentBlock::ToolCall(AssistantToolCall {

187 + id: "call-a".to_string(),

188 + name: "weather".to_string(),

189 + arguments: "{}".to_string(),

190 + }),

191 + AssistantContentBlock::ToolCall(AssistantToolCall {

192 + id: "call-b".to_string(),

193 + name: "search".to_string(),

194 + arguments: "{}".to_string(),

195 + }),

196 + ]),

197 + ChatMessage::tool_response("result-b", "call-b"),

198 + ];

199 + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;

200 +

201 + let rendered = render_request(&request);

202 +

203 + assert!(rendered.contains("<|open|>call tool=\"search\" index=\"2\"<|sep|>"));

204 + assert!(

205 + rendered

206 + .contains("<|open|>message role=\"tool\" tool=\"search\" index=\"2\"<|sep|>result-b")

207 + );

208 +}

209 +

210 #[test]

211 fn defaults_thinking_effort_to_max() {

212 let rendered = render_request(&crate::request::ChatRequest::for_test());

rust/src/parser/src/unified/kimi_k3.rs

+14 -37

@@ -148,9 +148,8 @@ enum KimiK3Mode {

148 pub struct KimiK3UnifiedParser {

149 buffer: String,

150 mode: KimiK3Mode,

151 - /// Parser-provided tool-call IDs (`{tool}:{zero_based_index}`) by tool

152 - /// index; its length is also the count of emitted calls.

153 - call_ids: Vec<String>,

151 + /// Number of calls emitted in the current response.

152 + emitted_call_count: usize,

153 tokenizer: DynTokenizer,

154 open_token_id: u32,

155 sep_token_id: u32,

@@ -165,7 +164,7 @@ impl KimiK3UnifiedParser {

164 Ok(Self {

165 buffer: String::new(),

166 mode: KimiK3Mode::default(),

168 - call_ids: Vec::new(),

167 + emitted_call_count: 0,

168 tokenizer,

169 open_token_id,

170 sep_token_id,

@@ -226,7 +225,7 @@ impl KimiK3UnifiedParser {

225 }

226 KimiK3Event::CallComplete { arguments } => {

227 let mode = std::mem::replace(&mut self.mode, KimiK3Mode::Tools);

229 - let KimiK3Mode::Call { name, index, .. } = mode else {

228 + let KimiK3Mode::Call { name, .. } = mode else {

229 return Err(parsing_failed!(

230 "Kimi K3 call completion without an active tool call"

231 ));

@@ -237,8 +236,8 @@ impl KimiK3UnifiedParser {

236 return Ok(());

237 }

238

240 - let tool_index = self.call_ids.len();

241 - self.call_ids.push(tool_call_id_for(&name, index.as_deref()));

239 + let tool_index = self.emitted_call_count;

240 + self.emitted_call_count += 1;

241 output.push_call(ToolCallDelta {

242 tool_index,

243 name: Some(name),

@@ -251,7 +250,7 @@ impl KimiK3UnifiedParser {

250

251 fn reset_state(&mut self) -> String {

252 self.mode = KimiK3Mode::Idle;

254 - self.call_ids.clear();

253 + self.emitted_call_count = 0;

254 std::mem::take(&mut self.buffer)

255 }

256 }

@@ -266,7 +265,7 @@ impl UnifiedParser for KimiK3UnifiedParser {

265

266 fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {

267 self.buffer.clear();

269 - self.call_ids.clear();

268 + self.emitted_call_count = 0;

269 self.initialize_mode(prompt_token_ids);

270 Ok(())

271 }

@@ -279,10 +278,6 @@ impl UnifiedParser for KimiK3UnifiedParser {

278 Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER)

279 }

280

282 - fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

283 - self.call_ids.get(tool_index).map(String::as_str)

284 - }

285 -

281 fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {

282 self.buffer.push_str(chunk);

283

@@ -313,7 +308,6 @@ impl UnifiedParser for KimiK3UnifiedParser {

308 }

309 }

310

316 - // Keep call_ids so tool_call_id() stays available after the stream ends.

311 self.mode = KimiK3Mode::Idle;

312 Ok(output)

313 }

@@ -323,20 +317,6 @@ impl UnifiedParser for KimiK3UnifiedParser {

317 }

318 }

319

326 -/// Build the API-side tool-call ID from the XTML one-based `index` attribute.

327 -///

328 -/// The ID uses the zero-based call ordinal; XTML's message index stays

329 -/// one-based when rendering tool result messages.

330 -fn tool_call_id_for(name: &str, index: Option<&str>) -> String {

331 - match index {

332 - None => name.to_string(),

333 - Some(raw) => match raw.parse::<i64>() {

334 - Ok(one_based) => format!("{name}:{}", one_based - 1),

335 - Err(_) => format!("{name}:{raw}"),

336 - },

337 - }

338 -}

339 -

320 /// Parse one Kimi K3 event from buffered streaming input.

321 fn parse_next_kimi_k3_event(

322 input: &mut KimiK3Input<'_>,

@@ -762,7 +742,6 @@ mod tests {

742 "hours": [8, 20],

743 })

744 );

765 - assert_eq!(parser.tool_call_id(0), Some("get_weather:0"));

745 }

746

747 #[test]

@@ -929,8 +908,6 @@ mod tests {

908 assert_eq!(calls[1].tool_index, 1);

909 assert_eq!(calls[1].name.as_deref(), Some("get_time"));

910 assert_eq!(calls[1].arguments, "{}");

932 - assert_eq!(parser.tool_call_id(0), Some("get_weather:0"));

933 - assert_eq!(parser.tool_call_id(1), Some("get_time:1"));

911 }

912

913 #[test]

@@ -1018,11 +995,10 @@ mod tests {

995 assert_eq!(calls.len(), 1);

996 assert_eq!(calls[0].tool_index, 0);

997 assert_eq!(calls[0].name.as_deref(), Some("real"));

1021 - assert_eq!(parser.tool_call_id(0), Some("real:1"));

998 }

999

1000 #[test]

1025 - fn kimi_k3_tool_call_id_follows_index_attribute() {

1001 + fn kimi_k3_tool_indices_ignore_xtml_index_attribute() {

1002 let tools_body = format!(

1003 "{}{}{}",

1004 call("tool=\"first\" index=\"3\"", ""),

@@ -1032,11 +1008,12 @@ mod tests {

1008 let text = thinking_output("t", "", &tools_body);

1009

1010 let mut parser = test_parser();

1035 - parser.parse_complete(&text).unwrap();

1011 + let output = parser.parse_complete(&text).unwrap();

1012

1037 - assert_eq!(parser.tool_call_id(0), Some("first:2"));

1038 - assert_eq!(parser.tool_call_id(1), Some("second"));

1039 - assert_eq!(parser.tool_call_id(2), Some("third:x"));

1013 + assert_eq!(

1014 + output.calls().iter().map(|call| call.tool_index).collect::<Vec<_>>(),

1015 + [0, 1, 2]

1016 + );

1017 }

1018

1019 #[test]

tests/tool_use/test_kimi_k3_tool_parser.py

+9 -3

@@ -141,7 +141,6 @@ def test_extract_tool_calls_with_response_and_typed_arguments():

141 assert extracted.content == "answer"

142 assert len(extracted.tool_calls) == 1

143 tool_call = extracted.tool_calls[0]

144 - assert tool_call.id == "calc:0"

144 assert tool_call.function.name == "calc"

145 assert json.loads(tool_call.function.arguments) == {

146 "x": 1,

@@ -168,7 +167,6 @@ def test_delegating_parser_preserves_tool_calls_after_reasoning():

167 assert content == "answer"

168 assert tool_calls is not None

169 assert len(tool_calls) == 1

171 - assert tool_calls[0].id == "calc:0"

170 assert tool_calls[0].name == "calc"

171 assert json.loads(tool_calls[0].arguments) == {"x": 1}

172

@@ -364,11 +362,19 @@ def test_streaming_split_markers_do_not_leak():

362 assert OPEN not in content

363 assert SEP not in content

364 assert len(tool_deltas) == 1

367 - assert tool_deltas[0].id == "calc:0"

365 assert tool_deltas[0].function.name == "calc"

366 assert json.loads(tool_deltas[0].function.arguments) == {"x": 1}

367

368

369 +def test_tool_call_ids_are_unique_across_messages():

370 + output = _tools(_call("calc", 1))

371 +

372 + first = KimiK3ToolParser(DummyTokenizer()).extract_tool_calls(output, _request())

373 + second = KimiK3ToolParser(DummyTokenizer()).extract_tool_calls(output, _request())

374 +

375 + assert first.tool_calls[0].id != second.tool_calls[0].id

376 +

377 +

378 def test_streaming_consumed_response_prefix_no_call_keeps_content():

379 parser = KimiK3ToolParser(DummyTokenizer())

380 request = _request()

vllm/tool_parsers/kimi_k3_tool_parser.py

-10

@@ -218,7 +218,6 @@ class KimiK3ToolParser(ToolParser):

218 """

219 call_attrs = self._attrs(attrs)

220 tool_name = call_attrs.get("tool", "")

221 - tool_index = call_attrs.get("index", "")

221 arguments: dict = {}

222 for arg_match in self._arg_re.finditer(body):

223 arg_attrs = self._attrs(arg_match["attrs"])

@@ -234,16 +233,7 @@ class KimiK3ToolParser(ToolParser):

233 arguments[key] = raw_value

234 if not tool_name:

235 return None

237 - tool_call_id = tool_name

238 - if tool_index:

239 - try:

240 - tool_call_id = f"{tool_name}:{int(tool_index) - 1}"

241 - except ValueError:

242 - tool_call_id = f"{tool_name}:{tool_index}"

243 - # id uses the API-side zero-based call ordinal; XTML's message index

244 - # stays one-based when rendering tool result messages.

236 return ToolCall(

246 - id=tool_call_id,

237 type="function",

238 function=FunctionCall(

239 name=tool_name,