[Rust Frontend] Handle continue_final_message with renderer sentinel (#47844)

Signed-off-by: Bugen Zhao <[email protected]>

.buildkite/test-amd.yaml

@@ -2026,7 +2026,7 @@ steps:
   - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
   - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
   - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
-  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
+  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"
 
 - label: Rust Frontend Core Correctness # TBD
   timeout_in_minutes: 180

.buildkite/test_areas/rust_frontend.yaml

@@ -64,7 +64,7 @@ steps:
   - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
   # - pytest -v -s entrypoints/serve/dev/test_sleep.py
   # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op).
-  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
+  - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"
 
 - label: Rust Frontend Core Correctness
   timeout_in_minutes: 30

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

@@ -155,11 +155,23 @@ impl HfChatRenderer {
         effective_template: &CompiledChatTemplate,
         request: &ChatRequest,
     ) -> Result<RenderedPrompt> {
-        let messages = to_template_messages(
+        let mut messages = to_template_messages(
             &request.messages,
             effective_template.content_format(),
             self.multimodal.as_ref(),
         )?;
+
+        // Handling of `continue_final_message`:
+        // Append a sentinel tag to the final message content, render as usual, then
+        // truncate the rendered prompt at the tag so any template suffix after the
+        // final message content (e.g. the end-of-turn marker) is dropped.
+        let final_message_text = if request.chat_options.continue_final_message() {
+            let final_message = messages.last_mut().ok_or(Error::EmptyMessages)?;
+            Some(append_continue_final_message_tag(final_message)?)
+        } else {
+            None
+        };
+
         let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools));
         trace!(
             message_count = messages.len(),
@@ -183,6 +195,13 @@ impl HfChatRenderer {
             })
             .map_err(|error| Error::ChatTemplate(error.to_report_string()))?;
 
+        let prompt = match &final_message_text {
+            Some(final_message_text) => {
+                truncate_prompt_at_continue_final_message_tag(prompt, final_message_text)?
+            }
+            None => prompt,
+        };
+
         trace!(
             prompt_len = prompt.len(),
             prompt, "rendered chat template prompt"
@@ -429,6 +448,74 @@ fn to_template_string_content(
     }
 }
 
+/// Sentinel appended to the final message content when `continue_final_message`
+/// is requested, used to locate the truncation point in the rendered prompt.
+///
+/// Same literal as `transformers`. Occurrences of this string earlier in the
+/// prompt are harmless because truncation uses the rightmost match, and the
+/// appended sentinel ends up last as long as the template renders messages in
+/// order.
+const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG ";
+
+/// Append [`CONTINUE_FINAL_MESSAGE_TAG`] to the trailing text of the final
+/// message, returning the original text for post-render validation.
+// TODO: transformers v5 also allows continuing a non-`content` field (e.g.
+// `reasoning_content`) by passing a field name; only the boolean form is
+// supported here.
+fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result<String> {
+    let text = match &mut message.content {
+        TemplateContent::String(text) => Some(text),
+        // Pick the last text part in the message.
+        TemplateContent::OpenAi(parts) => parts.iter_mut().rev().find_map(|part| match part {
+            TemplateContentPart::Text { text } => Some(text),
+            TemplateContentPart::Image => None,
+        }),
+    };
+    let text = text.ok_or_else(|| {
+        Error::ChatTemplate(
+            "continue_final_message is set but there is no text to continue \
+             in the final message"
+                .to_string(),
+        )
+    })?;
+
+    let original = text.clone();
+    text.push_str(CONTINUE_FINAL_MESSAGE_TAG);
+    Ok(original)
+}
+
+/// Truncate the rendered prompt at [`CONTINUE_FINAL_MESSAGE_TAG`] so that it
+/// ends exactly with the final message content, dropping any template suffix
+/// such as end-of-turn markers.
+fn truncate_prompt_at_continue_final_message_tag(
+    mut rendered: String,
+    final_message_text: &str,
+) -> Result<String> {
+    let tag_loc = rendered
+        .rfind(CONTINUE_FINAL_MESSAGE_TAG.trim_end())
+        .filter(|_| rendered.contains(final_message_text.trim()));
+    let Some(tag_loc) = tag_loc else {
+        return Err(Error::ChatTemplate(format!(
+            "continue_final_message is set but the final message does not appear \
+             in the prompt after applying the chat template! This can happen if \
+             the chat template deletes portions of the final message. Final \
+             message to continue: {}",
+            final_message_text.trim(),
+        )));
+    };
+
+    if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) {
+        // The template preserved spacing, so a plain cut at the tag suffices.
+        rendered.truncate(tag_loc);
+    } else {
+        // The template trimmed the trailing spacing of the message content, so
+        // apply the same trimming to the retained prefix.
+        rendered.truncate(tag_loc);
+        rendered.truncate(rendered.trim_end().len());
+    }
+    Ok(rendered)
+}
+
 fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
     tools
         .iter()
