From a1c2c975cb8456a60618302e03c937321d72c24f Mon Sep 17 00:00:00 2001
From: fdc310 <2213070223@qq.com>
Date: Thu, 6 Aug 2026 01:50:20 +0800
Subject: [PATCH] fix(provider): preserve think tags in streamed reasoning

---
 .../modelmgr/requesters/litellmchat.py        | 37 +++++++++++++---
 .../provider/test_reasoning_control.py        | 42 +++++++++++++++++++
 2 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
index 002027f0b..4267f0190 100644
--- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
+++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py
@@ -1135,6 +1135,8 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
         role = 'assistant'
         tool_call_state: dict[int, dict[str, typing.Any]] = {}
         think_state = _ThinkStripState() if remove_think else None
+        reasoning_started = False
+        reasoning_closed = False
 
         try:
             response = await acompletion(**args)
@@ -1165,18 +1167,36 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
                 if 'role' in delta and delta['role']:
                     role = delta['role']
 
-                delta_content = delta.get('content', '')
-                reasoning_content = delta.get('reasoning_content', '')
+                delta_content = delta.get('content') or ''
+                reasoning_content = delta.get('reasoning_content') or ''
                 provider_fields = dict(delta.get('provider_specific_fields') or {})
 
                 # Handle reasoning_content based on remove_think flag
                 if reasoning_content:
                     provider_fields['reasoning_content'] = reasoning_content
                     if remove_think:
-                        delta_content = None
+                        delta_content = delta_content or None
                     else:
-                        # Use reasoning_content as the displayed content
-                        delta_content = reasoning_content
+                        # Stream explicit markers so downstream adapters and
+                        # the debug page see the same format as non-streaming
+                        # responses.
+                        if not reasoning_started:
+                            delta_content = '<think>\n'
+                            reasoning_started = True
+                        else:
+                            delta_content = ''
+                        delta_content += reasoning_content
+                        if delta.get('content'):
+                            delta_content += f'\n</think>\n{delta.get("content")}'
+                            reasoning_closed = True
+
+                elif delta_content and not remove_think and reasoning_started and not reasoning_closed:
+                    delta_content = f'\n</think>\n{delta_content}'
+                    reasoning_closed = True
+
+                if finish_reason and not remove_think and reasoning_started and not reasoning_closed:
+                    delta_content = f'{delta_content}\n</think>\n'
+                    reasoning_closed = True
 
                 if think_state is not None and delta_content:
                     delta_content = think_state.feed(delta_content)
@@ -1205,6 +1225,13 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
                 yield provider_message.MessageChunk(**chunk_data)
                 chunk_idx += 1
 
+            if reasoning_started and not reasoning_closed:
+                yield provider_message.MessageChunk(
+                    role=role,
+                    content='\n</think>\n',
+                    is_final=True,
+                )
+
             if think_state is not None:
                 pending_content = think_state.flush()
                 if pending_content:
diff --git a/tests/unit_tests/provider/test_reasoning_control.py b/tests/unit_tests/provider/test_reasoning_control.py
index 35056a86c..2f4b24dc8 100644
--- a/tests/unit_tests/provider/test_reasoning_control.py
+++ b/tests/unit_tests/provider/test_reasoning_control.py
@@ -714,3 +714,45 @@ async def test_stream_reasoning_round_trip_with_hidden_display(monkeypatch):
     assert emitted is not None
     assert emitted.content == 'answer'
     assert emitted.provider_specific_fields == {'reasoning_content': 'private '}
+
+
+@pytest.mark.asyncio
+async def test_stream_reasoning_content_is_wrapped_for_display(monkeypatch):
+    request = _requester('deepseek')
+    request._build_completion_args = AsyncMock(return_value={})
+
+    async def chunks():
+        yield SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    delta=_Dumpable({'role': 'assistant', 'reasoning_content': 'private '}),
+                    finish_reason=None,
+                )
+            ],
+            usage=None,
+        )
+        yield SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    delta=_Dumpable({'content': 'answer'}),
+                    finish_reason='stop',
+                )
+            ],
+            usage=None,
+        )
+
+    monkeypatch.setattr(litellmchat, 'acompletion', AsyncMock(return_value=chunks()))
+    accumulator = _StreamAccumulator(remove_think=False)
+    emitted: provider_message.MessageChunk | None = None
+
+    async for chunk in request.invoke_llm_stream(
+        None,
+        _runtime_model(request),
+        [],
+        remove_think=False,
+    ):
+        emitted = accumulator.add(chunk) or emitted
+
+    assert emitted is not None
+    assert emitted.content == '<think>\nprivate \n</think>\nanswer'
+    assert emitted.provider_specific_fields == {'reasoning_content': 'private '}