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

[Rust Frontend] Accept custom chat roles for HF templates (#58311)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 24, 2026 at 13:43 UTC a1b6763df65981474ba5c4dc27e16566eddb41e3

14 files changed +215 -20

rust/src/chat/src/error.rs

+4 -1

@@ -21,6 +21,8 @@ pub enum Error {

21 InvalidReasoningEffort(String),

22 #[error("{message}")]

23 InvalidReasoningControl { message: String },

24 + #[error("chat role `{role}` is not supported by this chat renderer")]

25 + UnsupportedChatRole { role: String },

26 #[error("multimodal input is not supported by this chat renderer")]

27 UnsupportedMultimodalRenderer,

28 #[error("unsupported multimodal content: {0}")]

@@ -114,7 +116,8 @@ impl Error {

116 | Self::InvalidReasoningControl { .. }

117 | Self::DuplicateToolName { .. }

118 | Self::ToolChoiceRequiresTools

117 - | Self::ToolChoiceFunctionNotFound { .. } => true,

119 + | Self::ToolChoiceFunctionNotFound { .. }

120 + | Self::UnsupportedChatRole { .. } => true,

121 Self::Text(error) => error.is_request_validation_error(),

122 Self::UnsupportedMultimodalRenderer

123 | Self::UnsupportedMultimodalContent(_)

rust/src/chat/src/multimodal.rs

+4 -2

@@ -599,7 +599,8 @@ fn extract_media_parts(

599 ChatMessage::System { content }

600 | ChatMessage::Developer { content, .. }

601 | ChatMessage::User { content }

602 - | ChatMessage::ToolResponse { content, .. } => content,

602 + | ChatMessage::ToolResponse { content, .. }

603 + | ChatMessage::Custom { content, .. } => content,

604 ChatMessage::Assistant { .. } => {

605 bail_multimodal!("renderer reported multimodal assistant content")

606 }

@@ -622,7 +623,8 @@ fn extract_media_parts(

623 ChatMessage::System { content }

624 | ChatMessage::Developer { content, .. }

625 | ChatMessage::User { content }

625 - | ChatMessage::ToolResponse { content, .. } => Some(content),

626 + | ChatMessage::ToolResponse { content, .. }

627 + | ChatMessage::Custom { content, .. } => Some(content),

628 ChatMessage::Assistant { .. } => None,

629 };

630 match content {

rust/src/chat/src/renderer/deepseek.rs

@@ -227,6 +227,9 @@ pub(super) fn render_request_with_media_order(

227 );

228 }

229 }

230 + ChatMessage::Custom { role, .. } => {

231 + return Err(Error::UnsupportedChatRole { role: role.clone() });

232 + }

233 }

234

235 if is_user_like_entry(message, current_render_index as usize, dialect)

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

@@ -158,6 +158,7 @@ fn render_message(

158 thinking_mode == ThinkingMode::Thinking && after_or_at_last_user_turn,

159 content,

160 ),

161 + ChatMessage::Custom { role, .. } => Err(Error::UnsupportedChatRole { role: role.clone() }),

162 }

163 }

164

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

+18

@@ -5,6 +5,7 @@ use std::path::PathBuf;

6 use expect_test::{ExpectFile, expect, expect_file};

7 use serde_json::Value;

8 +use thiserror_ext::AsReport as _;

10 use super::DeepSeekV4ChatRenderer;

11 use crate::ChatRenderer;

@@ -764,3 +765,20 @@ fn last_developer_omits_generation_prompt_when_disabled() {

765

766 expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered);

767 }

768 +

769 +#[test]

770 +fn custom_roles_are_rejected_as_request_errors() {

771 + let request = ChatRequest {

772 + messages: vec![

773 + ChatMessage::custom("root", "identity"),

774 + ChatMessage::user("hello"),

775 + ],

776 + ..ChatRequest::for_test()

777 + };

778 +

779 + let error = DeepSeekV4ChatRenderer::default().render(&request).unwrap_err();

780 +

781 + assert!(error.is_request_validation_error());

782 + expect!["chat role `root` is not supported by this chat renderer"]

783 + .assert_eq(&error.to_report_string());

784 +}

