diff --git a/packages/sie_server/src/sie_server/api/openai_local.py b/packages/sie_server/src/sie_server/api/openai_local.py
index f708fdf1f..f1e135ca4 100644
--- a/packages/sie_server/src/sie_server/api/openai_local.py
+++ b/packages/sie_server/src/sie_server/api/openai_local.py
@@ -43,6 +43,7 @@
from sie_sdk.queue_types import denormalize_model_id
from sie_server.adapters._generation_base import (
+ REASONING_CONTROL_TOKENS,
GenerationAdapter,
GenerationError,
GenerationUnsupportedFieldError,
@@ -131,6 +132,7 @@
_MAX_CHAT_VIDEOS = 1
_MAX_U32 = (1 << 32) - 1
_ALLOWED_CHAT_ROLES = frozenset({"system", "user", "assistant", "tool", "developer"})
+_ANSWER_FINISH_REASONS = frozenset({"stop", "tool_calls"})
# Compatibility requests project onto the same native score bound.
_MAX_RERANK_DOCS = MAX_SCORE_ITEMS
@@ -1004,6 +1006,7 @@ def _sanitize_chat_payload(
initial_inside_thinking: bool,
reasoning_format: ReasoningFormat,
states: dict[int, ThinkingBlockStripper] | None = None,
+ held_contents: dict[int, list[str]] | None = None,
terminal: bool,
include_reasoning: bool = False,
) -> dict[str, Any]:
@@ -1016,6 +1019,7 @@ def _sanitize_chat_payload(
public_reasoning = include_reasoning and hide_thinking_blocks
state_map = states if states is not None else {}
+ held_map = held_contents if held_contents is not None else {}
for position, choice in enumerate(choices):
if not isinstance(choice, dict):
continue
@@ -1037,18 +1041,37 @@ def _sanitize_chat_payload(
if not hide_thinking_blocks:
continue
- stripper = state_map.setdefault(
- index,
- ThinkingBlockStripper(
+ raw_content = container.get("content")
+ content = raw_content if isinstance(raw_content, str) else ""
+ stripper = state_map.get(index)
+ if stripper is None:
+ if not content and not reasoning_was_present:
+ continue
+ stripper = state_map[index] = ThinkingBlockStripper(
inside=initial_inside_thinking,
reasoning_format=reasoning_format,
keep_reasoning=public_reasoning,
- ),
- )
- raw_content = container.get("content")
- content = raw_content if isinstance(raw_content, str) else ""
- visible = stripper.feed(content) if content else ""
+ )
+ if initial_inside_thinking and reasoning_was_present:
+ held_map[index] = []
is_terminal_choice = terminal or choice_obj.get("finish_reason") is not None
+ text = content
+ held = held_map.get(index)
+ if held is not None:
+ close = REASONING_CONTROL_TOKENS[reasoning_format][1]
+ closes = close in "".join(held[1 - len(close) :]) + content
+ if content:
+ held.append(content)
+ if closes or is_terminal_choice:
+ del held_map[index]
+ text = "".join(held)
+ if not closes and choice_obj.get("finish_reason") in _ANSWER_FINISH_REASONS:
+ stripper = state_map[index] = ThinkingBlockStripper(
+ reasoning_format=reasoning_format, keep_reasoning=public_reasoning
+ )
+ else:
+ text = ""
+ visible = stripper.feed(text) if text else ""
if is_terminal_choice:
visible += stripper.finish()
if isinstance(raw_content, str) or visible:
@@ -1093,6 +1116,7 @@ async def _sanitize_sse_stream(
include_reasoning: bool = False,
) -> AsyncIterator[bytes]:
states: dict[int, ThinkingBlockStripper] = {}
+ held_contents: dict[int, list[str]] = {}
strict_contents: dict[int, list[str]] = {}
strict_tool_choices: set[int] = set()
tool_calls: dict[int, set[object]] = {}
@@ -1135,6 +1159,7 @@ async def _sanitize_sse_stream(
initial_inside_thinking=initial_inside_thinking,
reasoning_format=reasoning_format,
states=states,
+ held_contents=held_contents,
terminal=False,
include_reasoning=include_reasoning,
)
diff --git a/packages/sie_server/tests/api/test_openai_local.py b/packages/sie_server/tests/api/test_openai_local.py
index 0b25ccd3c..035d756a9 100644
--- a/packages/sie_server/tests/api/test_openai_local.py
+++ b/packages/sie_server/tests/api/test_openai_local.py
@@ -2418,3 +2418,243 @@ def _handler(request: httpx.Request) -> httpx.Response:
reasoning = "".join(delta.get("reasoning_content", "") for delta in deltas)
assert reasoning == ("private default reasoning" if include_reasoning else "")
assert "" not in response.text
+
+
+_MLX_REASONING = "The user wants one word.\n"
+_MLX_ANSWER = "\n\nready"
+
+
+@pytest.mark.parametrize("include_reasoning", [False, True])
+@pytest.mark.parametrize(
+ ("message", "finish_reason", "content", "reasoning", "hidden"),
+ [
+ (
+ {"role": "assistant", "content": _MLX_ANSWER, "reasoning": _MLX_REASONING},
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ {"role": "assistant", "content": _MLX_ANSWER, "reasoning_content": _MLX_REASONING},
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ {"role": "assistant", "content": _MLX_REASONING + "" + _MLX_ANSWER},
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ {"role": "assistant", "content": "one word.\n" + _MLX_ANSWER, "reasoning": "The user wants "},
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ {"role": "assistant", "content": "readyone word", "reasoning": _MLX_REASONING},
+ "stop",
+ "ready",
+ _MLX_REASONING + "one word",
+ "one word",
+ ),
+ (
+ {"role": "assistant", "content": "still thinking", "reasoning": _MLX_REASONING},
+ "length",
+ "",
+ _MLX_REASONING + "still thinking",
+ "still thinking",
+ ),
+ ({"role": "assistant", "content": _MLX_ANSWER}, "stop", "", _MLX_ANSWER, "ready"),
+ ],
+ ids=[
+ "reasoning",
+ "reasoning_content",
+ "raw_tags",
+ "split_and_raw_close",
+ "split_then_raw_open",
+ "split_cut_off",
+ "no_reasoning_field",
+ ],
+)
+def test_mlx_thinking_chat_keeps_the_answer_beside_split_reasoning(
+ monkeypatch: pytest.MonkeyPatch,
+ message: dict[str, Any],
+ finish_reason: str,
+ content: str,
+ reasoning: str,
+ hidden: str,
+ include_reasoning: bool,
+) -> None:
+ payload = {
+ "id": "chatcmpl-mlx",
+ "system_fingerprint": "0.32.0",
+ "object": "chat.completion",
+ "model": "mlx-community/Qwen3.5-4B-4bit",
+ "created": 1,
+ "choices": [{"index": 0, "finish_reason": finish_reason, "message": message}],
+ "usage": {"prompt_tokens": 14, "completion_tokens": 12, "total_tokens": 26},
+ }
+ client, _ = _mlx_chat_client(
+ monkeypatch, lambda _request: httpx.Response(200, json=payload), chat_template_kwargs={"enable_thinking": True}
+ )
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "Qwen/Qwen3.5-4B:thinking",
+ "messages": [{"role": "user", "content": "Reply with one word: ready"}],
+ "include_reasoning": include_reasoning,
+ },
+ )
+
+ assert response.status_code == 200
+ returned = response.json()["choices"][0]["message"]
+ assert returned["content"] == content
+ assert "reasoning" not in returned
+ assert returned.get("reasoning_content") == (reasoning if include_reasoning else None)
+ assert include_reasoning or hidden not in response.text
+
+
+@pytest.mark.parametrize("include_reasoning", [False, True])
+@pytest.mark.parametrize(
+ ("deltas", "finish_reason", "content", "reasoning", "hidden"),
+ [
+ (
+ [
+ {"role": "assistant", "reasoning": "The user wants "},
+ {"role": "assistant", "reasoning": "one word.\n"},
+ {"role": "assistant", "content": "\n\n"},
+ {"role": "assistant", "content": "ready"},
+ ],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [
+ {"role": "assistant", "content": ""},
+ {"reasoning_content": "The user wants "},
+ {"reasoning_content": "one word.\n"},
+ {"content": "\n\nready"},
+ ],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [
+ {"role": "assistant", "content": "The user wants "},
+ {"role": "assistant", "content": "one word.\n\n\nready"},
+ ],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [
+ {"role": "assistant", "content": ""},
+ {"reasoning_content": "The user wants "},
+ {"content": "one word.\n\n\nready"},
+ ],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [
+ {"role": "assistant", "reasoning": "The user wants "},
+ {"role": "assistant", "content": "one word.\n\n\nready"},
+ ],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [{"role": "assistant", "reasoning": _MLX_REASONING, "content": _MLX_ANSWER}],
+ "stop",
+ _MLX_ANSWER,
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [
+ {"role": "assistant", "reasoning": _MLX_REASONING},
+ {"role": "assistant", "content": "Checking."},
+ {"role": "assistant", "tool_calls": [{"index": 0, **_PARIS_CALL}]},
+ ],
+ "tool_calls",
+ "Checking.",
+ _MLX_REASONING,
+ "one word",
+ ),
+ (
+ [{"role": "assistant", "reasoning": _MLX_REASONING}, {"role": "assistant", "content": "still thinking"}],
+ "length",
+ "",
+ _MLX_REASONING + "still thinking",
+ "still thinking",
+ ),
+ ([{"role": "assistant", "content": _MLX_ANSWER}], "stop", "", _MLX_ANSWER, "ready"),
+ ],
+ ids=[
+ "reasoning",
+ "reasoning_content_after_role",
+ "raw_tags",
+ "split_then_raw_close",
+ "split_then_raw_close_across_deltas",
+ "split_and_answer_in_one_delta",
+ "answer_before_tool_call",
+ "split_cut_off",
+ "no_reasoning_field",
+ ],
+)
+def test_mlx_thinking_chat_stream_keeps_the_answer_beside_split_reasoning(
+ monkeypatch: pytest.MonkeyPatch,
+ deltas: list[dict[str, Any]],
+ finish_reason: str,
+ content: str,
+ reasoning: str,
+ hidden: str,
+ include_reasoning: bool,
+) -> None:
+ wire = _sse_wire(
+ [*(_stream_event(delta) for delta in deltas), _stream_event({"role": "assistant"}, finish_reason=finish_reason)]
+ )
+ client, _ = _mlx_chat_client(
+ monkeypatch,
+ lambda _request: httpx.Response(
+ 200, stream=_ChunkedAsyncStream([wire[:83], wire[83:]]), headers={"content-type": "text/event-stream"}
+ ),
+ chat_template_kwargs={"enable_thinking": True},
+ )
+
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "Qwen/Qwen3.5-4B:thinking",
+ "messages": [{"role": "user", "content": "Reply with one word: ready"}],
+ "tools": [_WEATHER_TOOL],
+ "stream": True,
+ "include_reasoning": include_reasoning,
+ },
+ )
+
+ assert response.status_code == 200
+ returned = _sse_deltas(response.text)
+ assert "".join(delta.get("content") or "" for delta in returned) == content
+ assert not any("reasoning" in delta for delta in returned)
+ assert "".join(delta.get("reasoning_content", "") for delta in returned) == (reasoning if include_reasoning else "")
+ assert include_reasoning or hidden not in response.text