[#17156][fix] Flush buffered text in DeepSeekR1Parser.finish()

tensorrt_llm/llmapi/reasoning_parser.py

@@ -210,6 +210,31 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult:
         raise RuntimeError(
             "Unreachable code reached in `DeepSeekR1Parser.parse_delta`")
 
+    def finish(self) -> ReasoningParserResult:
+        """Flush text withheld by `parse_delta` when the stream ends.
+
+        `parse_delta` holds back a trailing fragment that could still grow
+        into a `<think>` / `</think>` tag. If the stream ends while such a
+        fragment is buffered - because the response was truncated mid-tag, or
+        simply ends with a literal `<` - the fragment is ordinary model output
+        and must still be emitted, otherwise it is silently dropped. The
+        buffered text is attributed to the block it was withheld in: inside a
+        reasoning block it is reasoning content, otherwise it is visible
+        content.
+
+        A buffer holding exactly a complete tag is a delimiter rather than
+        model output, so it is discarded as it would have been had more text
+        followed.
+        """
+        remaining = self._buffer
+        self._buffer = ""
+        if not remaining or remaining in (self.reasoning_start,
+                                          self.reasoning_end):
+            return ReasoningParserResult()
+        if self.in_reasoning:
+            return ReasoningParserResult(reasoning_content=remaining)
+        return ReasoningParserResult(content=remaining)
+
 
 @register_reasoning_parser("deepseek_v4")
 class DeepSeekV4ReasoningParser(BaseReasoningParser):

tests/unittest/llmapi/test_reasoning_parser.py

@@ -62,6 +62,79 @@ def test_deepseek_r1_reasoning_parser_stream(delta_texts: list, content: list,
         assert result.reasoning_content == reasoning_context[i]
 
 
+@pytest.mark.parametrize(
+    ("parser_key", "text"),
+    [
+        # `finish()` flushes as reasoning: the stream starts inside the
+        # reasoning block, so the withheld `<` is reasoning output.
+        ("deepseek-r1", "a <"),
+        # `finish()` flushes as content: no reasoning block was entered.
+        ("qwen3", "a<"),
+        # `finish()` discards: the buffer holds exactly a delimiter. Passes on
+        # `main` too - it guards against a fix that leaks the tag instead.
+        ("deepseek-r1", f"a{R1_END}"),
+    ])
+def test_deepseek_r1_reasoning_parser_stream_matches_non_stream(
+        parser_key: str, text: str) -> None:
+    """Streaming char-by-char then finishing must match a non-streaming parse.
+
+    One `(parser_key, text)` pair per branch of `finish()`. This is the
+    contract the missing flush violated, so it subsumes example-based tests
+    of the individual branches.
+    """
+    expected = ReasoningParserFactory.create_reasoning_parser(parser_key).parse(
+        text)
+    reasoning_parser = ReasoningParserFactory.create_reasoning_parser(
+        parser_key)
+    content, reasoning_context = "", ""
+    for char in text:
+        result = reasoning_parser.parse_delta(char)
+        content += result.content
+        reasoning_context += result.reasoning_content
+    result = reasoning_parser.finish()
+    content += result.content
+    reasoning_context += result.reasoning_content
+    assert content == expected.content
+    assert reasoning_context == expected.reasoning_content
+
+
+def test_deepseek_r1_reasoning_parser_finish_flushes_partial_tag() -> None:
+    """A partial tag arriving alongside text must still be flushed.
+
+    Such a delta fills `_buffer` through the `rfind` branch of `parse_delta`,
+    which one-character-at-a-time streaming never reaches - and it is the
+    shape a real stream delivers.
+    """
+    reasoning_parser = ReasoningParserFactory.create_reasoning_parser(
+        "deepseek-r1")
+    assert reasoning_parser.parse_delta("a </thin").reasoning_content == "a "
+    assert reasoning_parser.finish().reasoning_content == "</thin"
+
+
+@pytest.mark.parametrize(("chat_template_kwargs", "flushed"), [
+    ({
+        "thinking": True
+    }, "<"),
+    ({
+        "thinking": False
+    }, ""),
+])
+def test_deepseek_v4_reasoning_parser_finish_delegates(
+        chat_template_kwargs: dict[str, bool], flushed: str) -> None:
+    """`finish()` must reach whichever parser the thinking flag selected.
+
+    `DeepSeekV4ReasoningParser` delegates to two different targets:
+    `DeepSeekR1Parser`, which now flushes, and `IdentityReasoningParser`,
+    which withholds nothing to flush.
+    """
+    reasoning_parser = ReasoningParserFactory.create_reasoning_parser(
+        "deepseek_v4", chat_template_kwargs)
+    reasoning_parser.parse_delta("a <")
+    result = reasoning_parser.finish()
+    assert result.reasoning_content == flushed
+    assert result.content == ""
+
+
 @pytest.mark.parametrize("chat_template_kwargs", [{
     "thinking": True
 }, {