rust/src/chat/src/renderer/harmony/mod.rs

+5 -1

@@ -188,7 +188,8 @@ fn peel_leading_instructions(

188 )),

189 ChatMessage::User { .. }

190 | ChatMessage::Assistant { .. }

191 - | ChatMessage::ToolResponse { .. } => Ok((None, None, messages)),

191 + | ChatMessage::ToolResponse { .. }

192 + | ChatMessage::Custom { .. } => Ok((None, None, messages)),

193 }

194 }

195

@@ -328,6 +329,9 @@ fn to_harmony_message(

329 .with_recipient("assistant"),

330 ]

331 }

332 + ChatMessage::Custom { role, .. } => {

333 + return Err(Error::UnsupportedChatRole { role: role.clone() });

334 + }

335 })

336 }

337

rust/src/chat/src/renderer/hf/mod.rs

+35 -9

@@ -258,8 +258,8 @@ impl ChatRenderer for HfChatRenderer {

258 // TODO: borrow more fields directly from the original `ChatMessage`.

259 #[serde_with::skip_serializing_none]

260 #[derive(Debug, Serialize)]

261 -struct TemplateMessage {

262 - role: &'static str,

261 +struct TemplateMessage<'a> {

262 + role: &'a str,

263 content: TemplateContent,

264 // Developer-role messages may provide message-local tools in the same shape

265 // as top-level request tools.

@@ -322,22 +322,22 @@ struct TemplateToolDefinition {

322 }

323

324 /// Convert chat messages into the JSON shape expected by Jinja chat templates.

325 -fn to_template_messages(

326 - messages: &[ChatMessage],

325 +fn to_template_messages<'a>(

326 + messages: &'a [ChatMessage],

327 content_format: ChatTemplateContentFormat,

328 multimodal: Option<&MultimodalRenderInfo>,

329 -) -> Result<Vec<TemplateMessage>> {

329 +) -> Result<Vec<TemplateMessage<'a>>> {

330 messages

331 .iter()

332 .map(|message| to_template_message(message, content_format, multimodal))

333 .collect()

334 }

335

336 -fn to_template_message(

337 - message: &ChatMessage,

336 +fn to_template_message<'a>(

337 + message: &'a ChatMessage,

338 content_format: ChatTemplateContentFormat,

339 multimodal: Option<&MultimodalRenderInfo>,

340 -) -> Result<TemplateMessage> {

340 +) -> Result<TemplateMessage<'a>> {

341 Ok(match message {

342 ChatMessage::System { content } => TemplateMessage {

343 role: "system",

@@ -394,6 +394,15 @@ fn to_template_message(

394 tool_calls: None,

395 tool_call_id: Some(tool_call_id.clone()),

396 },

397 + ChatMessage::Custom { role, content } => TemplateMessage {

398 + role,

399 + content: to_template_content(content, content_format, multimodal)?,

400 + tools: None,

401 + reasoning: None,

402 + reasoning_content: None,

403 + tool_calls: None,

404 + tool_call_id: None,

405 + },

406 })

407 }

408

@@ -524,7 +533,7 @@ const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG ";

533 // TODO: transformers v5 also allows continuing a non-`content` field (e.g.

534 // `reasoning_content`) by passing a field name; only the boolean form is

535 // supported here.

527 -fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result<String> {

536 +fn append_continue_final_message_tag(message: &mut TemplateMessage<'_>) -> Result<String> {

537 let text = match &mut message.content {

538 TemplateContent::String(text) => Some(text),

539 // Pick the last text part in the message.

@@ -1001,6 +1010,23 @@ mod tests {

1010 assert_eq!(rendered, "developer|policy|get_weather|city");

1011 }

1012

1013 + #[test]

1014 + fn chat_template_passes_custom_roles_unchanged() {

1015 + let request = sample_request(vec![

1016 + ChatMessage::custom("root", "identity"),

1017 + ChatMessage::system("policy"),

1018 + ChatMessage::user("hello"),

1019 + ]);

1020 +

1021 + let rendered = render(

1022 + Some("{% for m in messages %}{{ m.role }}={{ m.content }};{% endfor %}"),

1023 + &request,

1024 + )

1025 + .unwrap();

1026 +

1027 + assert_eq!(rendered, "root=identity;system=policy;user=hello;");

1028 + }

1029 +

1030 #[test]

1031 fn chat_template_keeps_string_text_for_openai_detected_templates() {

1032 let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);

rust/src/chat/src/renderer/hf/template.rs

+1 -1

@@ -45,7 +45,7 @@ fn build_environment(template: String) -> Result<Environment<'static>> {

45 #[serde_with::skip_serializing_none]

46 #[derive(Default, Serialize)]

47 pub(super) struct TemplateContext<'a> {

48 - pub(super) messages: &'a [TemplateMessage],

48 + pub(super) messages: &'a [TemplateMessage<'a>],

49 pub(super) add_generation_prompt: bool,

50 pub(super) continue_final_message: bool,

51 pub(super) tools: Option<&'a [TemplateTool]>,

rust/src/chat/src/renderer/inkling/mod.rs

@@ -345,6 +345,9 @@ impl ChatRenderer for InklingChatRenderer {

345 &tool_call_id_to_name,

346 )?;

347 }

348 + ChatMessage::Custom { role, .. } => {

349 + return Err(Error::UnsupportedChatRole { role: role.clone() });

350 + }

351 }

352 }

