Hermit/scripts/e2e_webui_ask.py

"""WebUI ask_user round trip against a fake OpenAI-compatible API.

  python3 scripts/e2e_webui_ask.py      (CST_AGENT= picks the binary)

Checks: the turn really waits for the answer, a page opened while the
question is open still receives it, and Last-Event-ID resumes the stream.
"""
import json, os, shutil, subprocess, sys, tempfile, threading, time, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

BIN = os.environ.get("CST_AGENT", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cst-agent"))
REQS = []


class H(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def do_GET(self):
        body = json.dumps({"data": [{"id": "deepseek-flash"}]}).encode()
        self.send_response(200)
        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):
        n = int(self.headers.get("Content-Length", 0))
        REQS.append(json.loads(self.rfile.read(n)))
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.end_headers()
        base = {"id": "x", "object": "chat.completion.chunk", "model": "deepseek-flash"}

        def ev(o):
            self.wfile.write(b"data: " + json.dumps({**base, **o}).encode() + b"\n\n")
            self.wfile.flush()
        if len(REQS) == 1:
            args = json.dumps({"question": "Pick a colour?", "options": "red\nblue"})
            ev({"choices": [{"index": 0, "delta": {"role": "assistant", "tool_calls": [
                {"index": 0, "id": "call_1", "type": "function",
                 "function": {"name": "ask_user", "arguments": args}}]}}]})
            ev({"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]})
        else:
            ev({"choices": [{"index": 0, "delta": {"role": "assistant", "content": "done"}}]})
            ev({"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


fails = 0


def check(ok, msg):
    global fails
    print(("PASS " if ok else "FAIL ") + msg)
    fails += 0 if ok else 1


class Sse:
    def __init__(self, base, last_id=None):
        self.lines = []
        req = urllib.request.Request(base + "/api/events")
        if last_id is not None:
            req.add_header("Last-Event-ID", str(last_id))
        self.r = urllib.request.urlopen(req, timeout=60)
        threading.Thread(target=self.pump, daemon=True).start()

    def pump(self):
        try:
            for raw in self.r:
                self.lines.append(raw.decode())
        except Exception:
            pass

    def events(self, name):
        out = []
        for i, l in enumerate(self.lines):
            if l.strip() == "event: " + name and i + 1 < len(self.lines):
                out.append(json.loads(self.lines[i + 1][6:]))
        return out

    def ids(self):
        return [int(l[4:]) for l in self.lines if l.startswith("id: ")]

    def close(self):
        try:
            self.r.close()
        except Exception:
            pass


def wait(fn, t=15):
    end = time.time() + t
    while time.time() < end:
        v = fn()
        if v:
            return v
        time.sleep(0.1)
    return None


def post(base, path, obj):
    req = urllib.request.Request(base + path, data=json.dumps(obj).encode(),
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=5) as r:
            return r.status
    except urllib.error.HTTPError as e:
        return e.code


api = ThreadingHTTPServer(("127.0.0.1", 0), H)
threading.Thread(target=api.serve_forever, daemon=True).start()
ws = tempfile.mkdtemp(prefix="cst-ask-")
log = open(ws + "/server.log", "w+")
p = subprocess.Popen([BIN, "--web", "--web-port", "0", "--workspace", ws, "--new",
                      "--api-key", "sk-x", "--provider", "deepseek", "--model", "deepseek-flash",
                      "--api-base", "http://127.0.0.1:%d/v1" % api.server_address[1],
                      "--protocol", "native-eos", "--approval", "auto"],
                     stdout=log, stderr=subprocess.STDOUT)
try:
    def port():
        log.seek(0)
        for l in log.read().splitlines():
            if l.startswith("webui: http://127.0.0.1:"):
                return l.split(":")[3].split("/")[0]
    pt = wait(port)
    base = "http://127.0.0.1:%s" % pt
    a = Sse(base)
    time.sleep(0.3)
    check(post(base, "/api/turn", {"text": "ask me"}) == 202, "turn accepted")
    ask = wait(lambda: a.events("ask_user"))
    check(bool(ask), "ask_user event raised")
    time.sleep(1.5)
    check(len(REQS) == 1, "the turn waits for the answer (requests so far: %d)" % len(REQS))
    # a page opened now must still see the open question
    b = Sse(base)
    check(bool(wait(lambda: b.events("ask_user"), 5)), "a late page gets the open question")
    # reconnect with Last-Event-ID: nothing before that id is resent
    last = a.ids()[-1]
    c = Sse(base, last_id=last - 1)
    got = wait(lambda: c.ids(), 5)
    check(bool(got) and got[0] == last, "Last-Event-ID resumes after that id (%s, want %d)" % (got[:1] if got else got, last))
    if ask:
        check(post(base, "/api/ask_user", {"id": ask[0]["id"], "text": "blue"}) == 200, "answer accepted")
    end = wait(lambda: a.events("turn_end"), 20)
    check(bool(end) and end[0].get("rc") == 0, "turn ended rc=0")
    tool_msgs = [m for m in (REQS[1]["messages"] if len(REQS) > 1 else []) if m.get("role") == "tool"]
    check(bool(tool_msgs) and "blue" in json.dumps(tool_msgs), "the answer reached the model: %s" % json.dumps(tool_msgs)[:200])
    for s in (a, b, c):
        s.close()
finally:
    p.terminate()
    try:
        p.wait(20)
    except subprocess.TimeoutExpired:
        p.kill()
        check(False, "server hung on SIGTERM")
    check(p.returncode is not None and p.returncode < 128 and p.returncode >= 0, "clean exit (%s)" % p.returncode)
    log.close()
    shutil.rmtree(ws, ignore_errors=True)
print("FAILED:", fails)
sys.exit(1 if fails else 0)