← back to Exo
Add thinking in thinking blocks, and fix DeepSeek interleaved tool calls (#1548)
cf648a53b84bd0bf50a3e61d7fcd0bbb9b9ac457 · 2026-02-19 18:44:49 +0000 · rltakashige
## Motivation
OpenCode shows <think> tags and not thinking blocks as we aren't
following the API specs properly.
Claude was also getting horrible prefix cache hits because it sends
headers.
## Changes
Handle thinking tokens properly by placing them in think tags for each
of the API endpoints.
Also support DeepSeekV3.2 tool calling properly as a minor feature.
Strips Claude headers at the API level.
## Test Plan
### Manual Testing
Tested OpenCode manually
Needs testing with Claude.
### Automated Testing
All CI and tests passing - added a new e2e test for DeepSeekV32 tool
parsing.
Files touched
M dashboard/src/lib/stores/app.svelte.tsM src/exo/master/adapters/chat_completions.pyM src/exo/master/adapters/claude.pyM src/exo/master/adapters/responses.pyM src/exo/master/tests/test_claude_tool_use.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/claude_api.pyM src/exo/shared/types/openai_responses.pyM src/exo/shared/types/worker/runner_response.pyA src/exo/worker/engines/mlx/dsml_encoding.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/runner.pyA src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.pyM src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
Diff
commit cf648a53b84bd0bf50a3e61d7fcd0bbb9b9ac457
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Feb 19 18:44:49 2026 +0000
Add thinking in thinking blocks, and fix DeepSeek interleaved tool calls (#1548)
## Motivation
OpenCode shows <think> tags and not thinking blocks as we aren't
following the API specs properly.
Claude was also getting horrible prefix cache hits because it sends
headers.
## Changes
Handle thinking tokens properly by placing them in think tags for each
of the API endpoints.
Also support DeepSeekV3.2 tool calling properly as a minor feature.
Strips Claude headers at the API level.
## Test Plan
### Manual Testing
Tested OpenCode manually
Needs testing with Claude.
### Automated Testing
All CI and tests passing - added a new e2e test for DeepSeekV32 tool
parsing.
---
dashboard/src/lib/stores/app.svelte.ts | 86 +-
src/exo/master/adapters/chat_completions.py | 21 +-
src/exo/master/adapters/claude.py | 114 ++-
src/exo/master/adapters/responses.py | 237 ++++-
src/exo/master/tests/test_claude_tool_use.py | 9 +-
src/exo/shared/types/api.py | 2 +-
src/exo/shared/types/chunks.py | 1 +
src/exo/shared/types/claude_api.py | 35 +-
src/exo/shared/types/openai_responses.py | 74 +-
src/exo/shared/types/worker/runner_response.py | 1 +
src/exo/worker/engines/mlx/dsml_encoding.py | 72 ++
src/exo/worker/engines/mlx/utils_mlx.py | 27 +-
src/exo/worker/runner/runner.py | 220 ++++-
.../tests/unittests/test_runner/test_dsml_e2e.py | 967 +++++++++++++++++++++
.../unittests/test_runner/test_event_ordering.py | 1 +
.../unittests/test_runner/test_parse_gpt_oss.py | 23 +-
16 files changed, 1772 insertions(+), 118 deletions(-)
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index d13e7f2b..3e0074e3 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -1652,11 +1652,12 @@ class AppStore {
if (!reader) throw new Error("No response body");
let fullContent = prefixText;
+ let streamedThinking = "";
const collectedTokens: TokenData[] = [...tokensToKeep];
interface ChatCompletionChunk {
choices?: Array<{
- delta?: { content?: string };
+ delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1677,6 +1678,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
+ const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1695,7 +1697,11 @@ class AppStore {
}
}
- if (delta) {
+ if (thinkingDelta) {
+ streamedThinking += thinkingDelta;
+ }
+
+ if (delta || thinkingDelta) {
if (firstTokenTime === null) {
firstTokenTime = performance.now();
this.ttftMs = firstTokenTime - requestStartTime;
@@ -1709,9 +1715,14 @@ class AppStore {
this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000;
}
- fullContent += delta;
- const { displayContent, thinkingContent } =
+ if (delta) {
+ fullContent += delta;
+ }
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
+ const combinedThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
if (this.activeConversationId === targetConversationId) {
this.currentResponse = displayContent;
@@ -1723,7 +1734,7 @@ class AppStore {
messageId,
(m) => {
m.content = displayContent;
- m.thinking = thinkingContent || undefined;
+ m.thinking = combinedThinking || undefined;
m.tokens = [...collectedTokens];
},
);
@@ -1735,11 +1746,14 @@ class AppStore {
// Final update
if (this.conversationExists(targetConversationId)) {
- const { displayContent, thinkingContent } =
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
+ const finalThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
this.updateConversationMessage(targetConversationId, messageId, (m) => {
m.content = displayContent;
- m.thinking = thinkingContent || undefined;
+ m.thinking = finalThinking || undefined;
m.tokens = [...collectedTokens];
if (this.ttftMs !== null) m.ttftMs = this.ttftMs;
if (this.tps !== null) m.tps = this.tps;
@@ -1847,11 +1861,12 @@ class AppStore {
}
let streamedContent = "";
+ let streamedThinking = "";
const collectedTokens: TokenData[] = [];
interface ChatCompletionChunk {
choices?: Array<{
- delta?: { content?: string };
+ delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1872,6 +1887,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
+ const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1890,10 +1906,19 @@ class AppStore {
}
}
- if (delta) {
- streamedContent += delta;
- const { displayContent, thinkingContent } =
+ if (thinkingDelta) {
+ streamedThinking += thinkingDelta;
+ }
+
+ if (delta || thinkingDelta) {
+ if (delta) {
+ streamedContent += delta;
+ }
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
+ const combinedThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -1906,7 +1931,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
- msg.thinking = thinkingContent || undefined;
+ msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -1918,14 +1943,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
- const { displayContent, thinkingContent } =
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
+ const finalThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
- msg.thinking = thinkingContent || undefined;
+ msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2317,10 +2345,11 @@ class AppStore {
}
let streamedContent = "";
+ let streamedThinking = "";
interface ChatCompletionChunk {
choices?: Array<{
- delta?: { content?: string };
+ delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -2348,6 +2377,7 @@ class AppStore {
const choice = parsed.choices?.[0];
const tokenContent = choice?.delta?.content;
+ const thinkingContent = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -2366,7 +2396,11 @@ class AppStore {
}
}
- if (tokenContent) {
+ if (thinkingContent) {
+ streamedThinking += thinkingContent;
+ }
+
+ if (tokenContent || thinkingContent) {
// Track first token for TTFT
if (firstTokenTime === null) {
firstTokenTime = performance.now();
@@ -2383,11 +2417,16 @@ class AppStore {
this.tps = (tokenCount / elapsed) * 1000;
}
- streamedContent += tokenContent;
+ if (tokenContent) {
+ streamedContent += tokenContent;
+ }
- // Strip thinking tags for display and extract thinking content
- const { displayContent, thinkingContent } =
+ // Use stripThinkingTags as fallback for any <think> tags still in content
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
+ const combinedThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -2400,7 +2439,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
- msg.thinking = thinkingContent || undefined;
+ msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2436,14 +2475,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
- const { displayContent, thinkingContent } =
+ const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
+ const finalThinking = [streamedThinking, tagThinking]
+ .filter(Boolean)
+ .join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
- msg.thinking = thinkingContent || undefined;
+ msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
// Store performance metrics on the message
if (this.ttftMs !== null) {
diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py
index a8ca8500..9773806f 100644
--- a/src/exo/master/adapters/chat_completions.py
+++ b/src/exo/master/adapters/chat_completions.py
@@ -59,7 +59,11 @@ def chat_request_to_text_generation(
chat_template_messages.append({"role": "system", "content": content})
else:
# Skip messages with no meaningful content
- if msg.content is None and msg.thinking is None and msg.tool_calls is None:
+ if (
+ msg.content is None
+ and msg.reasoning_content is None
+ and msg.tool_calls is None
+ ):
continue
if msg.role in ("user", "assistant", "developer"):
@@ -111,6 +115,11 @@ def chunk_to_response(
]
)
+ if chunk.is_thinking:
+ delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text)
+ else:
+ delta = ChatCompletionMessage(role="assistant", content=chunk.text)
+
return ChatCompletionResponse(
id=command_id,
created=int(time.time()),
@@ -118,7 +127,7 @@ def chunk_to_response(
choices=[
StreamingChoiceResponse(
index=0,
- delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+ delta=delta,
logprobs=logprobs,
finish_reason=chunk.finish_reason,
)
@@ -208,6 +217,7 @@ async def collect_chat_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ChatCompletionResponse."""
text_parts: list[str] = []
+ thinking_parts: list[str] = []
tool_calls: list[ToolCall] = []
logprobs_content: list[LogprobsContentItem] = []
model: str | None = None
@@ -228,7 +238,10 @@ async def collect_chat_response(
if model is None:
model = chunk.model
last_usage = chunk.usage or last_usage
- text_parts.append(chunk.text)
+ if chunk.is_thinking:
+ thinking_parts.append(chunk.text)
+ else:
+ text_parts.append(chunk.text)
if chunk.logprob is not None:
logprobs_content.append(
LogprobsContentItem(
@@ -258,6 +271,7 @@ async def collect_chat_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
+ combined_thinking = "".join(thinking_parts) if thinking_parts else None
assert model is not None
yield ChatCompletionResponse(
@@ -270,6 +284,7 @@ async def collect_chat_response(
message=ChatCompletionMessage(
role="assistant",
content=combined_text,
+ reasoning_content=combined_thinking,
tool_calls=tool_calls if tool_calls else None,
),
logprobs=Logprobs(content=logprobs_content)
diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py
index d5246bbc..7a31f6fd 100644
--- a/src/exo/master/adapters/claude.py
+++ b/src/exo/master/adapters/claude.py
@@ -29,6 +29,8 @@ from exo.shared.types.claude_api import (
ClaudeStopReason,
ClaudeTextBlock,
ClaudeTextDelta,
+ ClaudeThinkingBlock,
+ ClaudeThinkingDelta,
ClaudeToolResultBlock,
ClaudeToolUseBlock,
ClaudeUsage,
@@ -104,12 +106,15 @@ def claude_request_to_text_generation(
# Process structured content blocks
text_parts: list[str] = []
+ thinking_parts: list[str] = []
tool_calls: list[dict[str, Any]] = []
tool_results: list[ClaudeToolResultBlock] = []
for block in msg.content:
if isinstance(block, ClaudeTextBlock):
text_parts.append(block.text)
+ elif isinstance(block, ClaudeThinkingBlock):
+ thinking_parts.append(block.thinking)
elif isinstance(block, ClaudeToolUseBlock):
tool_calls.append(
{
@@ -125,6 +130,7 @@ def claude_request_to_text_generation(
tool_results.append(block)
content = "".join(text_parts)
+ reasoning_content = "".join(thinking_parts) if thinking_parts else None
# Build InputMessage from text content
if msg.role in ("user", "assistant"):
@@ -132,9 +138,14 @@ def claude_request_to_text_generation(
# Build chat_template_messages preserving tool structure
if tool_calls:
- chat_template_messages.append(
- {"role": "assistant", "content": content, "tool_calls": tool_calls}
- )
+ chat_msg: dict[str, Any] = {
+ "role": "assistant",
+ "content": content,
+ "tool_calls": tool_calls,
+ }
+ if reasoning_content:
+ chat_msg["reasoning_content"] = reasoning_content
+ chat_template_messages.append(chat_msg)
elif tool_results:
for tr in tool_results:
chat_template_messages.append(
@@ -145,7 +156,10 @@ def claude_request_to_text_generation(
}
)
else:
- chat_template_messages.append({"role": msg.role, "content": content})
+ chat_msg = {"role": msg.role, "content": content}
+ if reasoning_content:
+ chat_msg["reasoning_content"] = reasoning_content
+ chat_template_messages.append(chat_msg)
# Convert Claude tool definitions to OpenAI-style function tools
tools: list[dict[str, Any]] | None = None
@@ -162,6 +176,10 @@ def claude_request_to_text_generation(
for tool in request.tools
]
+ enable_thinking: bool | None = None
+ if request.thinking is not None:
+ enable_thinking = request.thinking.type in ("enabled", "adaptive")
+
return TextGenerationTaskParams(
model=request.model,
input=input_messages
@@ -175,6 +193,7 @@ def claude_request_to_text_generation(
stop=request.stop_sequences,
stream=request.stream,
tools=tools,
+ enable_thinking=enable_thinking,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
@@ -192,6 +211,7 @@ async def collect_claude_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ClaudeMessagesResponse."""
text_parts: list[str] = []
+ thinking_parts: list[str] = []
tool_use_blocks: list[ClaudeToolUseBlock] = []
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
@@ -219,7 +239,10 @@ async def collect_claude_response(
stop_reason = "tool_use"
continue
- text_parts.append(chunk.text)
+ if chunk.is_thinking:
+ thinking_parts.append(chunk.text)
+ else:
+ text_parts.append(chunk.text)
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -228,9 +251,12 @@ async def collect_claude_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
+ combined_thinking = "".join(thinking_parts)
# Build content blocks
content: list[ClaudeContentBlock] = []
+ if combined_thinking:
+ content.append(ClaudeThinkingBlock(thinking=combined_thinking))
if combined_text:
content.append(ClaudeTextBlock(text=combined_text))
content.extend(tool_use_blocks)
@@ -275,16 +301,16 @@ async def generate_claude_stream(
start_event = ClaudeMessageStartEvent(message=initial_message)
yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
- # content_block_start for text block at index 0
- block_start = ClaudeContentBlockStartEvent(
- index=0, content_block=ClaudeTextBlock(text="")
- )
- yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
-
output_tokens = 0
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
- next_block_index = 1 # text block is 0, tool blocks start at 1
+ next_block_index = 0
+
+ # Track whether we've started thinking/text blocks
+ thinking_block_started = False
+ thinking_block_index = -1
+ text_block_started = False
+ text_block_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -329,12 +355,45 @@ async def generate_claude_stream(
output_tokens += 1 # Count each chunk as one token
- # content_block_delta
- delta_event = ClaudeContentBlockDeltaEvent(
- index=0,
- delta=ClaudeTextDelta(text=chunk.text),
- )
- yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+ if chunk.is_thinking:
+ # Start thinking block on first thinking token
+ if not thinking_block_started:
+ thinking_block_started = True
+ thinking_block_index = next_block_index
+ next_block_index += 1
+ block_start = ClaudeContentBlockStartEvent(
+ index=thinking_block_index,
+ content_block=ClaudeThinkingBlock(thinking=""),
+ )
+ yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+ delta_event = ClaudeContentBlockDeltaEvent(
+ index=thinking_block_index,
+ delta=ClaudeThinkingDelta(thinking=chunk.text),
+ )
+ yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+ else:
+ # Close thinking block when transitioning to text
+ if thinking_block_started and text_block_index == -1:
+ block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
+ yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+ # Start text block on first text token
+ if not text_block_started:
+ text_block_started = True
+ text_block_index = next_block_index
+ next_block_index += 1
+ block_start = ClaudeContentBlockStartEvent(
+ index=text_block_index,
+ content_block=ClaudeTextBlock(text=""),
+ )
+ yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+ delta_event = ClaudeContentBlockDeltaEvent(
+ index=text_block_index,
+ delta=ClaudeTextDelta(text=chunk.text),
+ )
+ yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -343,9 +402,22 @@ async def generate_claude_stream(
if last_usage is not None:
output_tokens = last_usage.completion_tokens
- # content_block_stop for text block
- block_stop = ClaudeContentBlockStopEvent(index=0)
- yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+ # Close any open blocks
+ if thinking_block_started and text_block_index == -1:
+ block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
+ yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+ if text_block_started:
+ block_stop = ClaudeContentBlockStopEvent(index=text_block_index)
+ yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+ if not thinking_block_started and not text_block_started:
+ empty_start = ClaudeContentBlockStartEvent(
+ index=0, content_block=ClaudeTextBlock(text="")
+ )
+ yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n"
+ empty_stop = ClaudeContentBlockStopEvent(index=0)
+ yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n"
# message_delta
message_delta = ClaudeMessageDeltaEvent(
diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py
index 90fa7732..dc726ba7 100644
--- a/src/exo/master/adapters/responses.py
+++ b/src/exo/master/adapters/responses.py
@@ -29,6 +29,12 @@ from exo.shared.types.openai_responses import (
ResponseOutputItemAddedEvent,
ResponseOutputItemDoneEvent,
ResponseOutputText,
+ ResponseReasoningItem,
+ ResponseReasoningSummaryPartAddedEvent,
+ ResponseReasoningSummaryPartDoneEvent,
+ ResponseReasoningSummaryText,
+ ResponseReasoningSummaryTextDeltaEvent,
+ ResponseReasoningSummaryTextDoneEvent,
ResponsesRequest,
ResponsesResponse,
ResponsesStreamEvent,
@@ -141,7 +147,9 @@ async def collect_responses_response(
"""Collect all token chunks and return a single ResponsesResponse."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
+ reasoning_id = f"rs_{command_id}"
accumulated_text = ""
+ thinking_parts: list[str] = []
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
error_message: str | None = None
@@ -168,6 +176,10 @@ async def collect_responses_response(
)
continue
+ if chunk.is_thinking:
+ thinking_parts.append(chunk.text)
+ continue
+
accumulated_text += chunk.text
if error_message is not None:
@@ -182,13 +194,21 @@ async def collect_responses_response(
total_tokens=last_usage.total_tokens,
)
- output: list[ResponseItem] = [
+ output: list[ResponseItem] = []
+ if thinking_parts:
+ output.append(
+ ResponseReasoningItem(
+ id=reasoning_id,
+ summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))],
+ )
+ )
+ output.append(
ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text=accumulated_text)],
status="completed",
)
- ]
+ )
output.extend(function_call_items)
yield ResponsesResponse(
@@ -212,6 +232,7 @@ async def generate_responses_stream(
"""Generate OpenAI Responses API streaming events from TokenChunks."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
+ reasoning_id = f"rs_{command_id}"
seq = count(1)
# response.created
@@ -233,32 +254,17 @@ async def generate_responses_stream(
)
yield _format_sse(in_progress_event)
- # response.output_item.added
- initial_item = ResponseMessageItem(
- id=item_id,
- content=[ResponseOutputText(text="")],
- status="in_progress",
- )
- item_added = ResponseOutputItemAddedEvent(
- sequence_number=next(seq), output_index=0, item=initial_item
- )
- yield _format_sse(item_added)
-
- # response.content_part.added
- initial_part = ResponseOutputText(text="")
- part_added = ResponseContentPartAddedEvent(
- sequence_number=next(seq),
- item_id=item_id,
- output_index=0,
- content_index=0,
- part=initial_part,
- )
- yield _format_sse(part_added)
-
accumulated_text = ""
+ accumulated_thinking = ""
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
- next_output_index = 1 # message item is at 0
+ next_output_index = 0
+
+ # Track dynamic block creation
+ reasoning_started = False
+ reasoning_output_index = -1
+ message_started = False
+ message_output_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -327,23 +333,184 @@ async def generate_responses_stream(
next_output_index += 1
continue
+ if chunk.is_thinking:
+ # Start reasoning block on first thinking token
+ if not reasoning_started:
+ reasoning_started = True
+ reasoning_output_index = next_output_index
+ next_output_index += 1
+
+ # response.output_item.added for reasoning
+ reasoning_item = ResponseReasoningItem(
+ id=reasoning_id,
+ summary=[],
+ status="in_progress",
+ )
+ rs_added = ResponseOutputItemAddedEvent(
+ sequence_number=next(seq),
+ output_index=reasoning_output_index,
+ item=reasoning_item,
+ )
+ yield _format_sse(rs_added)
+
+ # response.reasoning_summary_part.added
+ part_added = ResponseReasoningSummaryPartAddedEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ part=ResponseReasoningSummaryText(text=""),
+ )
+ yield _format_sse(part_added)
+
+ accumulated_thinking += chunk.text
+
+ # response.reasoning_summary_text.delta
+ rs_delta = ResponseReasoningSummaryTextDeltaEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ delta=chunk.text,
+ )
+ yield _format_sse(rs_delta)
+ continue
+
+ # Close reasoning block when transitioning to text
+ if reasoning_started and not message_started:
+ # response.reasoning_summary_text.done
+ rs_text_done = ResponseReasoningSummaryTextDoneEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ text=accumulated_thinking,
+ )
+ yield _format_sse(rs_text_done)
+
+ # response.reasoning_summary_part.done
+ rs_part_done = ResponseReasoningSummaryPartDoneEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ part=ResponseReasoningSummaryText(text=accumulated_thinking),
+ )
+ yield _format_sse(rs_part_done)
+
+ # response.output_item.done for reasoning
+ rs_item_done = ResponseOutputItemDoneEvent(
+ sequence_number=next(seq),
+ output_index=reasoning_output_index,
+ item=ResponseReasoningItem(
+ id=reasoning_id,
+ summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+ ),
+ )
+ yield _format_sse(rs_item_done)
+
+ # Start message block on first text token
+ if not message_started:
+ message_started = True
+ message_output_index = next_output_index
+ next_output_index += 1
+
+ initial_item = ResponseMessageItem(
+ id=item_id,
+ content=[ResponseOutputText(text="")],
+ status="in_progress",
+ )
+ item_added = ResponseOutputItemAddedEvent(
+ sequence_number=next(seq),
+ output_index=message_output_index,
+ item=initial_item,
+ )
+ yield _format_sse(item_added)
+
+ initial_part = ResponseOutputText(text="")
+ part_added = ResponseContentPartAddedEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=message_output_index,
+ content_index=0,
+ part=initial_part,
+ )
+ yield _format_sse(part_added)
+
accumulated_text += chunk.text
# response.output_text.delta
delta_event = ResponseTextDeltaEvent(
sequence_number=next(seq),
item_id=item_id,
- output_index=0,
+ output_index=message_output_index,
content_index=0,
delta=chunk.text,
)
yield _format_sse(delta_event)
+ # Close reasoning block if it was never followed by text
+ if reasoning_started and not message_started:
+ rs_text_done = ResponseReasoningSummaryTextDoneEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ text=accumulated_thinking,
+ )
+ yield _format_sse(rs_text_done)
+
+ rs_part_done = ResponseReasoningSummaryPartDoneEvent(
+ sequence_number=next(seq),
+ item_id=reasoning_id,
+ output_index=reasoning_output_index,
+ summary_index=0,
+ part=ResponseReasoningSummaryText(text=accumulated_thinking),
+ )
+ yield _format_sse(rs_part_done)
+
+ rs_item_done = ResponseOutputItemDoneEvent(
+ sequence_number=next(seq),
+ output_index=reasoning_output_index,
+ item=ResponseReasoningItem(
+ id=reasoning_id,
+ summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+ ),
+ )
+ yield _format_sse(rs_item_done)
+
+ # If no message block was started, create one now (empty text)
+ if not message_started:
+ message_output_index = next_output_index
+ next_output_index += 1
+
+ initial_item = ResponseMessageItem(
+ id=item_id,
+ content=[ResponseOutputText(text="")],
+ status="in_progress",
+ )
+ item_added = ResponseOutputItemAddedEvent(
+ sequence_number=next(seq),
+ output_index=message_output_index,
+ item=initial_item,
+ )
+ yield _format_sse(item_added)
+
+ initial_part = ResponseOutputText(text="")
+ part_added_evt = ResponseContentPartAddedEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=message_output_index,
+ content_index=0,
+ part=initial_part,
+ )
+ yield _format_sse(part_added_evt)
+
# response.output_text.done
text_done = ResponseTextDoneEvent(
sequence_number=next(seq),
item_id=item_id,
- output_index=0,
+ output_index=message_output_index,
content_index=0,
text=accumulated_text,
)
@@ -354,7 +521,7 @@ async def generate_responses_stream(
part_done = ResponseContentPartDoneEvent(
sequence_number=next(seq),
item_id=item_id,
- output_index=0,
+ output_index=message_output_index,
content_index=0,
part=final_part,
)
@@ -367,7 +534,9 @@ async def generate_responses_stream(
status="completed",
)
item_done = ResponseOutputItemDoneEvent(
- sequence_number=next(seq), output_index=0, item=final_message_item
+ sequence_number=next(seq),
+ output_index=message_output_index,
+ item=final_message_item,
)
yield _format_sse(item_done)
@@ -381,7 +550,15 @@ async def generate_responses_stream(
)
# response.completed
- output: list[ResponseItem] = [final_message_item]
+ output: list[ResponseItem] = []
+ if reasoning_started:
+ output.append(
+ ResponseReasoningItem(
+ id=reasoning_id,
+ summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
+ )
+ )
+ output.append(final_message_item)
output.extend(function_call_items)
final_response = ResponsesResponse(
id=response_id,
diff --git a/src/exo/master/tests/test_claude_tool_use.py b/src/exo/master/tests/test_claude_tool_use.py
index 77ce3492..3e814dce 100644
--- a/src/exo/master/tests/test_claude_tool_use.py
+++ b/src/exo/master/tests/test_claude_tool_use.py
@@ -261,7 +261,7 @@ class TestGenerateClaudeStreamToolUse:
parsed = _parse_sse_events(events)
- # Two tool block starts (at indices 1 and 2)
+ # Two tool block starts (at indices 0 and 1 — no text block when only tools)
tool_starts = [
e
for e in parsed
@@ -270,12 +270,11 @@ class TestGenerateClaudeStreamToolUse:
== "tool_use"
]
assert len(tool_starts) == 2
- assert tool_starts[0]["index"] == 1
- assert tool_starts[1]["index"] == 2
+ assert tool_starts[0]["index"] == 0
+ assert tool_starts[1]["index"] == 1
- # Two tool block stops (at indices 1 and 2), plus text block stop at 0
+ # Two tool block stops (at indices 0 and 1)
block_stops = [e for e in parsed if e.get("type") == "content_block_stop"]
stop_indices = [e["index"] for e in block_stops]
assert 0 in stop_indices
assert 1 in stop_indices
- assert 2 in stop_indices
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 824c2c91..23ca9b7b 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -77,7 +77,7 @@ class ChatCompletionMessage(BaseModel):
content: (
str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None
) = None
- thinking: str | None = None # Added for GPT-OSS harmony format support
+ reasoning_content: str | None = None
name: str | None = None
tool_calls: list[ToolCall] | None = None
tool_call_id: str | None = None
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index b97eee5f..64481bab 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -27,6 +27,7 @@ class TokenChunk(BaseChunk):
stats: GenerationStats | None = None
logprob: float | None = None
top_logprobs: list[TopLogprobItem] | None = None
+ is_thinking: bool = False
class ErrorChunk(BaseChunk):
diff --git a/src/exo/shared/types/claude_api.py b/src/exo/shared/types/claude_api.py
index dcf8c515..8eb18497 100644
--- a/src/exo/shared/types/claude_api.py
+++ b/src/exo/shared/types/claude_api.py
@@ -47,6 +47,14 @@ class ClaudeImageBlock(BaseModel, frozen=True):
source: ClaudeImageSource
+class ClaudeThinkingBlock(BaseModel, frozen=True):
+ """Thinking content block in Claude Messages API."""
+
+ type: Literal["thinking"] = "thinking"
+ thinking: str
+ signature: str | None = None
+
+
class ClaudeToolUseBlock(BaseModel, frozen=True):
"""Tool use content block in Claude Messages API."""
@@ -66,11 +74,17 @@ class ClaudeToolResultBlock(BaseModel, frozen=True):
cache_control: dict[str, str] | None = None
-ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock
+ClaudeContentBlock = (
+ ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
+)
# Input content blocks can also include tool_result (sent by user after tool_use)
ClaudeInputContentBlock = (
- ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock
+ ClaudeTextBlock
+ | ClaudeImageBlock
+ | ClaudeThinkingBlock
+ | ClaudeToolUseBlock
+ | ClaudeToolResultBlock
)
@@ -82,6 +96,11 @@ class ClaudeMessage(BaseModel, frozen=True):
content: str | list[ClaudeInputContentBlock]
+class ClaudeThinkingConfig(BaseModel, frozen=True):
+ type: Literal["enabled", "disabled", "adaptive"]
+ budget_tokens: int | None = None
+
+
class ClaudeMessagesRequest(BaseModel):
"""Request body for Claude Messages API."""
@@ -96,6 +115,7 @@ class ClaudeMessagesRequest(BaseModel):
top_k: int | None = None
tools: list[ClaudeToolDefinition] | None = None
metadata: dict[str, str] | None = None
+ thinking: ClaudeThinkingConfig | None = None
# Response types
@@ -145,7 +165,7 @@ class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
type: Literal["content_block_start"] = "content_block_start"
index: int
- content_block: ClaudeTextBlock | ClaudeToolUseBlock
+ content_block: ClaudeTextBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
class ClaudeTextDelta(BaseModel, frozen=True):
@@ -155,6 +175,13 @@ class ClaudeTextDelta(BaseModel, frozen=True):
text: str
+class ClaudeThinkingDelta(BaseModel, frozen=True):
+ """Delta for thinking content block."""
+
+ type: Literal["thinking_delta"] = "thinking_delta"
+ thinking: str
+
+
class ClaudeInputJsonDelta(BaseModel, frozen=True):
"""Delta for tool use input JSON content block."""
@@ -167,7 +194,7 @@ class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
type: Literal["content_block_delta"] = "content_block_delta"
index: int
- delta: ClaudeTextDelta | ClaudeInputJsonDelta
+ delta: ClaudeTextDelta | ClaudeThinkingDelta | ClaudeInputJsonDelta
class ClaudeContentBlockStopEvent(BaseModel, frozen=True):
diff --git a/src/exo/shared/types/openai_responses.py b/src/exo/shared/types/openai_responses.py
index e29fb253..331f68e4 100644
--- a/src/exo/shared/types/openai_responses.py
+++ b/src/exo/shared/types/openai_responses.py
@@ -145,7 +145,23 @@ class ResponseFunctionCallItem(BaseModel, frozen=True):
status: ResponseStatus = "completed"
-ResponseItem = ResponseMessageItem | ResponseFunctionCallItem
+class ResponseReasoningSummaryText(BaseModel, frozen=True):
+ """Summary text part in a reasoning output item."""
+
+ type: Literal["summary_text"] = "summary_text"
+ text: str
+
+
+class ResponseReasoningItem(BaseModel, frozen=True):
+ """Reasoning output item in response output array."""
+
+ type: Literal["reasoning"] = "reasoning"
+ id: str
+ summary: list[ResponseReasoningSummaryText] = Field(default_factory=list)
+ status: ResponseStatus = "completed"
+
+
+ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem
class ResponseUsage(BaseModel, frozen=True):
@@ -273,6 +289,58 @@ class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True):
arguments: str
+class ResponseReasoningSummaryPartAddedEvent(BaseModel, frozen=True):
+ """Event sent when a reasoning summary part is added."""
+
+ type: Literal["response.reasoning_summary_part.added"] = (
+ "response.reasoning_summary_part.added"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ summary_index: int
+ part: ResponseReasoningSummaryText
+
+
+class ResponseReasoningSummaryTextDeltaEvent(BaseModel, frozen=True):
+ """Event sent for reasoning summary text delta during streaming."""
+
+ type: Literal["response.reasoning_summary_text.delta"] = (
+ "response.reasoning_summary_text.delta"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ summary_index: int
+ delta: str
+
+
+class ResponseReasoningSummaryTextDoneEvent(BaseModel, frozen=True):
+ """Event sent when reasoning summary text is done."""
+
+ type: Literal["response.reasoning_summary_text.done"] = (
+ "response.reasoning_summary_text.done"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ summary_index: int
+ text: str
+
+
+class ResponseReasoningSummaryPartDoneEvent(BaseModel, frozen=True):
+ """Event sent when a reasoning summary part is done."""
+
+ type: Literal["response.reasoning_summary_part.done"] = (
+ "response.reasoning_summary_part.done"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ summary_index: int
+ part: ResponseReasoningSummaryText
+
+
class ResponseCompletedEvent(BaseModel, frozen=True):
"""Event sent when response is completed."""
@@ -292,5 +360,9 @@ ResponsesStreamEvent = (
| ResponseOutputItemDoneEvent
| ResponseFunctionCallArgumentsDeltaEvent
| ResponseFunctionCallArgumentsDoneEvent
+ | ResponseReasoningSummaryPartAddedEvent
+ | ResponseReasoningSummaryTextDeltaEvent
+ | ResponseReasoningSummaryTextDoneEvent
+ | ResponseReasoningSummaryPartDoneEvent
| ResponseCompletedEvent
)
diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py
index a66f762c..19118e15 100644
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -28,6 +28,7 @@ class GenerationResponse(BaseRunnerResponse):
finish_reason: FinishReason | None = None
stats: GenerationStats | None = None
usage: Usage | None
+ is_thinking: bool = False
class ImageGenerationResponse(BaseRunnerResponse):
diff --git a/src/exo/worker/engines/mlx/dsml_encoding.py b/src/exo/worker/engines/mlx/dsml_encoding.py
new file mode 100644
index 00000000..d5a5e210
--- /dev/null
+++ b/src/exo/worker/engines/mlx/dsml_encoding.py
@@ -0,0 +1,72 @@
+import json
+import re
+from typing import Any
+
+from mlx_lm.chat_templates import deepseek_v32
+
+from exo.shared.types.api import ToolCallItem
+
+BOS_TOKEN: str = deepseek_v32.bos_token
+EOS_TOKEN: str = deepseek_v32.eos_token
+DSML_TOKEN: str = deepseek_v32.dsml_token
+THINKING_START: str = deepseek_v32.thinking_start_token
+THINKING_END: str = deepseek_v32.thinking_end_token
+USER_TOKEN = "<\uff5cUser\uff5c>"
+ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
+TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
+TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
+encode_messages = deepseek_v32.encode_messages
+
+_INVOKE_PATTERN = re.compile(
+ rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"
+ rf"(.*?)"
+ rf"</{re.escape(DSML_TOKEN)}invoke>",
+ re.DOTALL,
+)
+
+_PARAM_PATTERN = re.compile(
+ rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">"
+ rf"(.*?)"
+ rf"</{re.escape(DSML_TOKEN)}parameter>",
+ re.DOTALL,
+)
+
+
+def parse_dsml_output(text: str) -> list[ToolCallItem] | None:
+ """Parse DSML function_calls block from model output text.
+
+ Args:
+ text: The text containing the DSML function_calls block
+ (including the start/end markers).
+
+ Returns:
+ List of ToolCallItem, or None if parsing fails.
+ """
+ tool_calls: list[ToolCallItem] = []
+
+ for invoke_match in _INVOKE_PATTERN.finditer(text):
+ func_name = invoke_match.group(1)
+ invoke_body = invoke_match.group(2)
+
+ args: dict[str, Any] = {}
+ for param_match in _PARAM_PATTERN.finditer(invoke_body):
+ param_name = param_match.group(1)
+ is_string = param_match.group(2) == "true"
+ param_value = param_match.group(3)
+
+ if is_string:
+ args[param_name] = param_value
+ else:
+ try:
+ args[param_name] = json.loads(param_value)
+ except (json.JSONDecodeError, ValueError):
+ args[param_name] = param_value
+
+ tool_calls.append(
+ ToolCallItem(
+ name=func_name,
+ arguments=json.dumps(args),
+ )
+ )
+
+ return tool_calls if tool_calls else None
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 567b0f91..01bfbe50 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -458,6 +458,19 @@ def _patch_lossy_chat_template(template: str) -> str | None:
return patched if n > 0 else None
+def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
+ if "deepseek-v3.2" not in task_params.model.lower():
+ return False
+ # Use DSML encoding when tools are provided or tool results are in the conversation
+ if task_params.tools:
+ return True
+ if task_params.chat_template_messages:
+ return any(
+ msg.get("role") == "tool" for msg in task_params.chat_template_messages
+ )
+ return False
+
+
def apply_chat_template(
tokenizer: TokenizerWrapper,
task_params: TextGenerationTaskParams,
@@ -469,7 +482,6 @@ def apply_chat_template(
When chat_template_messages is available (from Chat Completions API),
uses those directly to preserve tool_calls, thinking, and other fields.
- Otherwise builds messages from the task params input/instructions.
"""
formatted_messages: list[dict[str, Any]] = []
if task_params.chat_template_messages is not None:
@@ -497,6 +509,19 @@ def apply_chat_template(
partial_assistant_content = cast(str, formatted_messages[-1].get("content", ""))
formatted_messages = formatted_messages[:-1]
+ if _needs_dsml_encoding(task_params):
+ from exo.worker.engines.mlx.dsml_encoding import encode_messages
+
+ prompt = encode_messages(
+ messages=formatted_messages,
+ thinking_mode="thinking" if task_params.enable_thinking else "chat",
+ tools=task_params.tools,
+ )
+ if partial_assistant_content:
+ prompt += partial_assistant_content
+ logger.info(prompt)
+ return prompt
+
extra_kwargs: dict[str, Any] = {}
if task_params.enable_thinking is not None:
# Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 51950192..f22ae5c8 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -7,6 +7,7 @@ from functools import cache
from typing import Literal
import mlx.core as mx
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
from mlx_lm.models.gpt_oss import Model as GptOssModel
from mlx_lm.tokenizer_utils import TokenizerWrapper
from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
@@ -353,16 +354,22 @@ def main(
group=group,
)
- # For other thinking models (GLM, etc.), check if we need to
- # prepend the thinking tag that was consumed by the chat template
- if detect_thinking_prompt_suffix(prompt, tokenizer):
+ if tokenizer.has_thinking:
mlx_generator = parse_thinking_models(
- mlx_generator, tokenizer
+ mlx_generator,
+ tokenizer,
+ # For other thinking models (GLM, etc.), check if we need to
+ # prepend the thinking tag that was consumed by the chat template
+ starts_in_thinking=detect_thinking_prompt_suffix(
+ prompt, tokenizer
+ ),
)
- # GPT-OSS specific parsing to match other model formats.
+ # Model-specific output parsing for tool calls.
if isinstance(inference_model, GptOssModel):
mlx_generator = parse_gpt_oss(mlx_generator)
+ elif isinstance(inference_model, DeepseekV32Model):
+ mlx_generator = parse_deepseek_v32(mlx_generator)
elif tool_parser:
mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
@@ -414,6 +421,7 @@ def main(
stats=response.stats,
logprob=response.logprob,
top_logprobs=response.top_logprobs,
+ is_thinking=response.is_thinking,
),
)
)
@@ -675,44 +683,208 @@ def parse_gpt_oss(
if ch == "analysis" and not thinking:
thinking = True
- yield response.model_copy(update={"text": "<think>"})
if ch != "analysis" and thinking:
thinking = False
- yield response.model_copy(update={"text": "</think>"})
if delta:
- yield response.model_copy(update={"text": delta})
+ yield response.model_copy(update={"text": delta, "is_thinking": thinking})
if response.finish_reason is not None:
- if thinking:
- yield response.model_copy(update={"text": "</think>"})
yield response
+def parse_deepseek_v32(
+ responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse | ToolCallResponse]:
+ """Parse DeepSeek V3.2 DSML tool calls from the generation stream.
+
+ Uses accumulated-text matching (not per-token marker checks) because
+ DSML markers like <｜DSML｜function_calls> may span multiple tokens.
+ Also handles <think>...</think> blocks for thinking mode.
+ """
+ from exo.worker.engines.mlx.dsml_encoding import (
+ THINKING_END,
+ THINKING_START,
+ TOOL_CALLS_END,
+ TOOL_CALLS_START,
+ parse_dsml_output,
+ )
+
+ accumulated = ""
+ in_tool_call = False
+ thinking = False
+ # Tokens buffered while we detect the start of a DSML block
+ pending_buffer: list[GenerationResponse] = []
+ # Text accumulated during a tool call block
+ tool_call_text = ""
+
+ for response in responses:
+ assert isinstance(response, GenerationResponse)
+
+ # ── Handle thinking tags ──
+ if not thinking and THINKING_START in response.text:
+ thinking = True
+ # Yield any text before the <think> tag
+ before = response.text[: response.text.index(THINKING_START)]
+ if before:
+ yield response.model_copy(update={"text": before})
+ continue
+
+ if thinking and THINKING_END in response.text:
+ thinking = False
+ # Yield any text after the </think> tag
+ after = response.text[
+ response.text.index(THINKING_END) + len(THINKING_END) :
+ ]
+ if after:
+ yield response.model_copy(update={"text": after, "is_thinking": False})
+ continue
+
+ if thinking:
+ yield response.model_copy(update={"is_thinking": True})
+ continue
+
+ # ── Handle tool call accumulation ──
+ if in_tool_call:
+ tool_call_text += response.text
+ if TOOL_CALLS_END in tool_call_text:
+ # Parse the accumulated DSML block
+ parsed = parse_dsml_output(tool_call_text)
+ if parsed is not None:
+ logger.info(f"parsed DSML tool calls: {parsed}")
+ yield ToolCallResponse(
+ tool_calls=parsed,
+ usage=response.usage,
+ stats=response.stats,
+ )
+ else:
+ logger.warning(
+ f"DSML tool call parsing failed for: {tool_call_text}"
+ )
+ yield response.model_copy(update={"text": tool_call_text})
+ in_tool_call = False
+ tool_call_text = ""
+ continue
+
+ # EOS reached before end marker — yield buffered text as-is
+ if response.finish_reason is not None:
+ logger.info("DSML tool call parsing interrupted by EOS")
+ yield response.model_copy(update={"text": tool_call_text})
+ in_tool_call = False
+ tool_call_text = ""
+ continue
+
+ # ── Detect start of tool call block ──
+ accumulated += response.text
+
+ if TOOL_CALLS_START in accumulated:
+ # The start marker might be split across pending_buffer + current token
+ start_idx = accumulated.index(TOOL_CALLS_START)
+ # Yield any pending tokens that are purely before the marker
+ pre_text = accumulated[:start_idx]
+ if pre_text:
+ # Flush pending buffer tokens that contributed text before the marker
+ for buf_resp in pending_buffer:
+ if pre_text:
+ chunk = buf_resp.text
+ if len(chunk) <= len(pre_text):
+ yield buf_resp
+ pre_text = pre_text[len(chunk) :]
+ else:
+ yield buf_resp.model_copy(update={"text": pre_text})
+ pre_text = ""
+ pending_buffer = []
+ tool_call_text = accumulated[start_idx:]
+ accumulated = ""
+
+ # Check if the end marker is already present (entire tool call in one token)
+ if TOOL_CALLS_END in tool_call_text:
+ parsed = parse_dsml_output(tool_call_text)
+ if parsed is not None:
+ logger.info(f"parsed DSML tool calls: {parsed}")
+ yield ToolCallResponse(
+ tool_calls=parsed,
+ usage=response.usage,
+ stats=response.stats,
+ )
+ else:
+ logger.warning(
+ f"DSML tool call parsing failed for: {tool_call_text}"
+ )
+ yield response.model_copy(update={"text": tool_call_text})
+ tool_call_text = ""
+ else:
+ in_tool_call = True
+ continue
+
+ # Check if accumulated text might be the start of a DSML marker
+ # Buffer tokens if we see a partial match at the end
+ if _could_be_dsml_prefix(accumulated):
+ pending_buffer.append(response)
+ continue
+
+ # No partial match — flush all pending tokens and the current one
+ for buf_resp in pending_buffer:
+ yield buf_resp
+ pending_buffer = []
+ accumulated = ""
+ yield response
+
+ # Flush any remaining pending buffer at generator end
+ for buf_resp in pending_buffer:
+ yield buf_resp
+
+
+def _could_be_dsml_prefix(text: str) -> bool:
+ """Check if the end of text could be the start of a DSML function_calls marker.
+
+ We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern.
+ This allows us to buffer tokens until we can determine if a tool call is starting.
+ """
+ from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START
+
+ # Only check the last portion of text that could overlap with the marker
+ max_check = len(TOOL_CALLS_START)
+ tail = text[-max_check:] if len(text) > max_check else text
+
+ # Check if any suffix of tail is a prefix of TOOL_CALLS_START
+ for i in range(len(tail)):
+ suffix = tail[i:]
+ if TOOL_CALLS_START.startswith(suffix):
+ return True
+ return False
+
+
def parse_thinking_models(
responses: Generator[GenerationResponse],
tokenizer: TokenizerWrapper,
+ starts_in_thinking: bool = True,
) -> Generator[GenerationResponse]:
+ """Route thinking tokens via is_thinking flag.
+
+ Swallows think tag tokens, sets is_thinking on all others.
+ Always yields tokens with finish_reason to avoid hanging the chunk stream.
"""
- For models that inject thinking tags in the prompt (like GLM-4.7),
- prepend the thinking tag to the output stream so the frontend
- can properly parse thinking content.
- """
- first = True
+ in_thinking = starts_in_thinking
for response in responses:
if isinstance(response, ToolCallResponse):
yield response
continue
- if first:
- first = False
- yield response.model_copy(
- update={
- "text": tokenizer.think_start,
- "token": tokenizer.think_start_id,
- }
- )
- yield response
+
+ is_think_tag = (
+ tokenizer.think_end is not None and response.text == tokenizer.think_end
+ ) or (
+ tokenizer.think_start is not None and response.text == tokenizer.think_start
+ )
+
+ if is_think_tag:
+ in_thinking = response.text != tokenizer.think_end
+ # Never swallow finish_reason — the chunk stream needs it to terminate.
+ if response.finish_reason is not None:
+ yield response.model_copy(update={"text": "", "is_thinking": False})
+ continue
+ yield response.model_copy(update={"is_thinking": in_thinking})
def _send_image_chunk(
diff --git a/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py
new file mode 100644
index 00000000..a59383e5
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py
@@ -0,0 +1,967 @@
+import json
+from collections.abc import Generator
+from typing import Any
+
+from exo.shared.types.worker.runner_response import (
+ GenerationResponse,
+ ToolCallResponse,
+)
+from exo.worker.engines.mlx.dsml_encoding import (
+ ASSISTANT_TOKEN,
+ BOS_TOKEN,
+ DSML_TOKEN,
+ EOS_TOKEN,
+ THINKING_END,
+ THINKING_START,
+ TOOL_CALLS_END,
+ TOOL_CALLS_START,
+ USER_TOKEN,
+ encode_messages,
+ parse_dsml_output,
+)
+from exo.worker.runner.runner import parse_deepseek_v32
+
+# ── Shared fixtures ──────────────────────────────────────────────
+
+_WEATHER_TOOLS: list[dict[str, Any]] = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get the current weather in a given city",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "city": {"type": "string", "description": "The city name"},
+ "units": {
+ "type": "string",
+ "enum": ["celsius", "fahrenheit"],
+ "description": "Temperature units",
+ },
+ },
+ "required": ["city"],
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "description": "Get the current time in a timezone",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "timezone": {"type": "string"},
+ },
+ "required": ["timezone"],
+ },
+ },
+ },
+]
+
+
+def _simulate_tokens(
+ texts: list[str],
+ finish_on_last: bool = True,
+) -> Generator[GenerationResponse]:
+ """Simulate a model producing tokens from a list of text strings."""
+ for i, text in enumerate(texts):
+ is_last = i == len(texts) - 1
+ yield GenerationResponse(
+ text=text,
+ token=i,
+ finish_reason="stop" if (is_last and finish_on_last) else None,
+ usage=None,
+ )
+
+
+# ── Test: Standard text response (no tool calls) ────────────────
+
+
+class TestE2EStandardResponse:
+ """Model generates a plain text response — no tool calling involved."""
+
+ def test_plain_text_passthrough(self):
+ """Simulate model producing: 'The weather in NYC is 72°F and sunny.'"""
+ # Step 1: Encode the prompt (with tools available)
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "What's the weather in NYC?"},
+ ]
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+ # Verify prompt structure
+ assert BOS_TOKEN in prompt
+ assert "## Tools" in prompt
+ assert "get_weather" in prompt
+ assert f"{USER_TOKEN}What's the weather in NYC?{ASSISTANT_TOKEN}" in prompt
+
+ # Step 2: Simulate model response — plain text tokens (no DSML)
+ model_tokens = [
+ "The weather",
+ " in NYC",
+ " is 72",
+ "°F",
+ " and sunny",
+ ".",
+ ]
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+ # Step 3: Verify all tokens pass through as GenerationResponse
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 0
+ assert len(gen_results) == 6
+ full_text = "".join(r.text for r in gen_results)
+ assert full_text == "The weather in NYC is 72°F and sunny."
+ assert gen_results[-1].finish_reason == "stop"
+
+
+# ── Test: Tool call response ─────────────────────────────────────
+
+
+class TestE2EToolCallResponse:
+ """Model generates a DSML tool call — realistic token boundaries."""
+
+ def test_realistic_tool_call_tokens(self):
+ """Simulate model generating a get_weather tool call with realistic token splits.
+
+ Real models split DSML markers across tokens unpredictably.
+ This simulates how DeepSeek V3.2 actually tokenizes DSML output.
+ """
+ # Step 1: Encode prompt
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "What's the weather in San Francisco?"},
+ ]
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+ assert "get_weather" in prompt
+
+ # Step 2: Simulate realistic token-by-token model output
+ # The model first produces some text, then a DSML tool call block
+ model_tokens = [
+ "I'll check the weather for you.",
+ "\n\n",
+ f"<{DSML_TOKEN}", # marker split across tokens
+ "function_calls>\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">',
+ "San Francisco",
+ f"</{DSML_TOKEN}parameter>\n",
+ f'<{DSML_TOKEN}parameter name="units" string="false">',
+ '"celsius"',
+ f"</{DSML_TOKEN}parameter>\n",
+ f"</{DSML_TOKEN}invoke>\n",
+ f"</{DSML_TOKEN}function_calls>",
+ ]
+
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+ # Step 3: Verify
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ # Should have text tokens before tool call + one ToolCallResponse
+ assert len(tool_results) == 1
+ assert len(tool_results[0].tool_calls) == 1
+
+ tc = tool_results[0].tool_calls[0]
+ assert tc.name == "get_weather"
+ args = json.loads(tc.arguments) # pyright: ignore[reportAny]
+ assert args["city"] == "San Francisco"
+ assert args["units"] == "celsius"
+
+ # The text before the tool call should still be yielded
+ text_before = "".join(r.text for r in gen_results if not r.is_thinking)
+ assert "check the weather" in text_before
+
+ def test_multiple_tool_calls_in_one_block(self):
+ """Model generates two tool calls in a single function_calls block."""
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are helpful."},
+ {"role": "user", "content": "Weather in NYC and time in EST?"},
+ ]
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+ assert "get_weather" in prompt
+ assert "get_time" in prompt
+
+ # Simulate model output with two invocations
+ model_tokens = [
+ "Let me check both.\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ f'<{DSML_TOKEN}invoke name="get_time">\n',
+ f'<{DSML_TOKEN}parameter name="timezone" string="true">EST</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 1
+ assert len(tool_results[0].tool_calls) == 2
+ assert tool_results[0].tool_calls[0].name == "get_weather"
+ assert tool_results[0].tool_calls[1].name == "get_time"
+
+ args0 = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny]
+ args1 = json.loads(tool_results[0].tool_calls[1].arguments) # pyright: ignore[reportAny]
+ assert args0 == {"city": "NYC"}
+ assert args1 == {"timezone": "EST"}
+
+
+# ── Test: Multi-turn tool use flow ───────────────────────────────
+
+
+class TestE2EMultiTurnToolUse:
+ """Full multi-turn: user asks → model calls tool → tool result → model answers."""
+
+ def test_encode_multi_turn_with_tool_results(self):
+ """Verify the prompt for turn 2 (after tool results) is correctly encoded."""
+ # Turn 1: user asks, model calls tool
+ # Turn 2: tool result provided, model answers
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are a weather assistant."},
+ {"role": "user", "content": "What's the weather in NYC?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "NYC"}',
+ },
+ }
+ ],
+ },
+ {"role": "tool", "content": '{"temperature": 72, "condition": "sunny"}'},
+ ]
+
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+ # Verify multi-turn structure
+ assert BOS_TOKEN in prompt
+ assert "You are a weather assistant." in prompt
+ assert "## Tools" in prompt
+
+ # The assistant's tool call should be encoded as DSML
+ assert TOOL_CALLS_START in prompt
+ assert f'<{DSML_TOKEN}invoke name="get_weather">' in prompt
+ assert EOS_TOKEN in prompt
+
+ # The tool result should be wrapped in function_results
+ assert "<function_results>" in prompt
+ assert "<result>" in prompt
+ assert "72" in prompt
+ assert "</function_results>" in prompt
+
+ # Now simulate model answering after seeing the tool result
+ model_tokens = [
+ "The current",
+ " weather in NYC",
+ " is 72°F",
+ " and sunny.",
+ ]
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 0
+ full_text = "".join(r.text for r in gen_results)
+ assert full_text == "The current weather in NYC is 72°F and sunny."
+
+ def test_multi_tool_results_encoding(self):
+ """Verify encoding when model called two tools and both return results."""
+ messages: list[dict[str, Any]] = [
+ {"role": "user", "content": "Weather and time?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "LA"}',
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "arguments": '{"timezone": "PST"}',
+ },
+ },
+ ],
+ },
+ {"role": "tool", "content": "85F, clear skies"},
+ {"role": "tool", "content": "3:42 PM PST"},
+ ]
+
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+ # Should have one function_results block with two results
+ assert prompt.count("<function_results>") == 1
+ assert prompt.count("</function_results>") == 1
+ assert "<result>85F, clear skies</result>" in prompt
+ assert "<result>3:42 PM PST</result>" in prompt
+
+
+# ── Test: Thinking + tool call ───────────────────────────────────
+
+
+class TestE2EThinkingAndToolCall:
+ """Model uses thinking mode, reasons, then makes a tool call."""
+
+ def test_thinking_then_tool_call(self):
+ """Model thinks first, then produces a DSML tool call block."""
+ messages: list[dict[str, Any]] = [
+ {"role": "user", "content": "What's the weather?"},
+ ]
+ prompt = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ # Thinking mode: prompt should end with <think>
+ assert prompt.endswith(THINKING_START)
+
+ # Simulate: model outputs <think>, thinks, closes thinking, then tool call.
+ # In the full pipeline, parse_thinking_models handles the case where
+ # <think> is in the prompt. Here we test parse_deepseek_v32 directly,
+ # which detects <think>/<think> markers in the stream.
+ model_tokens = [
+ THINKING_START,
+ "The user wants weather",
+ " information. I should use",
+ " the get_weather tool.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">',
+ "San Francisco",
+ f"</{DSML_TOKEN}parameter>\n",
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ # Should have thinking tokens + tool call
+ thinking_results = [r for r in gen_results if r.is_thinking]
+
+ assert len(thinking_results) >= 1
+ thinking_text = "".join(r.text for r in thinking_results)
+ assert "get_weather tool" in thinking_text
+
+ assert len(tool_results) == 1
+ assert tool_results[0].tool_calls[0].name == "get_weather"
+ args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny]
+ assert args["city"] == "San Francisco"
+
+ def test_thinking_prompt_encoding(self):
+ """Verify thinking mode affects prompt encoding correctly."""
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "Be thorough."},
+ {"role": "user", "content": "What's the weather?"},
+ ]
+
+ # With thinking enabled
+ prompt_think = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ assert prompt_think.endswith(THINKING_START)
+
+ # With thinking disabled
+ prompt_no_think = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="chat"
+ )
+ assert prompt_no_think.endswith(THINKING_END)
+
+ # Both should have the same tool definitions
+ assert "get_weather" in prompt_think
+ assert "get_weather" in prompt_no_think
+
+
+# ── Test: Round-trip encode → parse ──────────────────────────────
+
+
+class TestE2ERoundTrip:
+ """Verify that DSML we encode can be parsed back correctly."""
+
+ def test_encoded_tool_call_is_parseable(self):
+ """Encode an assistant tool call message, then parse the DSML output."""
+ messages: list[dict[str, Any]] = [
+ {"role": "user", "content": "Weather?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Tokyo", "units": "celsius"}',
+ },
+ }
+ ],
+ },
+ ]
+
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+ # Extract the DSML function_calls block from the prompt
+ start = prompt.index(TOOL_CALLS_START)
+ end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END)
+ dsml_block = prompt[start:end]
+
+ # Parse it back
+ parsed = parse_dsml_output(dsml_block)
+ assert parsed is not None
+ assert len(parsed) == 1
+ assert parsed[0].name == "get_weather"
+ args = json.loads(parsed[0].arguments) # pyright: ignore[reportAny]
+ assert args["city"] == "Tokyo"
+ assert args["units"] == "celsius"
+
+ def test_encoded_multi_tool_call_round_trips(self):
+ """Encode multiple tool calls, verify they parse back correctly."""
+ messages: list[dict[str, Any]] = [
+ {"role": "user", "content": "Both please"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Paris"}',
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "arguments": '{"timezone": "CET"}',
+ },
+ },
+ ],
+ },
+ ]
+
+ prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS)
+
+ start = prompt.index(TOOL_CALLS_START)
+ end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END)
+ dsml_block = prompt[start:end]
+
+ parsed = parse_dsml_output(dsml_block)
+ assert parsed is not None
+ assert len(parsed) == 2
+ assert parsed[0].name == "get_weather"
+ assert parsed[1].name == "get_time"
+ assert json.loads(parsed[0].arguments) == {"city": "Paris"}
+ assert json.loads(parsed[1].arguments) == {"timezone": "CET"}
+
+
+# ── Test: Edge cases with realistic token boundaries ─────────────
+
+
+class TestE2EEdgeCases:
+ """Edge cases that occur in real model inference."""
+
+ def test_dsml_marker_split_at_fullwidth_pipe(self):
+ """The fullwidth pipe character ｜ might be its own token."""
+ # This is a realistic tokenization: the DSML marker is split at the ｜ chars
+ model_tokens = [
+ "Let me help.\n\n",
+ "<\uff5c", # start of ｜DSML｜
+ "DSML\uff5c", # rest of DSML token
+ "function_calls>\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 1
+ assert tool_results[0].tool_calls[0].name == "get_weather"
+
+ def test_tool_call_with_nested_json_object(self):
+ """Model passes a complex JSON object as a non-string parameter."""
+ dsml_block = (
+ f"{TOOL_CALLS_START}\n"
+ f'<{DSML_TOKEN}invoke name="create_event">\n'
+ f'<{DSML_TOKEN}parameter name="title" string="true">Team Standup</{DSML_TOKEN}parameter>\n'
+ f'<{DSML_TOKEN}parameter name="config" string="false">'
+ f'{{"recurring": true, "days": ["mon", "wed", "fri"], "time": "09:00"}}'
+ f"</{DSML_TOKEN}parameter>\n"
+ f"</{DSML_TOKEN}invoke>\n"
+ f"{TOOL_CALLS_END}"
+ )
+
+ # Feed as single token (model might produce it all at once after prefill)
+ results = list(parse_deepseek_v32(_simulate_tokens([dsml_block])))
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 1
+ tc = tool_results[0].tool_calls[0]
+ assert tc.name == "create_event"
+ args = json.loads(tc.arguments) # pyright: ignore[reportAny]
+ assert args["title"] == "Team Standup"
+ assert args["config"]["recurring"] is True
+ assert args["config"]["days"] == ["mon", "wed", "fri"]
+
+ def test_text_with_angle_brackets_not_mistaken_for_dsml(self):
+ """Angle brackets in normal text should not trigger DSML buffering."""
+ model_tokens = [
+ "The formula is ",
+ "<x, y>",
+ " where x > 0",
+ " and y < 100.",
+ ]
+
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 0
+ full_text = "".join(r.text for r in gen_results)
+ assert "formula" in full_text
+ assert "<x, y>" in full_text
+
+ def test_empty_model_response(self):
+ """Model produces only EOS (empty response)."""
+ model_tokens = [""]
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+ gen_results = [r for r in results if isinstance(r, GenerationResponse)]
+ assert len(gen_results) == 1
+ assert gen_results[0].text == ""
+ assert gen_results[0].finish_reason == "stop"
+
+
+# ── Test: Full EPDP spec round-trip ──────────────────────────────
+
+
+class TestE2EFullRoundTrip:
+ """Full round-trip matching the vLLM EPDP spec.
+
+ Simulates the complete multi-turn flow:
+ Turn 1: user asks → think → tool call → tool result → think → answer
+ Turn 2: user asks again → old reasoning stripped → think → answer
+ """
+
+ def test_single_tool_full_flow_with_thinking(self):
+ """Complete flow: user → think → tool call → tool result → think → answer.
+
+ This is the core EPDP flow from the vLLM spec.
+ """
+ # ── Turn 1.1: User asks, encode prompt ──
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are a weather assistant."},
+ {"role": "user", "content": "How's the weather in Hangzhou?"},
+ ]
+ prompt_1 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ assert prompt_1.endswith(THINKING_START)
+ assert "## Tools" in prompt_1
+ assert "get_weather" in prompt_1
+
+ # ── Turn 1.1: Model thinks, then calls tool ──
+ model_tokens_1 = [
+ THINKING_START,
+ "The user wants to know the weather in Hangzhou.",
+ " I need to use the get_weather tool.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+ results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+
+ # Verify: thinking tokens + tool call
+ gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)]
+ tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+ thinking_1 = [r for r in gen_1 if r.is_thinking]
+
+ assert len(thinking_1) >= 1
+ assert "get_weather tool" in "".join(r.text for r in thinking_1)
+ assert len(tool_1) == 1
+ assert tool_1[0].tool_calls[0].name == "get_weather"
+ tc_args = json.loads(tool_1[0].tool_calls[0].arguments) # pyright: ignore[reportAny]
+ assert tc_args == {"city": "Hangzhou"}
+
+ # ── Turn 1.2: Add assistant response + tool result to messages ──
+ messages.append(
+ {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "The user wants to know the weather in Hangzhou. I need to use the get_weather tool.",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Hangzhou"}',
+ },
+ }
+ ],
+ }
+ )
+ messages.append(
+ {
+ "role": "tool",
+ "content": '{"temperature": "7~13°C", "condition": "Cloudy"}',
+ }
+ )
+
+ # Encode prompt for turn 1.2
+ prompt_2 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+
+ # Verify: prompt has the full conversation structure
+ assert TOOL_CALLS_START in prompt_2 # assistant's encoded tool call
+ assert EOS_TOKEN in prompt_2 # assistant turn ends with EOS
+ assert "<function_results>" in prompt_2
+ assert "<result>" in prompt_2
+ assert "Cloudy" in prompt_2
+ assert "</function_results>" in prompt_2
+ # After tool results with thinking enabled → <think> appended
+ assert prompt_2.endswith(THINKING_START)
+ # The assistant's reasoning_content should appear (it's after last_user_idx)
+ assert "get_weather tool" in prompt_2
+
+ # ── Turn 1.2: Model thinks about results, then answers ──
+ model_tokens_2 = [
+ THINKING_START,
+ "The weather in Hangzhou is Cloudy, 7~13°C.",
+ " I'll tell the user.",
+ THINKING_END,
+ "The weather in Hangzhou is currently cloudy with temperatures between 7°C and 13°C.",
+ ]
+ results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+
+ gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)]
+ tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+ thinking_2 = [r for r in gen_2 if r.is_thinking]
+ non_thinking_2 = [r for r in gen_2 if not r.is_thinking]
+
+ assert len(tool_2) == 0 # No more tool calls
+ assert len(thinking_2) >= 1
+ assert "Cloudy" in "".join(r.text for r in thinking_2)
+ assert len(non_thinking_2) >= 1
+ final_text = "".join(r.text for r in non_thinking_2)
+ assert "7°C" in final_text
+ assert "13°C" in final_text
+
+ def test_multi_tool_full_flow(self):
+ """Flow with two tools: user → think → 2 tool calls → 2 results → think → answer."""
+ # ── Initial prompt ──
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You help with weather and time."},
+ {"role": "user", "content": "Weather in NYC and time in EST?"},
+ ]
+ prompt_1 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ assert prompt_1.endswith(THINKING_START)
+
+ # ── Model thinks, calls both tools ──
+ model_tokens_1 = [
+ THINKING_START,
+ "Two requests: weather and time. I'll call both.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ f'<{DSML_TOKEN}invoke name="get_time">\n',
+ f'<{DSML_TOKEN}parameter name="timezone" string="true">EST</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+ results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+ tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_1) == 1
+ assert len(tool_1[0].tool_calls) == 2
+ assert tool_1[0].tool_calls[0].name == "get_weather"
+ assert tool_1[0].tool_calls[1].name == "get_time"
+
+ # ── Add assistant + both tool results ──
+ messages.append(
+ {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "Two requests: weather and time. I'll call both.",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "NYC"}',
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "arguments": '{"timezone": "EST"}',
+ },
+ },
+ ],
+ }
+ )
+ messages.append({"role": "tool", "content": "72°F, sunny"})
+ messages.append({"role": "tool", "content": "2:30 PM EST"})
+
+ prompt_2 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+
+ # Verify multi-tool result encoding
+ # Count is 2: 1 in _TOOLS_SYSTEM_TEMPLATE example + 1 in conversation
+ assert prompt_2.count("<function_results>") == 2
+ assert prompt_2.count("</function_results>") == 2
+ assert "<result>72°F, sunny</result>" in prompt_2
+ assert "<result>2:30 PM EST</result>" in prompt_2
+ assert prompt_2.endswith(THINKING_START)
+
+ # ── Model thinks about results, answers ──
+ model_tokens_2 = [
+ THINKING_START,
+ "Got both results. Weather is 72°F sunny, time is 2:30 PM.",
+ THINKING_END,
+ "In NYC it's currently 72°F and sunny. The time in EST is 2:30 PM.",
+ ]
+ results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+
+ tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+ gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)]
+ non_thinking_2 = [r for r in gen_2 if not r.is_thinking]
+
+ assert len(tool_2) == 0
+ final_text = "".join(r.text for r in non_thinking_2)
+ assert "72°F" in final_text
+ assert "2:30 PM" in final_text
+
+ def test_two_user_turns_reasoning_stripped(self):
+ """Turn 2: old reasoning_content is stripped from history.
+
+ Per the vLLM spec, clear_reasoning_content is called between user turns
+ to save bandwidth. Our _drop_old_thinking handles this.
+ """
+ # Full turn 1 conversation (already completed)
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are helpful."},
+ {"role": "user", "content": "Weather in Hangzhou?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "I need to call get_weather for Hangzhou.",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Hangzhou"}',
+ },
+ }
+ ],
+ },
+ {"role": "tool", "content": "Cloudy 7~13°C"},
+ {
+ "role": "assistant",
+ "content": "The weather in Hangzhou is cloudy, 7-13°C.",
+ "reasoning_content": "The tool returned cloudy weather. I'll summarize.",
+ },
+ # Turn 2: user asks again
+ {"role": "user", "content": "What about Beijing?"},
+ ]
+
+ prompt = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+
+ # Old reasoning_content from turn 1 assistants should be STRIPPED
+ # (they're before the last user message at index 5)
+ assert "I need to call get_weather" not in prompt
+ assert "tool returned cloudy" not in prompt
+
+ # But the assistant's content and tool calls should still be there
+ assert "cloudy, 7-13°C" in prompt
+ assert TOOL_CALLS_START in prompt
+
+ # Prompt ends with <think> for the new turn
+ assert prompt.endswith(THINKING_START)
+
+ # ── Turn 2: Model thinks, calls tool for Beijing ──
+ model_tokens = [
+ THINKING_START,
+ "Now the user wants Beijing weather.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">Beijing</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+ results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+ tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+
+ assert len(tool_results) == 1
+ assert tool_results[0].tool_calls[0].name == "get_weather"
+ args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny]
+ assert args == {"city": "Beijing"}
+
+ def test_chained_tool_calls_loop(self):
+ """Model calls tool, gets result, calls another tool, gets result, answers.
+
+ This simulates the inner while loop from the vLLM spec where the model
+ may need multiple sub-turns of tool calling before it has enough info.
+ """
+ # ── Sub-turn 1: user asks, model calls get_time ──
+ messages: list[dict[str, Any]] = [
+ {"role": "system", "content": "You are helpful."},
+ {"role": "user", "content": "What's the weather in Hangzhou tomorrow?"},
+ ]
+
+ prompt_1 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ assert prompt_1.endswith(THINKING_START)
+
+ # Model first calls get_time to figure out the date
+ model_tokens_1 = [
+ THINKING_START,
+ "I need the current date first to calculate tomorrow.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_time">\n',
+ f'<{DSML_TOKEN}parameter name="timezone" string="true">Asia/Shanghai</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+ results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+ tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)]
+ assert len(tool_1) == 1
+ assert tool_1[0].tool_calls[0].name == "get_time"
+
+ # ── Sub-turn 2: add tool result, model calls get_weather ──
+ messages.append(
+ {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "I need the current date first to calculate tomorrow.",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_time",
+ "arguments": '{"timezone": "Asia/Shanghai"}',
+ },
+ }
+ ],
+ }
+ )
+ messages.append({"role": "tool", "content": "2025-12-01 14:30 CST"})
+
+ prompt_2 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ assert "<result>2025-12-01 14:30 CST</result>" in prompt_2
+ assert prompt_2.endswith(THINKING_START)
+
+ # Model now knows the date, calls get_weather
+ model_tokens_2 = [
+ THINKING_START,
+ "Today is 2025-12-01, so tomorrow is 2025-12-02.",
+ " Now I can check weather for Hangzhou.",
+ THINKING_END,
+ "\n\n",
+ TOOL_CALLS_START,
+ "\n",
+ f'<{DSML_TOKEN}invoke name="get_weather">\n',
+ f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou</{DSML_TOKEN}parameter>\n',
+ f"</{DSML_TOKEN}invoke>\n",
+ TOOL_CALLS_END,
+ ]
+ results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+ tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
+ assert len(tool_2) == 1
+ assert tool_2[0].tool_calls[0].name == "get_weather"
+
+ # ── Sub-turn 3: add weather result, model answers ──
+ messages.append(
+ {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "Today is 2025-12-01, so tomorrow is 2025-12-02. Now I can check weather for Hangzhou.",
+ "tool_calls": [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "Hangzhou"}',
+ },
+ }
+ ],
+ }
+ )
+ messages.append({"role": "tool", "content": "Sunny, 5~12°C"})
+
+ prompt_3 = encode_messages(
+ messages, tools=_WEATHER_TOOLS, thinking_mode="thinking"
+ )
+ # Should have both function_results blocks (one per tool round)
+ # Count is 3: 1 in _TOOLS_SYSTEM_TEMPLATE example + 2 in conversation
+ assert prompt_3.count("<function_results>") == 3
+ assert prompt_3.count("</function_results>") == 3
+ assert "<result>2025-12-01 14:30 CST</result>" in prompt_3
+ assert "<result>Sunny, 5~12°C</result>" in prompt_3
+ assert prompt_3.endswith(THINKING_START)
+
+ # Model finally answers
+ model_tokens_3 = [
+ THINKING_START,
+ "I have the weather for tomorrow in Hangzhou.",
+ THINKING_END,
+ "Tomorrow in Hangzhou will be sunny with temperatures between 5°C and 12°C.",
+ ]
+ results_3 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_3)))
+
+ tool_3 = [r for r in results_3 if isinstance(r, ToolCallResponse)]
+ gen_3 = [r for r in results_3 if isinstance(r, GenerationResponse)]
+ non_thinking_3 = [r for r in gen_3 if not r.is_thinking]
+
+ assert len(tool_3) == 0 # No more tool calls — loop ends
+ final_text = "".join(r.text for r in non_thinking_3)
+ assert "sunny" in final_text.lower()
+ assert "5°C" in final_text
+ assert "12°C" in final_text
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index 38a0a921..864e424b 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -148,6 +148,7 @@ class MockTokenizer:
tool_call_start = None
tool_call_end = None
has_tool_calling = False
+ has_thinking = False
class MockGroup:
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
index 0a7ba102..080f0389 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
@@ -149,12 +149,23 @@ class TestParseGptOssThinkingThenToolCall:
def test_thinking_then_tool_call(self):
results = _collect(THINKING_THEN_TOOL_TOKENS)
- # Should have thinking tags + content + tool call
- text_parts = [r.text for r in results if isinstance(r, GenerationResponse)]
- combined = "".join(text_parts)
- assert "<think>" in combined
- assert "</think>" in combined
- assert "Let me think about this." in combined
+ # Thinking tokens should have is_thinking=True and no <think> tags
+ thinking_responses = [
+ r for r in results if isinstance(r, GenerationResponse) and r.is_thinking
+ ]
+ thinking_text = "".join(r.text for r in thinking_responses)
+ assert "Let me think about this." in thinking_text
+ assert "<think>" not in thinking_text
+ assert "</think>" not in thinking_text
+
+ # Non-thinking tokens should have is_thinking=False
+ non_thinking = [
+ r
+ for r in results
+ if isinstance(r, GenerationResponse) and not r.is_thinking
+ ]
+ non_thinking_text = "".join(r.text for r in non_thinking)
+ assert "<think>" not in non_thinking_text
# And the tool call
tc = _get_tool_call(results)
← 94b2ce69 feat: Mac Studio en2 RDMA port warning v2 (#1551)
·
back to Exo
·
Prevent common fatal crashes (#1555) 7031901a →