353

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

@@ -205,6 +205,9 @@ pub(super) fn render_request_with_media_order(

205 write_assistant_message(&mut out, content, thinking)?;

206 }

207 ChatMessage::ToolResponse { .. } => unreachable!("handled above"),

208 + ChatMessage::Custom { role, .. } => {

209 + return Err(Error::UnsupportedChatRole { role: role.clone() });

210 + }

211 }

212 }

213 flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?;

rust/src/chat/src/request.rs

+25 -4

@@ -24,6 +24,9 @@ pub enum ChatRole {

24 User,

25 Assistant,

26 ToolResponse,

27 + /// Role outside the standard set, passed to chat templates unchanged.

28 + #[serde(untagged)]

29 + Custom(String),

30 }

32 /// One text-only chat content part in OpenAI-style block format.

@@ -212,6 +215,10 @@ pub enum ChatMessage {

215 content: ChatContent,

216 tool_call_id: String,

217 },

218 + /// Message content under a role outside the standard set, passed to chat

219 + /// templates unchanged.

220 + #[serde(untagged)]

221 + Custom { role: String, content: ChatContent },

222 }

223

224 impl ChatMessage {

@@ -230,6 +237,7 @@ impl ChatMessage {

237 use ChatMessage::tool_response() instead"

238 )

239 }

240 + ChatRole::Custom(role) => Self::custom(role, content),

241 }

242 }

243

@@ -286,6 +294,14 @@ impl ChatMessage {

294 }

295 }

296

297 + /// Construct one chat message with a role outside the standard set.

298 + pub fn custom(role: impl Into<String>, content: impl Into<ChatContent>) -> Self {

299 + Self::Custom {

300 + role: role.into(),

301 + content: content.into(),

302 + }

303 + }

304 +

305 /// Return the chat role of this message.

