[Rust Frontend] Anchor tokens after pending UTF-8 bytes at their own first byte (#58357)

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

rust/src/tokenizer/src/incremental.rs

@@ -42,6 +42,10 @@ pub(crate) struct DecodeStream<'a, T: Tokenizer + ?Sized> {
     prefix: String,
     prefix_index: usize,
     prefix_seeded: bool,
+    /// Joint decode from the last `push_token` that left tokens pending, i.e.
+    /// `decode(context ++ pending)`. Used to split the pending group from the
+    /// next token when that token provably does not merge with it.
+    pending_string: Option<String>,
     decoded: AttributedTextBuffer,
 }
 
@@ -60,6 +64,7 @@ impl<'a, T: Tokenizer + ?Sized> DecodeStream<'a, T> {
             prefix: String::new(),
             prefix_index: 0,
             prefix_seeded: prompt_token_ids.is_empty(),
+            pending_string: None,
             decoded: AttributedTextBuffer::default(),
         }
     }
@@ -72,6 +77,44 @@ const SAFE_SUFFIX_MIN: usize = 4;
 const SAFE_SUFFIX_MAX: usize = 6;
 
 impl<T: Tokenizer + ?Sized> DecodeStream<'_, T> {
+    /// Decide whether `token_id` can be anchored at its own first byte inside
+    /// `new_chunk`, separately from the older tokens still pending.
+    ///
+    /// Pending tokens (an incomplete UTF-8 sequence, or text still empty in
+    /// context) normally share one anchor with the token that completes them,
+    /// because the tokenizer only defines their joint decode. That sharing is
+    /// wrong when the new token does *not* merge with them, e.g. a stray
+    /// incomplete byte followed by a special token or plain ASCII: the joint
+    /// decode is then exactly `decode(pending) ++ decode([token_id])`, and the
+    /// new token's first byte is known. Returns that split offset, or `None`
+    /// when the tokens merge (byte fallback completing a character), when
+    /// nothing is pending, or when the isolated decode is not provably
+    /// consistent with the joint one. `pending` is the joint decode kept from the
+    /// push that left tokens pending. Decoded text is never changed by this;
+    /// only anchor placement is.
+    fn pending_group_split(
+        &self,
+        pending: Option<&str>,
+        token_id: u32,
+        new_chunk: &str,
+    ) -> Result<Option<usize>> {
+        let Some(pending) = pending else {
+            return Ok(None);
+        };
+        let prefix_len = self.prefix.len();
+        let pending_chunk = &pending[pending.floor_char_boundary(prefix_len)..];
+        // Empty pending text: today's shared anchor is already this token's first
+        // byte. Not a prefix: the new token completed the pending character.
+        if pending_chunk.is_empty() || !new_chunk.starts_with(pending_chunk) {
+            return Ok(None);
+        }
+        let own = self.tokenizer.decode(&[token_id], self.skip_special_tokens)?;
+        let independent = !own.is_empty()
+            && new_chunk.len() == pending_chunk.len() + own.len()
+            && new_chunk.ends_with(&own);
+        Ok(independent.then_some(pending_chunk.len()))
+    }
+
     /// Decode prompt-only context for prefix seeding.
     ///
     /// Prompt ids may come from the model vocabulary rather than the local
@@ -145,16 +188,40 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
             || (!produced_text && self.tokenizer.id_to_token(token_id).is_none());
         if zero_width {
             self.decoded.record_zero_width_token(token_id);
-        } else {
-            self.decoded.record_pending_token(token_id);
         }
-
         if !produced_text {
+            if !zero_width {
+                self.decoded.record_pending_token(token_id);
+                self.pending_string = Some(string);
+            }
             return Ok(0);
         }
+
         // Ensure we split at a utf-8 char boundary.
         let new_chunk = &string[string.floor_char_boundary(prefix_len)..];
-        self.decoded.append_visible_text(new_chunk);
+        // Text is released, so the pending group resolves in this push either
+        // way; never carry its joint decode into a later push.
+        let pending = self.pending_string.take();
+        let split = if zero_width {
+            None
+        } else {
+            self.pending_group_split(pending.as_deref(), token_id, new_chunk)?
+        };
+        match split {
+            Some(split) => {
+                // Older pending tokens own the prefix (including any U+FFFD from
+                // an incomplete character); this token owns exactly its own text.
+                self.decoded.append_visible_text(&new_chunk[..split]);
+                self.decoded.record_pending_token(token_id);
+                self.decoded.append_visible_text(&new_chunk[split..]);
+            }
+            None => {
+                if !zero_width {
+                    self.decoded.record_pending_token(token_id);
+                }
+                self.decoded.append_visible_text(new_chunk);
+            }
+        }
         self.ids.drain(..self.prefix_index);
         self.prefix = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
         self.prefix_index = self.ids.len();