@@ -548,28 +635,124 @@ mod tests {
             ChatRole::Assistant,
             "The capital of",
         )]);
+        let template =
+            "{% if continue_final_message %}continue:{% endif %}{{ messages[0].content }}";
 
-        assert_eq!(
-            render(
-                Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
-                &request,
-            )
-            .unwrap(),
-            "new"
-        );
+        assert_eq!(render(Some(template), &request).unwrap(), "The capital of");
 
         request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
 
         assert_eq!(
-            render(
-                Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
-                &request,
-            )
-            .unwrap(),
-            "continue"
+            render(Some(template), &request).unwrap(),
+            "continue:The capital of"
         );
     }
 
+    #[test]
+    fn continue_final_message_truncates_template_suffix() {
+        let mut request = sample_request(vec![
+            ChatMessage::text(ChatRole::User, "What is the capital of France?"),
+            ChatMessage::text(ChatRole::Assistant, "The capital of"),
+        ]);
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
+
+        // The Qwen3 template is unaware of `continue_final_message`; the
+        // end-of-turn marker it appends must still be stripped.
+        let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap();
+
+        expect![[r#"
+            <|im_start|>user
+            What is the capital of France?<|im_end|>
+            <|im_start|>assistant
+            <think>
+
+            </think>
+
+            The capital of"#]]
+        .assert_eq(&rendered);
+    }
+
+    #[test]
+    fn continue_final_message_trims_like_the_template_does() {
+        let mut request = sample_request(vec![ChatMessage::text(ChatRole::Assistant, "Sure, ")]);
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
+
+        // The template trims the trailing spacing of the message content, so
+        // the truncated prompt must be trimmed the same way.
+        let rendered = render(
+            Some("{{ messages[0].content.strip() }}<|im_end|>"),
+            &request,
+        )
+        .unwrap();
+
+        assert_eq!(rendered, "Sure,");
+    }
+
+    #[test]
+    fn continue_final_message_appends_to_last_text_part() {
+        // The renderer itself is role-agnostic like transformers (the
+        // assistant-final restriction is enforced by request validation
+        // upstream), so a multimodal user message exercises the part
+        // selection: the sentinel must attach to the last *text* part,
+        // skipping the trailing image.
+        let mut request = sample_request(vec![ChatMessage::user(vec![
+            ChatContentPart::text("Sure,"),
+            ChatContentPart::image_url("data:image/png;base64,test"),
+        ])]);
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
+
+        let rendered = render_mm(
+            "{% for item in messages[0].content %}{% if item.type == 'image' %}<image>{% else %}{{ item.text }}{% endif %}{% endfor %}<|im_end|>",
+            &request,
+            ChatTemplateContentFormatOption::OpenAi,
+        )
+        .unwrap()
+        .prompt;
+
+        // Anything rendered after the continued text (here the image
+        // placeholder and the end marker) is truncated away, matching
+        // transformers.
+        assert_eq!(rendered, Prompt::Text("Sure,".to_string()));
+    }
+
+    #[test]
+    fn continue_final_message_composes_with_aware_templates() {
+        // A template that reads `continue_final_message` and skips its own
+        // end-of-turn marker must produce the same prompt as an unaware one:
+        // the sentinel truncation degenerates to a cut at the very end.
+        let mut request = sample_request(vec![
+            ChatMessage::text(ChatRole::User, "hi"),
+            ChatMessage::text(ChatRole::Assistant, "Sure,"),
+        ]);
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
+
+        let aware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}{% if not (loop.last and continue_final_message) %}<|im_end|>\n{% endif %}{% endfor %}";
+        let unaware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}";
+
+        let expected = "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\nSure,";
+        assert_eq!(render(Some(aware), &request).unwrap(), expected);
+        assert_eq!(render(Some(unaware), &request).unwrap(), expected);
+    }
+
+    #[test]
+    fn continue_final_message_errors_when_template_drops_final_message() {
+        let mut request = sample_request(vec![
+            ChatMessage::text(ChatRole::User, "hi"),
+            ChatMessage::text(ChatRole::Assistant, "Sure,"),
+        ]);
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
+
+        let error = render(
+            Some(
+                "{% for m in messages %}{% if m.role == 'user' %}{{ m.content }}{% endif %}{% endfor %}",
+            ),
+            &request,
+        )
+        .unwrap_err();
+
+        assert!(matches!(error, Error::ChatTemplate(_)));
+    }
+
     #[test]
     fn chat_template_flattens_text_parts_for_string_templates() {
         let request = sample_request(vec![ChatMessage::user(vec![