fix(generate): start a chat reply on its own line after the prompt echo

src/commands/generate.rs

@@ -1270,13 +1270,31 @@ fn build_cli_sampling_config_with_flags(
     build_sampling_config(resolved_cli_sampling_params(args, stop_token_ids, flags))
 }
 
-pub(super) fn print_generation_preamble(user_prompt: &str) -> Result<()> {
+/// Echo the prompt before the reply streams in.
+///
+/// With a chat template the reply is a new assistant turn, so the echoed prompt
+/// ends with a newline and the reply starts on its own line. Without one (raw
+/// completion, `--no-chat-template`, or a checkpoint without a template) the
+/// reply continues the prompt text, so no newline is added and the prompt and
+/// completion read as one passage.
+pub(super) fn print_generation_preamble(user_prompt: &str, reply_is_new_turn: bool) -> Result<()> {
     println!("Generating...");
-    print!("{}", user_prompt);
+    print!(
+        "{}",
+        generation_preamble_echo(user_prompt, reply_is_new_turn)
+    );
     io::stdout().flush()?;
     Ok(())
 }
 
+fn generation_preamble_echo(user_prompt: &str, reply_is_new_turn: bool) -> String {
+    if reply_is_new_turn && !user_prompt.ends_with('\n') {
+        format!("{user_prompt}\n")
+    } else {
+        user_prompt.to_string()
+    }
+}
+
 fn generated_suffix<'a>(full_text: &'a str, prompt_text: &str) -> &'a str {
     full_text.strip_prefix(prompt_text).unwrap_or(full_text)
 }
@@ -3207,7 +3225,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> {
     if select_backend().name() == "xla" {
         validate_xla_output_audio(args.generation.output_audio.as_deref())?;
         let num_layers = xla_num_layers(&args.model.model);
-        print_generation_preamble(&user_prompt)?;
+        print_generation_preamble(&user_prompt, prompt != user_prompt)?;
         let (generated_tokens, stats) = generate_xla(
             &args.model.model,
             num_layers,
@@ -3254,7 +3272,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> {
             &[],
         );
         let num_layers = resolve_cli_pipeline_num_layers(&args.model.model)?;
-        print_generation_preamble(&user_prompt)?;
+        print_generation_preamble(&user_prompt, prompt != user_prompt)?;
         generate_pipeline_text(
             &args.model.model,
             num_layers,
@@ -3363,7 +3381,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> {
                 no_chat_template: args.generation.no_chat_template,
             },
         )?;
-        print_generation_preamble(&user_prompt)?;
+        print_generation_preamble(&user_prompt, prompt != user_prompt)?;
         let generation = run_generation_mode(
             &model,
             &args,

src/commands/generate_diffusion.rs

@@ -176,7 +176,7 @@ pub(crate) fn run_diffusion_generation(
             None => (prompt_tokens, None),
         };
 
-    print_generation_preamble(user_prompt)?;
+    print_generation_preamble(user_prompt, !args.generation.no_chat_template)?;
 
     // Stream through the shared incremental detokenizer (byte-fallback
     // safe); raw ids are collected in parallel so any held-back tail bytes

src/commands/generate_florence2.rs

@@ -79,7 +79,7 @@ pub(crate) fn run_florence2_generation(
         .pop()
         .ok_or_else(|| anyhow!("Image decoding returned no image for {image_path:?}"))?;
 
-    print_generation_preamble(user_prompt)?;
+    print_generation_preamble(user_prompt, !args.generation.no_chat_template)?;
     println!();
 
     let started = Instant::now();

src/commands/generate_llada2.rs

@@ -100,7 +100,7 @@ pub(crate) fn run_llada2_generation(
         mlxcel_core::random_seed(seed);
     }
 
-    print_generation_preamble(user_prompt)?;
+    print_generation_preamble(user_prompt, !args.generation.no_chat_template)?;
 
     // Stream through the shared incremental detokenizer; raw ids are collected
     // in parallel so any held-back tail bytes can be flushed afterward.

src/commands/generate_nemotron_parse.rs

@@ -85,7 +85,7 @@ pub(crate) fn run_nemotron_parse_generation(
         .pop()
         .ok_or_else(|| anyhow!("Image decoding returned no image for {image_path:?}"))?;
 
-    print_generation_preamble(prompt)?;
+    print_generation_preamble(prompt, true)?;
     println!();
 
     let started = Instant::now();

src/commands/generate_tests.rs

@@ -14,6 +14,7 @@
 
 #[cfg(feature = "xla-backend")]
 use super::decode_xla_cli_images;
+use super::generation_preamble_echo;
 use super::{
     CliPromptMedia, CliSamplingFlagState, CliVideoFrameGroup, CliVideoFrames, OfflineDraftRoute,
     apply_user_chat_template, apply_vlm_chat_template, build_cli_sampling_config_with_flags,
@@ -1707,3 +1708,26 @@ fn the_repl_uses_a_drafter_only_with_an_explicit_server_kind() {
     args.speculative.draft_kind = Some("mtp".to_string());
     assert_eq!(repl_drafter(&args), None, "a kind without a drafter");
 }
+
+#[test]
+fn preamble_echo_ends_the_prompt_line_before_a_chat_reply() {
+    // A chat-templated reply is a new assistant turn: it must not run on from
+    // the echoed prompt ("...on a laptop.Cores ablaze").
+    assert_eq!(
+        generation_preamble_echo("Write a haiku.", true),
+        "Write a haiku.\n"
+    );
+    assert_eq!(
+        generation_preamble_echo("Already ends\n", true),
+        "Already ends\n"
+    );
+}
+
+#[test]
+fn preamble_echo_lets_a_raw_completion_continue_the_prompt() {
+    // Without a chat template the completion continues the prompt text.
+    assert_eq!(
+        generation_preamble_echo("The capital of France is", false),
+        "The capital of France is"
+    );
+}