fix(server): accept tool_choice and parallel_tool_calls on the MLX chat path

[krisztian-gajdar]

Problem

On Apple Silicon the /v1/chat/completions route proxies to an mlx_lm.server child. It accepted tools but refused tool_choice and parallel_tool_calls with 400 unsupported_field. OpenCode, like most agent clients, sends tool_choice: "auto" on every turn that carries tools, so it failed at its first turn on a Mac. The SGLang route accepts the field, so the same client works on NVIDIA.

What mlx_lm does with these fields

mlx_lm.server reads neither field. Checked in mlx_lm/server.py at v0.30.7 (the sglang bundle's floor, mlx-lm>=0.30.7), v0.31.3 and v0.32.0 (the latest release, and the version the bundle resolves today). It reads tools from the body and renders every declared tool into the prompt through the chat template. It ignores unknown body keys, so it has no way to force a call or to cap the number of calls in a turn.

It also detects tool calls from the model's tokenizer, not from the request: the text state machine switches into its tool state on the tokenizer's tool-call markers, and the collected calls are parsed and returned whether or not the request sent tools. So withholding the tools from the prompt makes a call unlikely, not impossible, and a turn can carry any number of calls.

Change

The MLX route accepts both fields, resolves them itself, proxies neither to the child, and checks what the child sends back:

Request	MLX route
tool_choice omitted or "auto"	tools forwarded unchanged, every returned call passed through
parallel_tool_calls: true	accepted, every returned call passed through
parallel_tool_calls: false	accepted; a choice that carries a second call fails closed
tool_choice: "none", or no tools	tools withheld from the child; a response that still carries a call fails closed
tool_choice: "required" or a named function	400 unsupported_field, param: tool_choice, before any model load

"Fails closed" is 502 upstream_error on a buffered response. On a stream, the calls already within the limit are delivered, then an error event and [DONE], and the child stream is closed. This matches the remote chat route, which raises when calls arrive under tool_choice: "none" or when more than one arrives under parallel_tool_calls: false. The local queued route enforces the same two rules: it withholds tools under none, and its tool-call parser stops a stream at a second call.

The none and no-tools cases are one rule, because the child receives the same body for both. The tool_choice refusal names the device rather than the MLX backend, because the same validation runs on every non-CUDA device. A malformed tool_choice (for example required with no tools, or a function name that is not declared) still answers invalid_request, as before.

The CUDA (SGLang) and remote chat routes are unchanged.

Tests

packages/sie_server/tests/api/test_openai_local.py drives the real route through TestClient, so _validate_mlx_chat_body, _prepare_mlx_body and both response paths run unmodified. Only the mlx_lm.server child is faked, as an httpx.MockTransport. Its streams use mlx_lm's shape: one delta per completed call, carrying index, then the finish event.

• auto, auto + parallel_tool_calls: true, and parallel_tool_calls: true alone: the child receives the tools and neither field, and both returned calls reach the client as a valid ChatCompletion. A streaming auto turn delivers both calls as valid ChatCompletionChunks.
• parallel_tool_calls: false with one returned call, buffered and streamed: the child receives the tools without the field, and the call reaches the client.
• none: the child receives neither tools nor tool_choice.
• Calls the request ruled out, buffered and streamed, each after an earlier assistant tool call and its tool result: two calls under parallel_tool_calls: false, one call under tool_choice: "none", one call with no tools. A buffered response answers the exact 502 envelope. A stream delivers only the calls within the limit, then the error event, and closes the child stream.
• required and a named function, on mps and on cpu: the exact 400 envelope naming the field and the device, the child is never reached and no model is touched.

The 12 new cases that touch the changed behaviour fail against the previous head and pass now. The module passes 133. mise run test -- packages/sie_server/tests gives 11380 passed. Its 5 failures are in host-filesystem tests (config/test_serving_artifacts.py, core/test_model_loader_cache.py) that this change does not touch. mise run lint and mise run typecheck pass. Merged locally with #728, which also edits this route, the module passes 146.

Live run on Apple Silicon

mise run serve -- -b sglang on an M3 Max, model Qwen/Qwen3.5-4B (MLX repo mlx-community/Qwen3.5-4B-4bit), child mlx_lm 0.32.0, through sie_sdk.SIEClient, temperature: 0:

auto: finish_reason=tool_calls tool_calls=['get_weather({"city": "Paris"})'] content=''
auto, two cities: finish_reason=tool_calls tool_calls=['get_weather({"city": "Paris"})', 'get_weather({"city": "Tokyo"})']
stream, auto, two cities: finish_reason=tool_calls tool_calls=['get_weather({"city": "Paris"})', 'get_weather({"city": "Tokyo"})'] error=None
tool result turn, auto: finish_reason=stop tool_calls=[] content='The weather in Paris is currently 18°C with clear skies.'
auto, parallel_tool_calls=false, one city: finish_reason=tool_calls tool_calls=['get_weather({"city": "Paris"})'] content=''
auto, parallel_tool_calls=false, two cities: ServerError status=502 generation child returned more than one tool call with parallel_tool_calls=false
stream, auto, parallel_tool_calls=false, two cities: tool_calls=['get_weather({"city": "Paris"})'] error=ServerError generation child returned more than one tool call with parallel_tool_calls=false
none after a tool turn: ServerError status=502 generation child returned a tool call the request did not offer
stream, none after a tool turn: tool_calls=[] error=ServerError generation child returned a tool call the request did not offer
required: RequestError status=400 'tool_choice' must be 'auto' or 'none' on this device, which cannot force a tool call

"None after a tool turn" sends an earlier assistant get_weather call and its result, then asks for another city, with tools declared and tool_choice: "none". The child got no tools, yet the model emitted a get_weather call and mlx_lm returned it. Before this change the route passed that call to the client.

The agent-loop conformance check from the coding-agent example (#727), run against the same server and model: tool_choice_auto, streaming_tool_call_deltas, parallel_tool_calls, tool_result_round_trip and the finish-reason and usage cases pass (8 pass). tool_choice_required is recorded as refused, which that check treats as an honest answer. The one failure is include_reasoning, a field this PR does not touch.

Summary by CodeRabbit

• New Features
  • MLX chat requests support automatic tool selection and multiple tool calls, including in streamed responses.
  • Requests that disable tool selection no longer send tools to the model.
• Bug Fixes
  • Unsupported tool-selection options are rejected before model loading.
  • Responses are rejected if they include tool calls that were not allowed by the request, or multiple calls when parallel calls are disabled. These checks apply to streamed and non-streamed responses.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Team
• Run ID: e3f3beb0-7779-4406-b9dc-dbd8f7b45dc7

📥 Commits

Reviewing files that changed from the base of the PR and between 5f032fe and 35ce55b.

📒 Files selected for processing (2)

• packages/sie_server/src/sie_server/api/openai_local.py
• packages/sie_server/tests/api/test_openai_local.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 2 remain after this review.

📝 Walkthrough

The MLX chat route accepts supported tool-call controls, rejects unsupported choices, and removes controls that the MLX child does not consume. It checks returned tool calls against request limits in streamed and non-streamed responses.

Changes: MLX tool-call request and response handling

Layer / File(s)	Summary
Validate and prepare tool-call requests packages/sie_server/src/sie_server/api/openai_local.py, packages/sie_server/tests/api/test_openai_local.py	The route accepts supported tool_choice and parallel_tool_calls values. It removes parallel_tool_calls before forwarding and withholds tools when tool_choice is "none". Tests cover supported requests and rejection of unsupported choices before model loading.
Enforce tool-call limits in responses packages/sie_server/src/sie_server/api/openai_local.py, packages/sie_server/tests/api/test_openai_local.py	The route rejects returned tool calls when the request offered no tools and rejects multiple distinct calls per choice when parallel_tool_calls is false. It checks streamed and non-streamed responses. Tests cover permitted calls and upstream errors for prohibited calls.

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant MLXChatRoute
  participant MLXChild
  Client->>MLXChatRoute: Send chat request with tool controls
  MLXChatRoute->>MLXChatRoute: Validate controls and set tool-call limit
  MLXChatRoute->>MLXChild: Forward prepared request
  MLXChild-->>MLXChatRoute: Return streamed or non-streamed response
  MLXChatRoute->>MLXChatRoute: Check returned tool calls against limit
  MLXChatRoute-->>Client: Return response or upstream error

Loading

Suggested reviewers: svonava

Priority: ⬇️ Low

Merge Risk: ⚪ Minimal · up to 35ce5

No actionable issue remains from this review; the change is mergeable after normal checks.

Pre-merge checks · 4 · 1

Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 13.64% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely identifies the main change: support for tool_choice and parallel_tool_calls on the MLX chat path.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[krisztian-gajdar]

Review round 2, head 35ce55bf (cd144be7 plus a merge of main at 78fa1b48, which only touches Rust in the gateway and the sidecar).

Gaps, both fixed with tests:

1. parallel_tool_calls: false was refused. It is now accepted and dropped from the proxied body, since mlx_lm (0.30.7, 0.31.3, 0.32.0) never reads it. The response path enforces it: a choice that carries a second call answers 502 upstream_error on a buffered response. A stream delivers the first call, then an error event and [DONE]. This matches the remote chat route's _check_tools and the local tool-call parser's stop at a second call. Tests: test_mlx_chat_parallel_tool_calls_false_returns_a_single_call[False|True] and test_mlx_chat_fails_closed_on_tool_calls_the_request_ruled_out[parallel_tool_calls_false-False|True]. Live, Qwen/Qwen3.5-4B asked for two cities answers 502 buffered, and on a stream it delivers the Paris call and then the error.
2. tool_choice: "none" only changed the prompt. Confirmed in mlx_lm 0.32.0: the tool state comes from the tokenizer, and the calls are parsed with tools=None (the qwen3_coder parser handles None). The response path now fails closed when a call comes back and the child was offered no tools. The child's body is the same for tool_choice: "none" and for a request without tools, so both cases are covered. Tests: ...ruled_out[tool_choice_none-*] and ...ruled_out[no_tools-*], each with the input from the review (an earlier assistant tool_calls turn plus its result). Live, the input reproduces with temperature: 0: the model emitted get_weather with no tools in the prompt, and the route now answers 502 (buffered) or an error event with no call delivered (stream). Before this change that call reached the client.

Risks:

• Unpinned mlx-lm>=0.30.7. Not a gap. SIE resolves both fields itself and forwards neither, and the new checks read the child's output, not its version. A later mlx_lm that honours either field can only make the required refusal stricter than needed.
• No streaming test for the tools path. Added test_mlx_chat_streams_every_tool_call_under_auto, plus the streamed parallel_tool_calls: false cases and the three streamed fail-closed cases. All of them use mlx_lm's stream shape: one delta per completed call with index, then the finish event.
• Refusal named MLX on a CPU host. Fixed. The message now reads 'tool_choice' must be 'auto' or 'none' on this device, which cannot force a tool call. test_mlx_chat_refuses_a_forced_tool_choice_before_model_load runs on mps and cpu.
• Prompt-cache reuse under none. Left as is, because it costs speed and not correctness. The local queued route also withholds tools under none (effective_tools in processors/streaming.py). Keeping the tools in the prompt would raise the rate of the fail-closed answer above.
• Overlap with feat(api): opt-in reasoning_content on chat completions via include_reasoning #728. git merge-tree of cd144be7 and feat(api): opt-in reasoning_content on chat completions via include_reasoning #728's head 71f21cb4 is clean, and the merged test_openai_local.py passes 146. Whichever PR merges second still needs a CI rerun.

Validation at this head: module 133 passed, and 12 of the new cases fail against 5f032fee. mise run test -- packages/sie_server/tests gives 11380 passed. Its 5 failures are host-filesystem tests in config/test_serving_artifacts.py and core/test_model_loader_cache.py, which this change does not touch. mise run lint and mise run typecheck pass. The PR body has the full live run and the conformance check (8 pass, 1 fail on include_reasoning, 2 recorded).