z-bar-qt/scripts/mock-llama-server.py

#!/usr/bin/env python3
"""Mock llama-server for exercising zshell's LLM error handling.

Implements the OpenAI-compatible subset of llama-server that the C++
backend talks to, and lets you trigger specific failure modes from the
QML chat UI by typing a keyword as your message.

Wire format is taken from ggml-org/llama.cpp (tools/server):
  - SSE events are "data: <json>\n\n" (note the space after "data:").
  - A clean stream ends with "data: [DONE]\n\n".
  - Errors are {"error": {"code", "message", "type"}}.
  - A stream that dies mid-way sends the error event, then closes
    WITHOUT [DONE].
  - A request that errors before the first token is a plain HTTP 4xx/5xx
    with the error JSON body (not a stream).
  - /props is a JSON object; n_ctx lives under
    default_generation_settings.n_ctx.

Run:
    python3 scripts/mock-llama-server.py [--port 8080]

Then point zshell's LLM endpoint at http://localhost:8080.

Trigger words (type as the chat message, case-insensitive):
    normal            -> healthy stream, ends with [DONE]  (also the default)
    instant error     -> 500 before any token ("server_error")
    context overflow  -> 400 "exceeds the available context size (2048 ...)";
                          backend retries once at 2048, then fails
    error during      -> streams a few tokens, then an SSE error event,
                          closes without [DONE]
    no done           -> streams a few tokens, closes without [DONE] and
                          without an error (server killed mid-response)
    200 error         -> HTTP 200 with a plain JSON error body (not SSE)
    tool call         -> streams reasoning, then a tool_calls round calling
                          readfile on /etc/os-release (finish_reason
                          "tool_calls"); after you approve (or deny) the
                          permission dialog, the follow-up round that
                          carries the tool result gets a normal summary
                          answer. Exercises the tool permission dialog,
                          tool execution and the multi-round transcript.
    multiple tool     -> two tool calls in one round (readfile + webfetch
    calls                 of this mock's own / endpoint); approve both at
                          once
    unknown tool      -> a tool call for a tool the client does not have;
                          approving it yields an "unknown tool" error
                          result

Control endpoint (for the non-chat routes / global state), via curl:
    curl -XPOST localhost:8080/_mock -d '{"loading": true}'
        /props, /v1/models, /health all answer 503 "Loading model"
    curl -XPOST localhost:8080/_mock -d '{"no_models": true}'
        /v1/models returns an empty list
    curl -XPOST localhost:8080/_mock -d '{"chat_mode": "error_during"}'
        force a chat behaviour regardless of the message text
        (values: normal|instant_error|context_overflow|error_during|
                 no_done|200_error|tool_call|tool_call_multi|
                 tool_call_unknown, or null to go back to keywords)
    curl -XPOST localhost:8080/_mock -d '{"reset": true}'
"""

import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

MODEL_ID = "mock-model"
CTX_SIZE = 8192

# Set in main(); used by the tool-call behaviours so webfetch can point
# at this server itself.
_host = "127.0.0.1"
_port = 8080

# Mutable, shared across handler threads.
_state_lock = threading.Lock()
_state = {
    "loading": False,       # 503 "Loading model" on /props, /v1/models, /health
    "no_models": False,     # /v1/models returns an empty list
    "chat_mode": None,      # force a chat behaviour, overriding keywords
}


def _now():
    return int(time.time())


def _chunk(cid, delta, finish_reason=None):
    return {
        "id": cid,
        "object": "chat.completion.chunk",
        "created": _now(),
        "model": MODEL_ID,
        "system_fingerprint": "mock",
        "choices": [
            {
                "index": 0,
                "delta": delta,
                "finish_reason": finish_reason,
            }
        ],
    }


def _usage_chunk(cid, prompt_tokens, completion_tokens):
    return {
        "id": cid,
        "object": "chat.completion.chunk",
        "created": _now(),
        "model": MODEL_ID,
        "choices": [],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }


def _error(code, message, etype):
    return {"error": {"code": code, "message": message, "type": etype}}


def _sse(obj):
    # Exact llama.cpp format_oai_sse shape: "data: " + compact json + "\n\n".
    return ("data: " + json.dumps(obj, separators=(",", ":")) + "\n\n").encode()


def _sse_done():
    return b"data: [DONE]\n\n"


