#!/usr/bin/env python3
"""mock_server.py - 最小的 OpenAI 兼容服务端, 用于端到端测试 ligent. 

不依赖任何第三方包. 行为完全由请求内容决定, 保持确定性: 

  - 用户消息含 ERROR_TEST  -> 返回 HTTP 400 + error 对象
  - 尚未有 tool 结果        -> 流式返回一个 tool_call (run_command echo ...)
  - 用户消息含 twice 且只执行过一次工具 -> 再返回第二个 tool_call
  - 用户消息含 MANYSTEPS_TEST          -> 连要 4 次工具往返才收工
  - 其余                    -> 流式/非流式返回最终文本, 内容里回显最近一次工具输出
"""
import json
import os
import sys
import time
from http.server import BaseHTTPRequestHandler, HTTPServer


STATE = {}   # 按消息文本隔离的重试计数


def sse(wfile, obj):
    wfile.write(("data: " + json.dumps(obj, ensure_ascii=False) + "\n\n").encode("utf-8"))
    wfile.flush()


def pairing_error(msgs):
    """服务端的硬要求: 带 tool_calls 的 assistant 后面必须**紧跟**与每个
    tool_call_id 一一对应的 tool 消息, 中间不能夹别的东西. 

    真实服务端就是这么回 400 的, mock 也照做 —— 这样任何会写坏历史的回归
     (比如把图片附件插在两条 tool 消息中间、或者中断时丢下没回应的调用)
    都会在 e2e 里当场炸出来, 而不是等到用户真拿它去连服务端才发现. 
    返回错误文案, 没问题返回 None. """
    i, n = 0, len(msgs)
    while i < n:
        m = msgs[i]
        if m.get("role") == "assistant" and m.get("tool_calls"):
            want = [tc.get("id") for tc in m["tool_calls"]]
            got, j = [], i + 1
            while j < n and msgs[j].get("role") == "tool":
                got.append(msgs[j].get("tool_call_id"))
                j += 1
            for w in want:
                if w not in got:
                    return ("An assistant message with 'tool_calls' must be "
                            "followed by tool messages responding to each "
                            "'tool_call_id'. (insufficient tool messages "
                            "following tool_calls message)")
        i += 1
    return None


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

    def log_message(self, *args):
        pass

    # -- 工具调用 (流式增量) --
    def emit_tool_call(self, index, call_id, name, arguments):
        # 故意把 arguments 切成多片, 检验客户端的增量拼接
        half = len(arguments) // 2
        sse(self.wfile, {"choices": [{"index": 0, "delta": {
            "role": "assistant", "content": ""}, "finish_reason": None}]})
        sse(self.wfile, {"choices": [{"index": 0, "delta": {"tool_calls": [
            {"index": index, "id": call_id, "type": "function",
             "function": {"name": name, "arguments": ""}}]}, "finish_reason": None}]})
        for part in (arguments[:half], arguments[half:]):
            if not part:
                continue
            sse(self.wfile, {"choices": [{"index": 0, "delta": {"tool_calls": [
                {"index": index, "function": {"arguments": part}}]},
                "finish_reason": None}]})
        sse(self.wfile, {"choices": [{"index": 0, "delta": {},
                                      "finish_reason": "tool_calls"}]})

    def emit_text(self, text):
        for piece in (text[: len(text) // 2], text[len(text) // 2:]):
            if not piece:
                continue
            sse(self.wfile, {"choices": [{"index": 0, "delta": {"content": piece},
                                          "finish_reason": None}]})
        sse(self.wfile, {"choices": [{"index": 0, "delta": {},
                                      "finish_reason": "stop"}],
                         "usage": {"prompt_tokens": 1, "completion_tokens": 1,
                                   "total_tokens": 2}})

    def reject(self, status, message):
        """回一个错误响应 (探测/校验用)"""
        payload = json.dumps({"error": {"message": message}}).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def send_text(self, text, stream, model):
        """按流式/非流式把一段最终文本发回去. """
        if not stream:
            payload = json.dumps({"choices": [{"index": 0,
                                               "message": {"role": "assistant",
                                                           "content": text},
                                               "finish_reason": "stop"}],
                                  "model": model}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()
        self.emit_text(text)
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()

    def do_GET(self):
        # doctor 的连通性/重定向检查走 GET. 
        # /redirect 开头故意返回 307, 用来验证客户端能把重定向指出来. 
        # doctor 的认证探测: 端点要求认证, 且只认 "Bearer good-key". 
        #  带对 key -> 200 (认证过了)
        #  带错 key -> 401 (密钥被拒)
        #  不带 key -> 401 (端点要求认证, 而你没配)
        # 客户端要靠"有没有配 key"把后两种 401 区分开. 
        if self.path.startswith("/authed"):
            auth = self.headers.get("Authorization") or ""
            ok = (auth == "Bearer good-key")
            self.send_response(200 if ok else 401)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", "2")
            self.end_headers()
            self.wfile.write(b"{}")
            return

        # /model 不带参数时拉模型列表. 
        # 故意混进一个重复项和一个缺 id 的条目 —— 客户端要能去重、跳过坏条目. 
        # MODEL_LIST=bare 时改回裸数组, MODEL_LIST=empty 时不返回任何模型, 
        # 两种兼容/异常形状都由 e2e 覆盖. 
        if "/models" in self.path:
            # 模式从**文件**读, 不从环境变量读: mock 是脚本开头起好的, 
            # 给客户端进程设的 env 它根本看不见 (第一版就是这么假通过的). 
            mode = ""
            mode_file = os.environ.get("MODELS_MODE_FILE", "")
            if mode_file:
                try:
                    with open(mode_file, encoding="utf-8") as f:
                        mode = f.read().strip()
                except OSError:
                    mode = ""
            if mode == "empty":
                payload = json.dumps({"object": "list", "data": []}).encode()
            elif mode == "bare":
                payload = json.dumps(["alpha-model", "beta-model"]).encode()
            elif mode == "garbage":
                payload = b"<html>not json at all</html>"
            else:
                # 默认列表里**故意包含当前模型** (deepseek-flash): 
                # 这样才验得了"光标初始落在当前项上". 
                # 另外塞一个重复项和一个缺 id 的坏条目, 验去重与跳过. 
                payload = json.dumps({"object": "list", "data": [
                    {"id": "deepseek-flash", "object": "model", "owned_by": "mock"},
                    {"id": "alpha-model", "object": "model", "owned_by": "mock"},
                    {"id": "beta-model", "object": "model", "owned_by": "mock"},
                    {"id": "alpha-model", "object": "model", "owned_by": "mock"},
                    {"object": "model", "owned_by": "mock"},
                ]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        if self.path.startswith("/redirect"):
            self.send_response(307)
            self.send_header("Location", "https://example.invalid/v1/chat/completions")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        body = b"{}"
        self.send_response(404)          # 端点根路径本来就没有东西
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        length = int(self.headers.get("Content-Length", "0"))
        raw = self.rfile.read(length)

        # ---- UTF-8 严格校验: 真服务端就是这么拒的 ----
        # DeepSeek 对请求体里非法的 UTF-8 回的是
        #   400 Failed to parse the request body as JSON:
        #       messages[43].content: invalid unicode code point
        # 而坏数据一旦进了对话历史, 之后**每一次**请求都会失败, 除非丢掉整轮
        # 对话. 所以 mock 必须一样严格, 否则"序列化器吐出了非法 UTF-8"这类 bug
        # 在 e2e 里永远看不见. 
        #  (json.loads(bytes) 其实也会因此抛异常、被下面那个 except 兜成 400, 
        #   但那条笼统的 400 什么都没说明, 而且换个实现就可能漏掉, 所以显式判. )
        try:
            raw.decode("utf-8")
        except UnicodeDecodeError as exc:
            self.reject(400, "Failed to parse the request body as JSON: "
                             "invalid unicode code point at byte %d" % exc.start)
            return

        # ---- 假的 Files API ----
        # 只做够用的部分: 认 multipart 里有没有图片内容, 发一个递增的 file_id. 
        try:
            body = json.loads(raw or b"{}")
        except Exception:
            self.send_error(400)
            return

        messages = body.get("messages", [])
        model = body.get("model", "?")
        stream = bool(body.get("stream"))

        # 把 user 消息拍平成纯文本再匹配. 带附件的消息 content 是**块数组**
        #  (/include、read_image 的附件都是), 不拍平的话 first_user 会是个
        # list, `"标记" in list` 永远不成立 —— 分支会莫名其妙地匹配不上. 
        def user_text(m):
            c = m.get("content")
            if isinstance(c, str):
                return c
            if isinstance(c, list):
                return " ".join(p.get("text", "") for p in c
                                if isinstance(p, dict) and p.get("type") == "text")
            return ""

        users = [user_text(m) for m in messages if m.get("role") == "user"]
        first_user = users[0] if users else ""
        last_user = users[-1] if users else ""
        tool_msgs = [m for m in messages if m.get("role") == "tool"]
        tools_sent = body.get("tools")

        # ---- 结构校验: 和真服务端一样严 ----
        perr = pairing_error(messages)
        if perr:
            self.reject(400, perr)
            return

        # ---- 能力探测序列 ----
        # ligent 会在第一次真实请求之前发几个试探请求 (基线 / 图片 / tools), 
        # 它们会被真实处理, 所以这里必须排在所有标记分支之前. 
        # 辨认靠形状: 刚好一条 user 消息, 内容是 "1", 或者带一张 1x1 的图. 
        def looks_like_probe(msgs):
            if len(msgs) != 1 or msgs[0].get("role") != "user":
                return False
            c = msgs[0].get("content")
            if c == "1":
                return True
            if isinstance(c, list):
                for p in c:
                    if not isinstance(p, dict) or p.get("type") != "image_url":
                        continue
                    u = (p.get("image_url") or {}).get("url", "")
                    if "iVBORw0KGgo" in u:
                        return True
            return False

        if looks_like_probe(messages):
            has_image = isinstance(messages[0].get("content"), list)
            # 用模型名模拟"不支持某能力"的服务端: 400 是明确的拒绝信号
            if "no-vision" in model and has_image:
                return self.reject(400, "unsupported content type: image_url")
            if "no-tools" in model and tools_sent:
                return self.reject(400, "unsupported parameter: tools")
            if "broken-probe" in model:
                return self.reject(401, "invalid api key")
            self.send_text("1", stream, model)
            return

        # 提权: 让模型发一个带 elevated 的 run_command
        if "ELEVATED_TEST" in first_user:
            if not tool_msgs:
                args = json.dumps({"command": "echo 提权后的输出", "elevated": True})
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_elev", "run_command", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            got = ""
            for m in messages:
                if m.get("role") == "tool":
                    got = m.get("content") or ""
            self.send_text("ELEV_SAW[%s]" % got.replace("\n", " | "), stream, model)
            return

        # /include: 把收到的 user 内容块报回来. 
        # 用 first_user 判断——附带之后最新那条 user 是带块的消息, last_user 会漂. 
        if "INCLUDE_TEST" in first_user:
            texts, imgs, names = 0, 0, []
            for m in messages:
                if m.get("role") != "user":
                    continue
                c = m.get("content")
                # 没有附件时 content 是纯字符串, 那也算"一个文本块"
                if isinstance(c, str):
                    texts += 1
                    continue
                if not isinstance(c, list):
                    continue
                for p in c:
                    if not isinstance(p, dict):
                        continue
                    if p.get("type") == "text":
                        texts += 1
                        t = p.get("text", "")
                        for line in t.split("\n"):
                            if line.startswith("===== "):
                                names.append(line[6:].split(" (")[0])
                    elif p.get("type") == "image_url":
                        imgs += 1
            self.send_text("INC[text=%d img=%d files=%s]"
                           % (texts, imgs, ",".join(names)), stream, model)
            return

        # 回显收到的工具名, 用来验证"不该暴露的工具确实没发出去"
        if "TOOLS_ECHO_TEST" in first_user:
            names = [((t or {}).get("function") or {}).get("name")
                     for t in (tools_sent or [])]
            self.send_text("TOOLS[%s]" % ",".join(n for n in names if n),
                           stream, model)
            return

        # 权限/参数校验: 客户端必须把 tools 传上来
        assert tools_sent, "请求里没有 tools 字段"

        if "ERROR_TEST" in first_user:
            payload = json.dumps({"error": {"message": "mock 模拟的服务端错误",
                                            "type": "invalid_request_error"}}).encode()
            self.send_response(400)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 消息条数超过阈值就报"上下文超限", 用于验证客户端的自动裁剪. 
        # 客户端丢掉最老的一轮后条数会降下来, 于是下一轮就能成功. 
        if "TOOLONG_TEST" in first_user:
            if len(messages) > 3:
                payload = json.dumps({"error": {
                    "message": "This model's maximum context length is 8192 tokens. "
                               "However, your messages resulted in 9000 tokens. "
                               "Please reduce the length of the messages.",
                    "code": "context_length_exceeded"}}).encode()
                self.send_response(400)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return
            text = "当前历史 %d 条消息, 正常回复" % len(messages)
            payload = json.dumps({"choices": [{"index": 0,
                                               "message": {"role": "assistant",
                                                           "content": text},
                                               "finish_reason": "stop"}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 把收到的用户消息原样回显, 用于验证"一次请求里到底装了什么"
        if "ECHO_TEST" in first_user:
            text = "ECHO[" + first_user.replace("\n", "\\n") + "]"
            payload = json.dumps({"choices": [{"index": 0,
                                               "message": {"role": "assistant",
                                                           "content": text},
                                               "finish_reason": "stop"}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 复刻 OpenRouter 对"key 格式不对"的回应: 头其实带了, 
        # 但文案说的是 "Missing Authentication header", 容易把人引偏
        if "AUTHFAIL_TEST" in first_user:
            payload = json.dumps({"error": {"message": "Missing Authentication header",
                                            "code": 401}}).encode()
            self.send_response(401)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 重定向: 用于验证客户端不静默跟随 (跨源跟随会丢掉 Authorization, 
        # 服务端只会返回 401 "Missing Authentication header", 误导用户去查 key)
        if "REDIRECT_TEST" in first_user:
            self.send_response(307)
            self.send_header("Location", "https://example.invalid/v1/chat/completions")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        # 前几次返回 503, 用于验证重试与 Retry-After
        if "RETRY_TEST" in first_user:
            key = "retry:" + first_user
            STATE[key] = STATE.get(key, 0) + 1
            STATE["retry"] = STATE[key]
            if STATE["retry"] < 3:
                payload = json.dumps({"error": {
                    "message": "mock 临时不可用 (第 %d 次)" % STATE["retry"]}}).encode()
                self.send_response(503)
                self.send_header("Content-Type", "application/json")
                self.send_header("Retry-After", "1")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return
            text = "重试后成功 (共尝试 %d 次)" % STATE["retry"]
            payload = json.dumps({"choices": [{"index": 0,
                                               "message": {"role": "assistant",
                                                           "content": text},
                                               "finish_reason": "stop"}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 永远 503, 用于验证重试次数耗尽后正确报错
        if "ALWAYS503_TEST" in first_user:
            payload = json.dumps({"error": {"message": "mock 持续不可用"}}).encode()
            self.send_response(503)
            self.send_header("Content-Type", "application/json")
            self.send_header("Retry-After", "1")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        # 超过客户端响应体上限的超大响应: 客户端应当主动中止而不是被撑爆
        if "HUGE_TEST" in first_user:
            total = 40 * 1024 * 1024
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(total))
            self.end_headers()
            chunk = b'{"pad":"' + b'x' * (1024 * 1024 - 12) + b'"}'
            try:
                for _ in range(40):
                    self.wfile.write(chunk)
                    self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError, OSError):
                pass          # 客户端主动断开是预期行为
            return

        # 畸形 SSE: 各种不合法的片段, 客户端必须不崩溃且给出合理结果
        if "GARBAGE_TEST" in first_user:
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.end_headers()
            lines = [
                "data: {\n\n",
                "data: not json at all\n\n",
                "data: [DONE]\n\n",
                'data: {"choices":[]}\n\n',
                'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n',
                "\xe6\xb3\xa8\xe9\x87\x8a\xe8\xa1\x8c\n\n",   # 非 ASCII 裸行
                "data:\n\n",
                "\x01\x02\xff\xfe\n\n",
            ]
            for line in lines:
                self.wfile.write(line.encode("utf-8", "surrogateescape"))
            self.wfile.flush()
            return

        # 传输中途断掉: 声明了 Content-Length 却提前关闭连接
        if "TRUNCATE_TEST" in first_user:
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Content-Length", "100000")
            self.end_headers()
            self.wfile.write(b"data: {\"choices\":[{\"delta\":{\"content\":\"\xe5\x8d\x8a\xe9\x80\x94\xe5\xb4\xa9\xe6\xba\x83\"}}]}\n\n")
            self.wfile.flush()
            self.close_connection = True
            return

        # 固定的 Markdown 回复, 用于验证流式渲染
        if "MARKDOWN_TEST" in first_user:
            md = (
                "# 一级标题\n"
                "## 二级标题\n"
                "\n"
                "这是 **粗体** 和 *斜体* 和 _也是斜体_ 和 `code` 混排. \n"
                "参考 [官方文档](https://example.com/doc) 与 arr[i]. \n"
                "\n"
                "- 列表甲\n- 列表乙\n"
                "\n"
                "| 名称 | 数量 | 说明 |\n"
                "|:-----|-----:|:----:|\n"
                "| 苹果 | 3 | 常见 |\n"
                "| 香蕉 | 12 | 较长 |\n"
                "\n"
                "```c\n"
                "int x = 1;  // **不应加粗** 也不高亮 [方括号]\n"
                "```\n"
            )
            if not stream:
                payload = json.dumps({"choices": [{"index": 0,
                                                   "message": {"role": "assistant",
                                                               "content": md},
                                                   "finish_reason": "stop"}],
                                      "model": model}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.end_headers()
            # 故意按 1 字符一片下发, 最大化"标记被切断"的概率
            for ch in md:
                sse(self.wfile, {"choices": [{"index": 0,
                                              "delta": {"content": ch},
                                              "finish_reason": None}]})
            sse(self.wfile, {"choices": [{"index": 0, "delta": {},
                                          "finish_reason": "stop"}]})
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
            return

        # 思考模式: 验证 reasoning_content 既能被累积, 也能随 assistant 消息
        # 回传. DeepSeek 要求请求带 tools 时把历史里的 reasoning_content 全带上, 
        # 漏了下一轮就 400, 所以这里第二轮专门检查它有没有回来. 
        if "THINK_TEST" in first_user:
            if not tool_msgs:
                if not stream:
                    msg = {"role": "assistant", "content": None,
                           "reasoning_content": "先想一下要不要调用工具. ",
                           "tool_calls": [{"id": "call_1", "type": "function",
                                           "function": {
                                               "name": "run_command",
                                               "arguments": '{"command": "echo think-ok"}'}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                # 思考过程分片下发: 客户端必须累积, 而不是只留最后一片
                for part in ("先想一下", "要不要调用工具. "):
                    sse(self.wfile, {"choices": [{"index": 0,
                                                  "delta": {"reasoning_content": part},
                                                  "finish_reason": None}]})
                sse(self.wfile, {"choices": [{"index": 0, "delta": {"tool_calls": [
                    {"index": 0, "id": "call_1", "type": "function",
                     "function": {"name": "run_command",
                                  "arguments": '{"command": "echo think-ok"}'}}]},
                    "finish_reason": None}]})
                sse(self.wfile, {"choices": [{"index": 0, "delta": {},
                                              "finish_reason": "tool_calls"}]})
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return

            assistants = [m for m in messages if m.get("role") == "assistant"]
            got = "".join(m.get("reasoning_content", "") for m in assistants)
            th = body.get("thinking")
            th = th.get("type") if isinstance(th, dict) else "none"
            self.send_text("THINK[echo=%s thinking=%s effort=%s]" % (
                "yes" if "要不要调用工具" in got else "NO", th,
                body.get("reasoning_effort") or "none"), stream, model)
            return

        # 慢速响应: 先静默一段 (网络/排队), 再吐思维链, 再静默一段, 
        # 最后给正文. 用来验证等待动画的两段标签: "等待响应" → "思考中". 
        # 名字刻意不含 "THINK_TEST" —— 上面那些标记都是按子串匹配的, 
        # 叫 SLOWTHINK_TEST 会被 THINK_TEST 那个分支先截胡. 
        if "SLOWRESP_TEST" in first_user:
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.end_headers()
            time.sleep(0.8)          # 这期间一个字节都没发
            sse(self.wfile, {"choices": [{"index": 0,
                                          "delta": {"reasoning_content": "先想一下. "},
                                          "finish_reason": None}]})
            time.sleep(0.8)          # 这期间模型在推理
            self.send_text("答案", stream=True, model=model)
            return

        # meta 模式: 第一轮让 meta 模型调 delete_rounds 删掉第 2 轮, 
        # 第二轮收工. 用来端到端验证 /meta 的完整链路. 
        if "META_TEST" in first_user:
            if not tool_msgs:
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                import json as _j
                sse(self.wfile, {"choices": [{"index": 0, "delta": {"tool_calls": [
                    {"index": 0, "id": "meta_1", "type": "function",
                     "function": {"name": "delete_rounds",
                                  "arguments": _j.dumps({"from": 2, "to": 2})}}]},
                    "finish_reason": None}]})
                sse(self.wfile, {"choices": [{"index": 0, "delta": {},
                                              "finish_reason": "tool_calls"}]})
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("已删除第 2 轮. ", stream, model)
            return

        # 按"最后一条用户消息"触发: 上面那些模式都看 first_user, 
        # 同一个进程里的第二轮就再也触发不了了, 没法验证"作用域仅限本轮". 
        if "WRITE_LAST_TEST" in last_user:
            if messages[-1].get("role") == "user":      # 新一轮的开头
                args = json.dumps({"path": "build/tmp/confirm_probe_turn2.txt",
                                   "content": "第二轮的"}, ensure_ascii=False)
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": "call_t2", "type": "function",
                         "function": {"name": "write_file", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_t2", "write_file", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("第二轮写完了", stream, model)
            return

        # 发出一个 write_file 调用: 用来验证"在确认提示上按 Ctrl-C"到底
        # 会不会把文件写下去
        if "WRITE_TEST" in first_user:
            if not tool_msgs:
                args = json.dumps({"path": "build/tmp/confirm_probe.txt",
                                   "content": "应该被拒绝的内容"}, ensure_ascii=False)
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": "call_w", "type": "function",
                         "function": {"name": "write_file", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_w", "write_file", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("写完了", stream, model)
            return

        # 读一个**含非法 UTF-8 字节**的文件. 工具输出里出现坏字节的来源很日常: 
        # 某条命令在非 UTF-8 locale 下打出来的路径、被截断的多字节字符、二进制
        # 文件里混的一小段文本. 这段字节会进对话历史, 然后被序列化进请求体 ——
        # 如果序列化器原样吐出, 服务端就会 400, 而且之后每次请求都 400. 
        # 这个分支的作用就是让那条路径在 e2e 里真的被走一遍. 
        if "UTF8_TEST" in first_user:
            if not tool_msgs:
                args = json.dumps({"path": "build/tmp/bad_utf8.txt"})
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": "call_u8", "type": "function",
                         "function": {"name": "read_file", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_u8", "read_file", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            # 第二轮: 能走到这里就说明请求体过了上面的 UTF-8 校验. 
            # 把工具消息里有没有替换字符也顺带报出来, 便于断言. 
            got = ""
            for m in tool_msgs:
                c = m.get("content")
                if isinstance(c, str):
                    got += c
            self.send_text("读完了, 替换字符=%d" % got.count("\ufffd"), stream, model)
            return

        # 一轮里连写两个文件: 用来验证 Always/neVer 的作用域就是"这一轮"
        #  (名字刻意不含 "WRITE_TEST", 以免被上面那个分支先截胡)
        if "WRITE2_TEST" in first_user:
            n = len(tool_msgs)
            if n < 2:
                args = json.dumps({"path": "build/tmp/confirm_probe%d.txt" % (n + 1),
                                   "content": "第 %d 个" % (n + 1)})
                call_id = "call_w%d" % n
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": call_id, "type": "function",
                         "function": {"name": "write_file", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, call_id, "write_file", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("两个都写完了", stream, model)
            return

        # 让模型去读一个"扩展名是 .png、内容不是图片"的文件, 
        # 验证 read_image 按内容判断、并明确拒绝
        # 两种"读不了": 内容不是图片、以及大得内联不进去. 
        # 注意: mock 的标记靠子串匹配, 是平铺的命名空间 —— 上面第 249 行的
        # HUGE_TEST 分支会吃掉任何含 "HUGE_TEST" 的标记, 所以这里叫 OVERSIZE. 
        # 两者都该在**本地**被拒 —— 工具返回错误, 之后就没有图片块可发. 
        if "IMAGE_BAD_TEST" in last_user or "IMAGE_OVERSIZE_TEST" in last_user:
            if not tool_msgs:
                path = ("build/tmp/huge.png" if "IMAGE_OVERSIZE_TEST" in last_user
                        else "build/tmp/notanimage.png")
                args = json.dumps({"path": path})
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_bad", "read_image", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("读不了就算了", stream, model)
            return

        # 一条 assistant 里发**两个**工具调用, 第一个是 read_image. 
        # 这是"图片附件插在两条 tool 消息中间"那个 bug 的最小复现形状: 
        # 只要附件没被攒到最后一起推, 上面的结构校验就会回 400. 
        # 注意用 first_user: 图片附件本身也是一条 user 消息, 而且排在最后, 
        # 所以第二轮请求的 last_user 是那个附件 (一个 list), 不是这条标记. 
        if "IMAGE_PAIR_TEST" in first_user:
            if not tool_msgs:
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_img", "read_image",
                                    json.dumps({"path": "build/tmp/dot.png"}))
                self.emit_tool_call(1, "call_cmd", "run_command",
                                    json.dumps({"command": "echo pair-ok"}))
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            found = []
            for m in messages:
                c = m.get("content")
                if isinstance(c, list):
                    for part in c:
                        if isinstance(part, dict) and part.get("type") == "image_url":
                            found.append("img")
            self.send_text("PAIR[img=%d tool=%d]"
                           % (len(found), len(tool_msgs)), stream, model)
            return

        # 图片: 先让模型调一次 read_image, 下一轮把收到的图片内联块报回来. 
        # Files API 已经不用了 —— 请求里应该直接是 image_url + data: base64. 
        if "IMAGE_TEST" in first_user:
            if not tool_msgs:
                args = json.dumps({"path": "build/tmp/dot.png"})
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": "call_img", "type": "function",
                         "function": {"name": "read_image", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_img", "read_image", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return

            # 第二轮: 请求里应该已经带上内联图片了. 
            # 报文形状: {"type":"image_url","image_url":{"url":"data:image/png;base64,..."}}
            found, notes = [], []
            for m in messages:
                c = m.get("content")
                if isinstance(c, list):
                    for part in c:
                        if not isinstance(part, dict):
                            continue
                        if part.get("type") == "image_url":
                            url = (part.get("image_url") or {}).get("url", "")
                            if url.startswith("data:"):
                                head = url.split(";", 1)[0]      # data:image/png
                                b64 = url.split(",", 1)[1] if "," in url else ""
                                found.append("%s:%d" % (head, len(b64)))
                        # 图片被拒之后, 附件消息里应该换成一句给模型看的说明. 
                        # 把它回显出来, 端到端才能证明**模型真的收到了**这句话, 
                        # 而不只是我们自己的历史被改了. 
                        elif part.get("type") == "text" and m.get("role") == "user":
                            notes.append(part.get("text", ""))

            # IMAGE_REJECT_TEST: 第一次故意说服务端不收这张图, 
            # 客户端应该把图片从历史里摘掉再试 (而不是每轮都失败)
            if "IMAGE_REJECT_TEST" in last_user and not STATE.get("told_once"):
                STATE["told_once"] = 1
                payload = json.dumps({"error": {
                    "message": "image rejected: unsupported image variant"}}).encode()
                self.send_response(400)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return
            reply = "IMAGE[%s]" % (",".join(found) if found else "none")
            if notes:
                reply += " NOTE[" + "|".join(notes) + "]"
            self.send_text(reply, stream, model)
            return

        # edit_file 的**调用展示**: 成功时屏幕上该是"路径 + 行号"加 -/+ 新旧对照, 
        # 而不是一整段参数 JSON (old_string/new_string 原样打出来会刷屏). 
        # EDITTEST_OK 成功剧本, EDITTEST_FAIL 行号越界剧本. 
        if "EDITTEST_OK" in first_user or "EDITTEST_FAIL" in first_user:
            if tool_msgs:
                self.send_text("DONE[%s]" % tool_msgs[-1].get("content", "")[:60],
                               stream, model)
                return
            line = 999 if "EDITTEST_FAIL" in first_user else 2
            # 12 行的 old_string: 既覆盖"多行改动"的显示, 也覆盖"超过 8 行要折叠"
            old = "\n".join("OLD%d" % i for i in range(1, 13)) if line == 2 else "OLD"
            args = json.dumps({"path": "build/tmp/edit_probe.txt", "line": line,
                               "old_string": old, "new_string": "NEW"})
            if not stream:
                msg = {"role": "assistant", "content": None, "tool_calls": [
                    {"id": "call_edit", "type": "function",
                     "function": {"name": "edit_file", "arguments": args}}]}
                payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                   "finish_reason": "tool_calls"}],
                                      "model": model}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.end_headers()
            self.emit_tool_call(0, "call_edit", "edit_file", args)
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
            return

        # 多步 (4 次工具往返才收工): 用来钉住"往返次数默认不限". 
        # 默认配置下必须跑得完; 显式设了正数上限时要停在那个上限上并如实报出. 
        if "MANYSTEPS_TEST" in first_user:
            done = len(tool_msgs)
            if done < 4:
                args = json.dumps({"command": "echo step-%d" % (done + 1)})
                if not stream:
                    msg = {"role": "assistant", "content": None, "tool_calls": [
                        {"id": "call_ms%d" % (done + 1), "type": "function",
                         "function": {"name": "run_command", "arguments": args}}]}
                    payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                       "finish_reason": "tool_calls"}],
                                          "model": model}).encode()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                self.emit_tool_call(0, "call_ms%d" % (done + 1), "run_command", args)
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
                return
            self.send_text("STEPS_DONE[%d]" % done, stream, model)
            return

        want_tools = "twice" in first_user
        need = 2 if want_tools else 1

        if len(tool_msgs) < need:
            if not stream:
                args = json.dumps({"command": "echo hello-from-tool"}, ensure_ascii=False)
                msg = {"role": "assistant", "content": None, "tool_calls": [
                    {"id": "call_1", "type": "function",
                     "function": {"name": "run_command", "arguments": args}}]}
                payload = json.dumps({"choices": [{"index": 0, "message": msg,
                                                   "finish_reason": "tool_calls"}],
                                      "model": model}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)
                return

            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.end_headers()
            if len(tool_msgs) == 0:
                self.emit_tool_call(0, "call_1", "run_command",
                                    '{"command": "echo hello-from-tool"}')
            else:
                self.emit_tool_call(0, "call_2", "run_command",
                                    '{"command": "echo second-call"}')
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
            return

        last = tool_msgs[-1].get("content", "")
        text = "FINAL: " + last.strip().replace("\n", " / ")

        if not stream:
            payload = json.dumps({"choices": [{"index": 0,
                                               "message": {"role": "assistant",
                                                           "content": text},
                                               "finish_reason": "stop"}],
                                  "model": model}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
            return

        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()
        self.emit_text(text)
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8731
    server = HTTPServer(("127.0.0.1", port), Handler)
    print("mock server on %d" % port, flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()