@@ -184,6 +251,7 @@ impl<T: Tokenizer + ?Sized> IncrementalDecoder for DecodeStream<'_, T> {
             }
         }
         self.decoded.resolve_pending_zero_width();
+        self.pending_string = None;
         self.ids.clear();
         self.prefix.clear();
         self.prefix_index = 0;

rust/src/tokenizer/src/incremental/attribution/tests.rs

@@ -520,6 +520,106 @@ fn attribution_incomplete_utf8_tokens_share_replacement_anchor_on_flush() {
     });
 }
 
+#[test]
+fn attribution_special_token_after_incomplete_utf8_anchors_at_its_own_spelling() {
+    // A stray incomplete byte followed by a special token: the joint decode is
+    // exactly `decode([0xe4]) ++ decode([special])`, so the special token owns
+    // its own first byte instead of sharing the group's anchor at 0.
+    let tokenizer = tokenizer_with_special_token();
+    run_attribution_case(AttributionCase {
+        tokenizer: &tokenizer,
+        prompt_token_ids: vec![],
+        token_ids: vec![0xe4, SPECIAL_TOKEN_ID],
+        skip_special_tokens: false,
+        min_bytes_to_buffer: 0,
+        truncate_output_to: None,
+        expected_chunks: vec![decoded(
+            "�<special>",
+            &[0xe4, SPECIAL_TOKEN_ID],
+            &[visible(0), visible(3)],
+        )],
+        expected_full: decoded(
+            "�<special>",
+            &[0xe4, SPECIAL_TOKEN_ID],
+            &[visible(0), visible(3)],
+        ),
+    });
+}
+
+#[test]
+fn attribution_ascii_token_after_incomplete_utf8_anchors_at_its_own_byte() {
+    // Same rule for an ordinary token: the split is decided by the non-merge
+    // equality, not by the token being special.
+    let tokenizer = TestTokenizer::new();
+    run_attribution_case(AttributionCase {
+        tokenizer: &tokenizer,
+        prompt_token_ids: vec![],
+        token_ids: vec![0xe4, b'a' as u32],
+        skip_special_tokens: false,
+        min_bytes_to_buffer: 0,
+        truncate_output_to: None,
+        expected_chunks: vec![decoded(
+            "�a",
+            &[0xe4, b'a' as u32],
+            &[visible(0), visible(3)],
+        )],
+        expected_full: decoded("�a", &[0xe4, b'a' as u32], &[visible(0), visible(3)]),
+    });
+}
+
+#[test]
+fn attribution_byte_fallback_completion_keeps_shared_anchor_before_special_token() {
+    // The completing byte merges with the pending group (no split); the special
+    // token that follows a complete character is an ordinary independent token.
+    let tokenizer = tokenizer_with_special_token();
+    run_attribution_case(AttributionCase {
+        tokenizer: &tokenizer,
+        prompt_token_ids: vec![],
+        token_ids: vec![0xe4, 0xbd, 0xa0, SPECIAL_TOKEN_ID],
+        skip_special_tokens: false,
+        min_bytes_to_buffer: 0,
+        truncate_output_to: None,
+        expected_chunks: vec![
+            decoded(
+                "你",
+                &[0xe4, 0xbd, 0xa0],
+                &[visible(0), visible(0), visible(0)],
+            ),
+            decoded("<special>", &[SPECIAL_TOKEN_ID], &[visible(0)]),
+        ],
+        expected_full: decoded(
+            "你<special>",
+            &[0xe4, 0xbd, 0xa0, SPECIAL_TOKEN_ID],
+            &[visible(0), visible(0), visible(0), visible(3)],
+        ),
+    });
+}
+
+#[test]
+fn attribution_filtered_special_between_incomplete_utf8_and_independent_token() {
+    // A skipped special token sits between the pending byte and the independent
+    // token; it stays zero-width in generation order and the split still applies.
+    let tokenizer = tokenizer_with_special_token();
+    run_attribution_case(AttributionCase {
+        tokenizer: &tokenizer,
+        prompt_token_ids: vec![],
+        token_ids: vec![0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
+        skip_special_tokens: true,
+        min_bytes_to_buffer: 0,
+        truncate_output_to: None,
+        expected_chunks: vec![decoded(
+            "�a",
+            &[0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
+            &[visible(0), zero_width(0), visible(3)],
+        )],
+        expected_full: decoded(
+            "�a",
+            &[0xe4, SPECIAL_TOKEN_ID, b'a' as u32],
+            &[visible(0), zero_width(0), visible(3)],
+        ),
+    });
+}
+
 #[test]
 fn attribution_truncation_converts_fully_removed_token_to_zero_width() {
     run_attribution_case(AttributionCase {