306 pub fn role(&self) -> ChatRole {

307 match self {

@@ -294,6 +310,7 @@ impl ChatMessage {

310 Self::User { .. } => ChatRole::User,

311 Self::Assistant { .. } => ChatRole::Assistant,

312 Self::ToolResponse { .. } => ChatRole::ToolResponse,

313 + Self::Custom { role, .. } => ChatRole::Custom(role.clone()),

314 }

315 }

316

@@ -303,7 +320,8 @@ impl ChatMessage {

320 Self::System { content }

321 | Self::Developer { content, .. }

322 | Self::User { content }

306 - | Self::ToolResponse { content, .. } => content.try_flatten_to_text(),

323 + | Self::ToolResponse { content, .. }

324 + | Self::Custom { content, .. } => content.try_flatten_to_text(),

325 Self::Assistant { content } => Ok(content.text()),

326 }

327 }

@@ -315,7 +333,8 @@ impl ChatMessage {

333 Self::System { .. }

334 | Self::Developer { .. }

335 | Self::User { .. }

318 - | Self::ToolResponse { .. } => None,

336 + | Self::ToolResponse { .. }

337 + | Self::Custom { .. } => None,

338 }

339 }

340

@@ -325,7 +344,8 @@ impl ChatMessage {

344 Self::System { content }

345 | Self::Developer { content, .. }

346 | Self::User { content }

328 - | Self::ToolResponse { content, .. } => content.has_multimodal(),

347 + | Self::ToolResponse { content, .. }

348 + | Self::Custom { content, .. } => content.has_multimodal(),

349 Self::Assistant { .. } => false,

350 }

351 }

@@ -673,13 +693,14 @@ impl ChatRequest {

693 impl ChatRole {

694 /// Return the chat-template role string used by the current text-only chat

695 /// backend.

676 - pub fn as_str(&self) -> &'static str {

696 + pub fn as_str(&self) -> &str {

697 match self {

698 Self::System => "system",

699 Self::Developer => "developer",

700 Self::User => "user",

701 Self::Assistant => "assistant",

702 Self::ToolResponse => "tool_response",

703 + Self::Custom(role) => role,

704 }

705 }

706 }

rust/src/server/src/routes/openai/chat_completions/convert.rs

+41

@@ -317,6 +317,9 @@ pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, A

317 convert_content(content)?,

318 convert_message_tools(tools)?,

319 )),

320 + ChatMessage::Custom { role, content } => {

321 + Ok(VllmChatMessage::custom(role, convert_content(content)?))

322 + }

323 }

324 }

325