# Chat behaviours. Each returns ("json", status, body_bytes) for a
# non-streaming answer, or ("sse", 200, None) and then writes to the
# provided stream writer.
def _behaviour_for(text):
    t = text.lower()
    if "instant error" in t:
        return "instant_error"
    if "context overflow" in t:
        return "context_overflow"
    if "error during" in t:
        return "error_during"
    if "no done" in t or "truncate" in t:
        return "no_done"
    if "200 error" in t:
        return "200_error"
    # Order matters: "multiple tool calls" contains "tool call".
    if "two tool calls" in t or "two tool call" in t:
        return "tool_call_duo"
    if "some tool calls" in t:
        return "tool_call_some"
    if "many tool calls" in t:
        return "tool_call_many"
    if "unknown tool" in t:
        return "tool_call_unknown"
    if "tool call" in t:
        return "tool_call"
    return "normal"


_TOOL_MODES = ("tool_call", "tool_call_duo", "tool_call_some", "tool_call_many", "tool_call_unknown")


def _tool_call_delta(index, call_id=None, name=None, arguments=None):
    call = {"index": index}
    if call_id is not None:
        call["id"] = call_id
    if name is not None or arguments is not None:
        function = {}
        if name is not None:
            function["name"] = name
        if arguments is not None:
            function["arguments"] = arguments
        call["function"] = function
    return {"tool_calls": [call]}


def _stream_tool_calls(w, cid, reasoning, calls):
    # One tool round: optional reasoning, then the calls (llama.cpp
    # streams the id/name with the first fragment of each call and the
    # arguments as successive fragments), then finish_reason
    # "tool_calls". The client concatenates the argument fragments and
    # shows the permission dialog on stream end.
    w(_sse(_chunk(cid, {"role": "assistant"})))
    if reasoning:
        for piece in reasoning:
            w(_sse(_chunk(cid, {"reasoning_content": piece})))
            time.sleep(0.05)
    for i, (call_id, name, arguments) in enumerate(calls):
        w(_sse(_chunk(cid, _tool_call_delta(i, call_id=call_id, name=name))))
        time.sleep(0.05)
        for fragment in arguments:
            w(
                _sse(
                    _chunk(cid, _tool_call_delta(i, arguments=fragment))
                )
            )
            time.sleep(0.05)
    w(_sse(_chunk(cid, {}, finish_reason="tool_calls")))
    w(_sse(_usage_chunk(cid, 24, len(reasoning or []) + len(calls) * 2)))
    w(_sse_done())


