fix: prevent duplicate LLM stream snapshots

assistant_agent.py

@@ -6,6 +6,8 @@ import urllib.request
 from collections.abc import Iterator
 from typing import Any
 
+from llm_stream import OpenAIStreamAccumulator
+
 
 class ReviewAssistantError(RuntimeError):
     pass
@@ -41,6 +43,7 @@ def stream_review_assistant(
     try:
         with urllib.request.urlopen(request, timeout=timeout) as response:
             yielded = False
+            accumulator = OpenAIStreamAccumulator()
             for raw_line in response:
                 line = raw_line.decode("utf-8", errors="replace").strip()
                 if not line or line.startswith(":"):
@@ -57,10 +60,7 @@ def stream_review_assistant(
                 if not choices:
                     continue
                 choice = choices[0] or {}
-                delta = choice.get("delta") or {}
+                content = accumulator.feed(choice)
-                content = delta.get("content")
-                if content is None:
-                    content = (choice.get("message") or {}).get("content")
                 if content:
                     yielded = True
                     yield str(content)

llm_stream.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class OpenAIStreamAccumulator:
+    """Normalize incremental deltas and provider-specific full-message snapshots."""
+
+    def __init__(self) -> None:
+        self.text = ""
+        self.saw_delta = False
+
+    def feed(self, choice: dict[str, Any]) -> str:
+        delta = choice.get("delta")
+        if isinstance(delta, dict) and delta.get("content") is not None:
+            chunk = str(delta.get("content") or "")
+            if chunk:
+                self.saw_delta = True
+                self.text += chunk
+            return chunk
+
+        message = choice.get("message")
+        if not isinstance(message, dict) or message.get("content") is None:
+            return ""
+        snapshot = str(message.get("content") or "")
+        if not snapshot:
+            return ""
+        if not self.text:
+            self.text = snapshot
+            return snapshot
+        if snapshot == self.text or self.text.startswith(snapshot):
+            return ""
+        if snapshot.startswith(self.text):
+            suffix = snapshot[len(self.text):]
+            self.text = snapshot
+            return suffix
+        if self.saw_delta:
+            # A final full snapshot cannot safely replace chunks already delivered.
+            return ""
+        return ""

mentor_agent.py

@@ -10,6 +10,8 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
+from llm_stream import OpenAIStreamAccumulator
+
 
 class MentorAgentError(RuntimeError):
     pass
@@ -211,6 +213,7 @@ def stream_with_mentor(
     try:
         with urllib.request.urlopen(request, timeout=timeout) as response:
             yielded = False
+            accumulator = OpenAIStreamAccumulator()
             for raw_line in response:
                 line = raw_line.decode("utf-8", errors="replace").strip()
                 if not line or line.startswith(":"):
@@ -227,9 +230,7 @@ def stream_with_mentor(
                 if not choices:
                     continue
                 choice = choices[0] or {}
-                content = (choice.get("delta") or {}).get("content")
+                content = accumulator.feed(choice)
-                if content is None:
-                    content = (choice.get("message") or {}).get("content")
                 if content:
                     yielded = True
                     yield str(content)

tests/test_llm_stream.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import unittest
+
+from llm_stream import OpenAIStreamAccumulator
+
+
+class OpenAIStreamAccumulatorTests(unittest.TestCase):
+    def test_repeated_delta_chunks_are_preserved_as_model_output(self) -> None:
+        accumulator = OpenAIStreamAccumulator()
+        self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes")
+        self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes")
+        self.assertEqual(accumulator.text, "yesyes")
+
+    def test_final_snapshot_can_add_a_missing_suffix(self) -> None:
+        accumulator = OpenAIStreamAccumulator()
+        self.assertEqual(accumulator.feed({"delta": {"content": "first"}}), "first")
+        self.assertEqual(
+            accumulator.feed({"message": {"content": "first second"}}),
+            " second",
+        )
+        self.assertEqual(accumulator.text, "first second")
+
+    def test_incompatible_final_snapshot_is_not_appended_twice(self) -> None:
+        accumulator = OpenAIStreamAccumulator()
+        accumulator.feed({"delta": {"content": "streamed answer"}})
+        self.assertEqual(
+            accumulator.feed({"message": {"content": "rewritten final answer"}}),
+            "",
+        )
+        self.assertEqual(accumulator.text, "streamed answer")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_mentor_stream.py

@@ -66,6 +66,43 @@ class MentorStreamTests(unittest.TestCase):
             )
         self.assertEqual(result["answer"], "first second")
 
+    def test_final_full_message_does_not_duplicate_streamed_deltas(self):
+        lines = [
+            b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
+            b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
+            b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
+            b"data: [DONE]\n",
+        ]
+        with patch(
+            "mentor_agent.urllib.request.urlopen",
+            return_value=FakeStreamResponse(lines),
+        ):
+            chunks = list(
+                stream_with_mentor(
+                    self.skill, {}, "question", [], "key",
+                    "https://example.test/v1", "model",
+                )
+            )
+        self.assertEqual(chunks, ["first", " second"])
+
+    def test_snapshot_only_stream_emits_only_new_suffix(self):
+        lines = [
+            b'data: {"choices":[{"message":{"content":"first"}}]}\n',
+            b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
+            b"data: [DONE]\n",
+        ]
+        with patch(
+            "mentor_agent.urllib.request.urlopen",
+            return_value=FakeStreamResponse(lines),
+        ):
+            chunks = list(
+                stream_with_mentor(
+                    self.skill, {}, "question", [], "key",
+                    "https://example.test/v1", "model",
+                )
+            )
+        self.assertEqual(chunks, ["first", " second"])
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_review_assistant.py

@@ -45,6 +45,23 @@ class ReviewAssistantStreamingTests(unittest.TestCase):
                     stream_review_assistant({}, "question", [], "key", "https://example.test/v1", "model")
                 )
 
+    def test_final_full_message_does_not_duplicate_streamed_deltas(self):
+        response = StreamingResponse(
+            [
+                b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
+                b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
+                b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
+                b"data: [DONE]\n",
+            ]
+        )
+        with patch("assistant_agent.urllib.request.urlopen", return_value=response):
+            chunks = list(
+                stream_review_assistant(
+                    {}, "question", [], "key", "https://example.test/v1", "model"
+                )
+            )
+        self.assertEqual(chunks, ["first", " second"])
+
     def test_ndjson_event_writer_flushes_complete_line(self):
         handler = RequestHandler.__new__(RequestHandler)
         handler.wfile = io.BytesIO()