Repro for mlx-openai-server __value__ race (non-deterministic)

repro_value_keyerror.py

#!/usr/bin/env python3
"""
Reproducer for mlx-openai-server v1.7.1 bug:
    ERROR | Error in text response generation: '__value__'

Root cause hypothesis: prompt_cache.get() and pop() do unguarded
`current["__value__"]` lookups, while set()/peek() use defensive `.get(...)`.
Under concurrent requests that share a token-prefix, one request's pop()
can race with another's get(), raising KeyError('__value__').

Usage:
    # Start mlx-openai-server with --tool-call-parser qwen3_coder \\
    # --reasoning-parser qwen3_5 --context-length 32768
    python3 repro_value_keyerror.py

What it does:
    - Fires 1 long-gen request (asks for 800 tokens) and N short requests
      with overlapping token-prefix, all concurrently.
    - Repeats up to MAX_ROUNDS times.
    - Polls server log for the '__value__' error.
    - Prints first reproducing round.

Stdlib only — no pip deps.
"""

import json
import os
import sys
import time
import urllib.request
import urllib.error
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

ENDPOINT = os.environ.get("MLX_URL", "http://127.0.0.1:11436/v1/chat/completions")
MODEL = os.environ.get("MLX_MODEL", "qwen3.6-mlx")
LOG_PATH = Path(os.environ.get("MLX_LOG", "/tmp/mlx-server.log"))
ERROR_NEEDLE = "Error in text response generation: '__value__'"
MAX_ROUNDS = int(os.environ.get("MAX_ROUNDS", "8"))
CONCURRENCY = int(os.environ.get("CONCURRENCY", "4"))

# A long shared prefix maximizes the probability of trie-node contention.
# In the original incident the failing prefix was ~19k tokens; we synthesize
# a 5-8k token prefix here to keep the repro tractable.
_FILLER = " ".join(
    f"Punkt {i}: MLX nutzt unified memory und teilt Gewichte zwischen RAM und GPU."
    for i in range(1, 401)
)
SHARED_PREFIX = (
    "You are a careful technical assistant. The user is preparing extensive "
    "notes for a deep-dive on quantized LLM inference. Stay in German. "
    "Be precise. Reference background context below:\n\n" + _FILLER + "\n\n"
)

LONG_PROMPT = SHARED_PREFIX + (
    "Schreibe jetzt 500 Worte ueber Apple Silicon Inferenz mit MLX. "
    "Strukturiere mit 5 Headers: Architektur, Quantisierung, Cache-Management, "
    "Tool-Use-Integration, Open-Issues. Keine Code-Blocks."
)

SHORT_PROMPT = SHARED_PREFIX + "Nenne 3 MLX-Vorteile in einem Satz."

# Optional tools payload — original incident included tools in the request.
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read a file",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}},
                "required": ["path"],
            },
        },
    }
]


def fire_request(prompt: str, max_tokens: int) -> dict:
    payload = json.dumps({
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.3,
        "top_p": 0.85,
        "presence_penalty": 1.5,
        "chat_template_kwargs": {"enable_thinking": False},
        "tools": TOOLS,
        "tool_choice": "auto",
    }).encode("utf-8")
    req = urllib.request.Request(
        ENDPOINT,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            body = json.loads(resp.read())
        return {
            "status": resp.status,
            "elapsed": round(time.time() - t0, 1),
            "completion_tokens": body.get("usage", {}).get("completion_tokens"),
            "finish_reason": body.get("choices", [{}])[0].get("finish_reason"),
        }
    except urllib.error.HTTPError as e:
        try:
            detail = json.loads(e.read()).get("detail")
        except Exception:
            detail = None
        return {
            "status": e.code,
            "elapsed": round(time.time() - t0, 1),
            "error": detail or str(e),
        }
    except Exception as e:
        return {"status": "EXC", "elapsed": round(time.time() - t0, 1), "error": repr(e)}


def log_has_error_after(baseline_line: int) -> bool:
    try:
        with LOG_PATH.open() as fh:
            for i, line in enumerate(fh):
                if i >= baseline_line and ERROR_NEEDLE in line:
                    return True
    except FileNotFoundError:
        return False
    return False


def current_log_lines() -> int:
    try:
        with LOG_PATH.open() as fh:
            return sum(1 for _ in fh)
    except FileNotFoundError:
        return 0


def main():
    print(f"target: {ENDPOINT}")
    print(f"model:  {MODEL}")
    print(f"log:    {LOG_PATH}")
    print(f"max_rounds={MAX_ROUNDS}  concurrency={CONCURRENCY}")
    print()

    with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
        for r in range(1, MAX_ROUNDS + 1):
            baseline = current_log_lines()
            print(f"[round {r}/{MAX_ROUNDS}] firing 1 long + {CONCURRENCY - 1} short ...")
            futures = [pool.submit(fire_request, LONG_PROMPT, 800)]
            for _ in range(CONCURRENCY - 1):
                futures.append(pool.submit(fire_request, SHORT_PROMPT, 80))
            results = [f.result() for f in futures]
            for idx, res in enumerate(results):
                tag = "LONG" if idx == 0 else f"SHORT-{idx}"
                print(f"  {tag:8s}  {res}")
            if log_has_error_after(baseline):
                print()
                print(f"REPRODUCED on round {r}.")
                print(f"server log lines >{baseline} contain: {ERROR_NEEDLE!r}")
                return 0
            print("  no error this round.")
            time.sleep(0.5)

    print()
    print(f"NOT reproduced in {MAX_ROUNDS} rounds. Try higher CONCURRENCY or more rounds.")
    return 1


if __name__ == "__main__":
    sys.exit(main())