def _tool_result_snippets(messages, limit=160):
    snippets = []
    for msg in messages or []:
        if msg.get("role") != "tool":
            continue
        content = msg.get("content")
        if isinstance(content, list):
            content = " ".join(
                part.get("text", "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            )
        flat = " ".join(str(content or "").split())
        snippets.append(flat[:limit] + ("..." if len(flat) > limit else ""))
    return snippets


def _stream_tool_round_done(w, cid, messages):
    # Follow-up round: the client sent the tool results back; answer
    # with a plain content summary so the turn completes.
    w(_sse(_chunk(cid, {"role": "assistant"})))
    parts = ["mock:", "tool", "round", "complete."]
    for i, snippet in enumerate(_tool_result_snippets(messages), 1):
        parts.append(f"result {i}:")
        parts.append(snippet or "(empty)")
    for word in parts:
        w(_sse(_chunk(cid, {"content": word + " "})))
        time.sleep(0.02)
    w(_sse(_chunk(cid, {}, finish_reason="stop")))
    w(_sse(_usage_chunk(cid, 40, len(parts) + 1)))
    w(_sse_done())


def _stream_normal(w, cid, words):
    w(_sse(_chunk(cid, {"role": "assistant"})))
    for word in words:
        w(_sse(_chunk(cid, {"content": word + " "})))
        time.sleep(0.05)
    w(_sse(_chunk(cid, {}, finish_reason="stop")))
    w(_sse(_usage_chunk(cid, 12, len(words) + 1)))
    w(_sse_done())


def _stream_error_during(w, cid, words):
    w(_sse(_chunk(cid, {"role": "assistant"})))
    for word in words:
        w(_sse(_chunk(cid, {"content": word + " "})))
        time.sleep(0.05)
    # Mid-stream failure: error event, then the connection closes without
    # [DONE] (llama.cpp terminates the stream on a task error).
    w(_sse(_error(500, "Mock: generation failed mid-stream", "server_error")))


def _stream_no_done(w, cid, words):
    w(_sse(_chunk(cid, {"role": "assistant"})))
    for word in words:
        w(_sse(_chunk(cid, {"content": word + " "})))
        time.sleep(0.05)
    # Clean FIN, no [DONE], no error: server died mid-response.


def _last_user_text(messages):
    for msg in reversed(messages or []):
        if msg.get("role") == "user":
            content = msg.get("content")
            if isinstance(content, list):
                # multimodal shape: pull out the text parts
                parts = []
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "text":
                        parts.append(part.get("text", ""))
                content = " ".join(parts)
            return content or ""
    return ""


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "MockLlama/1.0"

    # --- low-level response helpers -----------------------------------
    def _send_json(self, status, obj):
        body = json.dumps(obj, separators=(",", ":")).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Connection", "close")
        self.end_headers()
        self.wfile.write(body)

    def _send_loading_error(self):
        # llama.cpp middleware: server not ready -> 503 on every route.
        self._send_json(
            503, _error(503, "Loading model", "unavailable_error")
        )

    def _not_ready(self):
        with _state_lock:
            return _state["loading"]

    def _stream(self, status, writer):
        self.send_response(status)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "close")
        self.end_headers()
        try:
            writer(self.wfile.write)
        except (BrokenPipeError, ConnectionResetError):
            pass
        finally:
            try:
                self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError):
                pass

    # --- routing -------------------------------------------------------
    def do_GET(self):
        path = self.path.split("?", 1)[0].rstrip("/") or "/"
        if path in ("/health", "/v1/health"):
            if self._not_ready():
                return self._send_loading_error()
            return self._send_json(200, {"status": "ok"})
        if path in ("/props", "/v1/props"):
            if self._not_ready():
                return self._send_loading_error()
            return self._send_json(
                200,
                {
                    "total_slots": 1,
                    "model_alias": MODEL_ID,
                    "model_path": "/mock/model.gguf",
                    "default_generation_settings": {
                        "params": {"n_predict": 2048},
                        "n_ctx": CTX_SIZE,
                    },
                    "is_sleeping": False,
                },
            )
        if path in ("/models", "/v1/models"):
            if self._not_ready():
                return self._send_loading_error()
            with _state_lock:
                no_models = _state["no_models"]
            data = [] if no_models else [
                {
                    "id": MODEL_ID,
                    "object": "model",
                    "created": _now(),
                    "owned_by": "mock",
                }
            ]
            return self._send_json(200, {"object": "list", "data": data})
        if path == "/":
            return self._send_json(200, {"status": "ok", "mock": True})
        return self._send_json(404, _error(404, "not found", "not_found_error"))

    def do_POST(self):
        path = self.path.split("?", 1)[0].rstrip("/") or "/"
        length = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(length) if length else b""
        try:
            body = json.loads(raw) if raw else {}
        except json.JSONDecodeError:
            body = {}

        if path == "/_mock":
            return self._mock_control(body)

        if path in ("/chat/completions", "/v1/chat/completions"):
            if self._not_ready():
                return self._send_loading_error()
            return self._chat(body)

        return self._send_json(404, _error(404, "not found", "not_found_error"))

    # --- /_mock control ------------------------------------------------
    def _mock_control(self, body):
        with _state_lock:
            if body.get("reset"):
                _state["loading"] = False
                _state["no_models"] = False
                _state["chat_mode"] = None
            if "loading" in body:
                _state["loading"] = bool(body["loading"])
            if "no_models" in body:
                _state["no_models"] = bool(body["no_models"])
            if "chat_mode" in body:
                _state["chat_mode"] = body["chat_mode"]
            snapshot = dict(_state)
        self.log_message("control -> %s", json.dumps(snapshot))
        return self._send_json(200, {"ok": True, "state": snapshot})

    # --- /chat/completions --------------------------------------------
    def _chat(self, body):
        text = _last_user_text(body.get("messages"))
        with _state_lock:
            forced = _state["chat_mode"]
        mode = forced or _behaviour_for(text)
        stream = bool(body.get("stream", True))
        cid = "chatcmpl-mock"
        words = ["mock", "says", "hello", "from", "the", "mock", "server"]
        self.log_message("chat mode=%s stream=%s text=%r", mode, stream, text)

        if mode == "instant_error":
            # Error before the first token: plain HTTP error, not a stream.
            return self._send_json(
                500, _error(500, "Mock: instant server error", "server_error")
            )

        if mode == "context_overflow":
            return self._send_json(
                400,
                _error(
                    400,
                    "request (100000 tokens) exceeds the available context "
                    "size (2048 tokens), try increasing it",
                    "exceed_context_size_error",
                ),
            )

        if mode == "200_error":
            # 200 with a plain JSON error body (not SSE) -- the case the
            # data: filter would otherwise swallow.
            return self._send_json(
                200, _error(400, "Mock: 200-with-json-error", "server_error")
            )

        if mode in _TOOL_MODES:
            # Continuation round: the previous round's tool results came
            # back through the client (role: "tool" messages).
            if any(m.get("role") == "tool" for m in body.get("messages") or []):
                return self._stream(
                    200,
                    lambda w: _stream_tool_round_done(
                        w, cid, body.get("messages")
                    ),
                )
            if mode == "tool_call":
                calls = [
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "os-release\"}"],
                    )
                ]
                reasoning = [
                    "The user wants a tool call. ",
                    "I will read /etc/os-release.",
                ]
            elif mode == "tool_call_duo":
                calls = [
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "hostname\"}"],
                    ),
                    (
                        "call_mock_2",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                ]
                reasoning = ["Two independent tool calls, ",
                             "batched together."]
            elif mode == "tool_call_some":
                calls = [
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "hostname\"}"],
                    ),
                    (
                        "call_mock_2",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_3",
                        "readfile",
                        ['{"path": "/etc/', "os-release\"}"],
                    ),
                    (
                        "call_mock_4",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_5",
                        "readfile",
                        ['{"path": "/proc/', "uptime\"}"],
                    ),
                    (
                        "call_mock_6",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                ]
                reasoning = ["Two independent tool calls, ",
                             "batched together."]
            elif mode == "tool_call_many":
                calls = [
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "hostname\"}"],
                    ),
                    (
                        "call_mock_2",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_3",
                        "readfile",
                        ['{"path": "/etc/', "os-release\"}"],
                    ),
                    (
                        "call_mock_4",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_5",
                        "readfile",
                        ['{"path": "/proc/', "uptime\"}"],
                    ),
                    (
                        "call_mock_6",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "hostname\"}"],
                    ),
                    (
                        "call_mock_2",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_3",
                        "readfile",
                        ['{"path": "/etc/', "os-release\"}"],
                    ),
                    (
                        "call_mock_4",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_5",
                        "readfile",
                        ['{"path": "/proc/', "uptime\"}"],
                    ),
                    (
                        "call_mock_6",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_1",
                        "readfile",
                        ['{"path": "/etc/', "hostname\"}"],
                    ),
                    (
                        "call_mock_2",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_3",
                        "readfile",
                        ['{"path": "/etc/', "os-release\"}"],
                    ),
                    (
                        "call_mock_4",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                    (
                        "call_mock_5",
                        "readfile",
                        ['{"path": "/proc/', "uptime\"}"],
                    ),
                    (
                        "call_mock_6",
                        "webfetch",
                        [
                            '{"url": "http://%s:%d/", ' % (_host, _port),
                            '"format": "text"}',
                        ],
                    ),
                ]
                reasoning = ["Two independent tool calls, ",
                             "batched together."]
            else:  # tool_call_unknown
                calls = [
                    (
                        "call_mock_3",
                        "nonexistent_tool",
                        ['{"target": "test"}'],
                    )
                ]
                reasoning = ["Calling a tool the client ", "does not have."]
            self.log_message(
                "tool mode=%s calls=%s", mode, [c[1] for c in calls]
            )
            return self._stream(
                200, lambda w: _stream_tool_calls(w, cid, reasoning, calls)
            )

        if not stream:
            # Non-streaming success: single chat.completion object.
            content = " ".join(words)
            return self._send_json(
                200,
                {
                    "id": cid,
                    "object": "chat.completion",
                    "created": _now(),
                    "model": MODEL_ID,
                    "choices": [
                        {
                            "index": 0,
                            "message": {"role": "assistant", "content": content},
                            "finish_reason": "stop",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 12,
                        "completion_tokens": len(words),
                        "total_tokens": 12 + len(words),
                    },
                },
            )

        if mode == "error_during":
            return self._stream(200, lambda w: _stream_error_during(w, cid, words))
        if mode == "no_done":
            return self._stream(200, lambda w: _stream_no_done(w, cid, words))
        return self._stream(200, lambda w: _stream_normal(w, cid, words))

    def log_message(self, format, *args):
        # Compact single-line logging to stdout.
        print("%s - %s" % (self.log_date_time_string(), format % args), flush=True)


def main():
    ap = argparse.ArgumentParser(
        description="Mock llama-server for exercising zshell LLM error handling"
    )
    ap.add_argument("--port", type=int, default=8080)
    ap.add_argument("--host", default="127.0.0.1")
    args = ap.parse_args()

    global _host, _port
    _host = args.host
    _port = args.port

    server = ThreadingHTTPServer((args.host, args.port), Handler)
    server.daemon_threads = True
    print(
        f"mock llama-server on http://{args.host}:{args.port} "
        f"(model={MODEL_ID}, ctx={CTX_SIZE})\n"
        "type trigger words in the chat UI, or curl /_mock to change state",
        flush=True,
    )
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()