fix: stop sequence truncation, parameter validation, special token stripping

Autoresearch deep testing (20 rounds across OpenAI SDK, Aider, LangChain,
LiteLLM, Cline, OpenCode patterns) found 8 bugs:

- P1: stop sequences not forwarded to stream_generate()/generate() in SimpleEngine
- P1: stop sequences included in output (OpenAI spec requires truncation)
- P2: n > 1 silently ignored — now returns 400
- P2: negative max_tokens accepted — now returns 400
- P2: temperature out of 0-2 range accepted — now returns 400
- P2: <|eom_id|>/<|python_tag|> Llama tokens leaking into responses
- P2: uvicorn keep-alive too short for agentic long-poll clients (now 30s)
- P3: completion_tokens off-by-one in non-streaming (re-encoding vs actual)

All 8 fixes verified against live server. 195/197 unit tests pass
(2 require live server on port 8000).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Your Name

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rapid-mlx"
-version = "0.3.7"
+version = "0.3.8"
 description = "Rapid-MLX — AI inference for Apple Silicon. Drop-in OpenAI API, 2-4x faster than Ollama."
 readme = "README.md"
 license = {text = "Apache-2.0"}

vllm_mlx/api/models.py

@@ -207,6 +207,8 @@ class ChatCompletionRequest(BaseModel):
     timeout: float | None = None
     # Thinking/reasoning control (Qwen3 style).  None = server default.
     enable_thinking: bool | None = None
+    # Number of completions (only n=1 supported)
+    n: int | None = None
 
 
 class AssistantMessage(BaseModel):

vllm_mlx/api/utils.py

@@ -15,7 +15,8 @@
 # Keeps <think>...</think> blocks intact for reasoning models
 SPECIAL_TOKENS_PATTERN = re.compile(
     r"<\|im_end\|>|<\|im_start\|>|<\|endoftext\|>|"
-    r"<\|end\|>|<\|eot_id\|>|<\|start_header_id\|>|<\|end_header_id\|>|"
+    r"<\|end\|>|<\|eot_id\|>|<\|eom_id\|>|<\|python_tag\|>|"
+    r"<\|start_header_id\|>|<\|end_header_id\|>|"
     r"<\|channel\|>|<\|message\|>|<\|start\|>|<\|return\|>|<\|call\|>|<\|constrain\|>|"
     r"</s>|<s>|<pad>|\[PAD\]|\[SEP\]|\[CLS\]|"
     r"\[e~\[|\]~b\][a-z]*|\]~!b\["

vllm_mlx/cli.py

@@ -357,7 +357,13 @@ def serve_command(args):
     print(f"  Ready: http://{host_display}:{args.port}/v1")
     print(f"  Docs:  http://{host_display}:{args.port}/docs")
     print()
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_level="info",
+        timeout_keep_alive=30,
+    )
 
 
 def bench_command(args):

vllm_mlx/engine/simple.py

@@ -392,6 +392,7 @@ async def chat(
                     max_tokens=max_tokens,
                     temperature=temperature,
                     top_p=top_p,
+                    stop=stop,
                     **kwargs_copy,
                 )
                 # Return raw text — server handles cleaning after
@@ -557,6 +558,7 @@ async def stream_chat(
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
+            stop=stop,
             **kwargs,
         ):
             yield output

vllm_mlx/models/llm.py

@@ -199,6 +199,8 @@ def generate(
         # preserved via skip_special_tokens=False decoding, and the
         # prompt cache is properly managed.
         output_text = ""
+        token_ids = []
+        finish_reason = "stop"
         for chunk in self.stream_generate(
             prompt=prompt,
             max_tokens=max_tokens,
@@ -207,18 +209,19 @@ def generate(
             stop=stop,
         ):
             output_text += chunk.text
+            if hasattr(chunk, "token") and chunk.token:
+                token_ids.append(chunk.token)
             if chunk.finished:
+                finish_reason = chunk.finish_reason or "stop"
                 break
 
-        # Tokenize output to get token IDs
-        tokens = self.tokenizer.encode(output_text)
-
-        # Determine finish reason
-        finish_reason = "length" if len(tokens) >= max_tokens else "stop"
+        # Fall back to re-encoding if no token IDs were collected
+        if not token_ids:
+            token_ids = self.tokenizer.encode(output_text)
 
         return GenerationOutput(
             text=output_text,
-            tokens=tokens,
+            tokens=token_ids,
             finish_reason=finish_reason,
         )
 
@@ -661,12 +664,18 @@ def _make_generator():
                 new_text = decoder.add_token(token_id)
                 accumulated_text += new_text
 
-                # Check for stop sequences
+                # Check for stop sequences — truncate at the stop point
+                # (OpenAI spec: stop sequence is not included in output)
                 should_stop = False
+                stop_truncate_text = None
                 if stop:
                     for stop_seq in stop:
-                        if stop_seq in accumulated_text:
+                        idx = accumulated_text.find(stop_seq)
+                        if idx != -1:
                             should_stop = True
+                            # Truncate new_text so accumulated ends just before the stop seq
+                            stop_truncate_text = new_text[: len(new_text) - (len(accumulated_text) - idx)]
+                            accumulated_text = accumulated_text[:idx]
                             break
 
                 # Check if mlx-lm signalled completion (EOS token hit)
@@ -689,7 +698,7 @@ def _make_generator():
                     cache_saved = True
 
                 yield StreamingOutput(
-                    text=new_text,
+                    text=stop_truncate_text if stop_truncate_text is not None else new_text,
                     token=response.token if hasattr(response, "token") else 0,
                     finished=finished,
                     finish_reason=finish_reason,

vllm_mlx/server.py

@@ -1710,6 +1710,29 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
                 detail=f"Invalid role '{msg.role}'. Must be one of: {', '.join(sorted(_valid_roles))}",
             )
 
+    # Validate n parameter (only n=1 supported)
+    if request.n is not None and request.n > 1:
+        raise HTTPException(
+            status_code=400,
+            detail="n > 1 is not supported. Rapid-MLX generates one completion per request.",
+        )
+
+    # Validate max_tokens (must be positive)
+    if request.max_tokens is not None and request.max_tokens < 1:
+        raise HTTPException(
+            status_code=400,
+            detail="max_tokens must be at least 1",
+        )
+
+    # Validate temperature range (OpenAI spec: 0-2)
+    if request.temperature is not None and (
+        request.temperature < 0 or request.temperature > 2
+    ):
+        raise HTTPException(
+            status_code=400,
+            detail="temperature must be between 0 and 2",
+        )
+
     # Validate top_logprobs range (OpenAI spec: 0-20)
     if request.top_logprobs is not None and (
         request.top_logprobs < 0 or request.top_logprobs > 20