fix(server): flush tool call when generation ends in "tool" state (Mistral/Devstral)

[rinaldofesta]

Summary

The OpenAI server drops tool calls for models whose tool_call_end is empty — notably Mistral / Devstral, which emit [TOOL_CALLS]<name>[ARGS]{json}. The response comes back as an empty assistant message (no content, no tool_calls) even though the model generated a perfectly valid tool call.

Root cause

For these models the state machine never transitions back to "normal" (there is no closing tool token), so generation finishes while still in the "tool" state. In APIHandler.handle_completion the end-of-generation flush is guarded on prev_state == "tool":

if prev_state == "tool" and tool_text:
    tool_calls.append(tool_text)
    made_tool_call = True

But the terminal generation event resets prev_state (it is None by the time the loop exits), so the accumulated tool_text is never appended to tool_calls. The ToolCallFormatter then receives an empty list and emits nothing.

I confirmed this with a debug print inside the flush:

prev_state=None  tool_text='read[ARGS]{"file_path": "/tmp/foo.txt"}'  tool_calls=[]

tool_text is collected correctly (with [ARGS] intact) — the only problem is the guard.

Fix

Flush whenever tool_text is non-empty. tool_text is only ever populated while in the "tool" state and is reset to "" on the tool -> normal transition, so this is also correct for models that do close their tool calls (no double flush):

if tool_text:
    tool_calls.append(tool_text)
    made_tool_call = True

Verification

End-to-end against a real mlx_lm.server serving mlx-community/Devstral-Small-2-24B-Instruct-2512-4bit:

• Before: finish_reason: "stop", empty message, tool_calls absent (both streaming and non-streaming).
• After: finish_reason: "tool_calls" with the structured call — non-streaming {"name":"read","arguments":"{\"file_path\": \"/tmp/foo.txt\"}"}, and the streaming deltas assemble to the same.

Added tests/test_server.py::TestMistralToolStateFlush, which drives the request handler with a generation stream ending in the "tool" state (the terminal event carries state=None, reproducing the reset). It fails on main and passes with the fix.

python -m unittest tests.test_server tests.test_tool_parsing   # 26 tests, OK
black mlx_lm/server.py tests/test_server.py                     # clean

AI assistance disclosure

Diagnosed and implemented with AI assistance (Claude). I reproduced the bug against a live Devstral server, traced the root cause with an instrumented build, and verified the fix end-to-end (streaming + non-streaming) and via the new regression test before submitting. Happy to adjust the test harness to your preferred style.

[angeloskath]

Closed due to #1501.