@@ -442,6 +445,7 @@ mod tests {

445 use expect_test::expect;

446 use llm_multimodal::ImageDetail;

447 use serde_json::json;

448 + use thiserror_ext::AsReport as _;

449 use validator::Validate;

450 use vllm_chat::{

451 AssistantContentBlock, AssistantToolCall, ChatContentPart, ChatMessage as VllmChatMessage,

@@ -1047,6 +1051,43 @@ mod tests {

1051 );

1052 }

1053

1054 + #[test]

1055 + fn prepare_chat_request_accepts_custom_role_messages() {

1056 + let request: ChatCompletionRequest = serde_json::from_value(json!({

1057 + "messages": [

1058 + {"role": "root", "content": "Custom identity."},

1059 + {"role": "user", "content": "hello"},

1060 + ],

1061 + }))

1062 + .expect("parse custom role message");

1063 +

1064 + let prepared = prepare_chat_request(

1065 + request,

1066 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

1067 + ResolvedRequestContext::default(),

1068 + )

1069 + .expect("request is valid");

1070 +

1071 + assert_eq!(

1072 + prepared.chat_request.messages,

1073 + vec![

1074 + VllmChatMessage::custom("root", "Custom identity."),

1075 + VllmChatMessage::user("hello"),

1076 + ]

1077 + );

1078 + }

1079 +

1080 + #[test]

1081 + fn chat_http_message_keeps_standard_role_errors() {

1082 + let error = serde_json::from_value::<ChatMessage>(json!({

1083 + "role": "tool",

1084 + "content": "Sunny",

1085 + }))

1086 + .unwrap_err();

1087 +

1088 + expect!["missing field `tool_call_id`"].assert_eq(&error.to_report_string());

1089 + }

1090 +

1091 #[test]

1092 fn prepare_chat_request_maps_image_url_content_parts() {

1093 let request = ChatCompletionRequest {

rust/src/server/src/routes/openai/utils/types.rs

+70 -1

@@ -410,9 +410,17 @@ impl ToolReference {

410 // Chat Messages

411 // ============================================================================

412

413 +/// One chat message, tagged by `role`.

414 +///

415 +/// The derived serde impls use `remote = "Self"` and cover only the standard

416 +/// roles; the trait impls below dispatch on the role tag and handle

417 +/// [`ChatMessage::Custom`] separately. A `#[serde(untagged)]` fallback variant

418 +/// would also accept malformed standard messages, such as a `tool` message

419 +/// without `tool_call_id`, and replace their field errors with serde's generic

420 +/// untagged-enum error.

421 #[serde_with::skip_serializing_none]

422 #[derive(Debug, Clone, Deserialize, Serialize)]

415 -#[serde(tag = "role")]

423 +#[serde(tag = "role", remote = "Self")]

424 pub enum ChatMessage {

425 #[serde(rename = "system")]

426 System {

@@ -447,6 +455,67 @@ pub enum ChatMessage {

455 tools: Option<Vec<Tool>>,

456 name: Option<String>,

457 },

458 + /// Message with a role outside the OpenAI set, such as `root`. Chat

459 + /// templates receive the role string unchanged.

460 + #[serde(skip)]

461 + Custom {

462 + role: String,

463 + content: MessageContent,

464 + },

465 +}

466 +

467 +/// Role tags of the standard [`ChatMessage`] variants.

468 +const STANDARD_CHAT_ROLES: &[&str] = &[

469 + "system",

470 + "user",

471 + "assistant",

472 + "tool",

473 + "function",

474 + "developer",

475 +];

476 +

477 +/// Wire shape of [`ChatMessage::Custom`].

478 +#[derive(Deserialize, Serialize)]

479 +struct CustomChatMessage<R, C> {

480 + role: R,

481 + content: C,

482 +}

483 +

484 +impl Serialize for ChatMessage {

485 + fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>

486 + where

487 + S: serde::Serializer,

488 + {

489 + match self {

490 + Self::Custom { role, content } => {

491 + CustomChatMessage { role, content }.serialize(serializer)

492 + }

493 + _ => Self::serialize(self, serializer),

494 + }

495 + }

496 +}

497 +

498 +impl<'de> Deserialize<'de> for ChatMessage {

499 + /// Dispatch on the role tag: standard roles keep their derived errors,

500 + /// and any other string role becomes [`ChatMessage::Custom`].

501 + fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>

502 + where

503 + D: Deserializer<'de>,

504 + {

505 + let value = Value::deserialize(deserializer)?;

506 + let is_custom = value

507 + .get("role")

508 + .and_then(Value::as_str)

509 + .is_some_and(|role| !STANDARD_CHAT_ROLES.contains(&role));

510 +

511 + if is_custom {

512 + let CustomChatMessage { role, content } =

513 + CustomChatMessage::deserialize(value).map_err(serde::de::Error::custom)?;

514 + Ok(Self::Custom { role, content })

515 + } else {

516 + Self::deserialize(value).map_err(serde::de::Error::custom)

517 + }

518 + }

519 }

520

521 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]

rust/src/server/src/routes/tests.rs

+2 -1

@@ -555,7 +555,8 @@ fn render_fake_message_content(

555 ChatMessage::System { content }

556 | ChatMessage::Developer { content, .. }

557 | ChatMessage::User { content }

558 - | ChatMessage::ToolResponse { content, .. } => render_fake_content(content, placeholder),

558 + | ChatMessage::ToolResponse { content, .. }

559 + | ChatMessage::Custom { content, .. } => render_fake_content(content, placeholder),

560 ChatMessage::Assistant { .. } => message.text_content(),

561 }

562 }