diff options
| -rw-r--r-- | assistant.py | 140 |
1 files changed, 107 insertions, 33 deletions
diff --git a/assistant.py b/assistant.py index 37e8ad8..a36d43e 100644 --- a/assistant.py +++ b/assistant.py @@ -228,6 +228,9 @@ def emit(event): now_state.clear() now_state.update({k: v for k, v in event.items() if k != "type"}) line = (json.dumps(event) + "\n").encode() + # ponytail: sends run under the caller's lock, so one stalled client holds + # everyone up for its 2 s timeout, once, before it is dropped; per-client + # send queues if that is ever felt. with clients_lock: for c in clients[:]: try: @@ -263,6 +266,11 @@ def reader(c): if not data: break buf += data + if len(buf) > 1 << 20: # a megabyte with no newline is not a command + log("client line over 1 MB, dropped") + emit({"type": "error", "message": "line too long"}) + buf = b"" + continue *lines, buf = buf.split(b"\n") for line in lines: if line.strip(): @@ -271,8 +279,18 @@ def reader(c): def serve(path=SOCK): + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + probe.connect(path) + except OSError: + pass # missing, or left over from a crash + else: + # Unlinking a live socket would orphan the running daemon's clients. + probe.close() + raise SystemExit(f"another assistant is serving {path}") + probe.close() try: - os.unlink(path) # left over from a crash + os.unlink(path) except FileNotFoundError: pass srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -285,7 +303,12 @@ def serve(path=SOCK): def accept(): while True: - c, _ = srv.accept() + try: + c, _ = srv.accept() + except OSError as e: # EMFILE and friends: log, back off, keep serving + log(f"accept failed: {e}") + time.sleep(1) + continue c.settimeout(2) # Under the lock: no event can fall between hello and the client # joining the list. @@ -325,38 +348,59 @@ def handle(raw): log(f"bad command: {raw[:200]!r}") emit({"type": "error", "message": "bad command"}) return - with lock: - if kind == "ask": - text = str(cmd.get("text", "")).strip() - if text: - interrupt() - threading.Thread(target=reply, args=(turn, text), daemon=True).start() - elif kind == "stop": + # Ids are integers: SQLite would match "1" to row 1 while Python's + # comparison with the current id would not, and a missing id would + # equal an unsaved conversation's None. + bad = ("bad id" if kind in ("open", "delete") and type(cmd.get("id")) is not int + else "bad text" if kind == "ask" and not isinstance(cmd.get("text"), str) + else "bad query" if kind == "search" and not isinstance(cmd.get("q"), str) + else None) + if bad: + emit({"type": "error", "message": bad}) + return + try: + with lock: + command(kind, cmd) + except Exception: + # A failed command must not kill the reader: the client would stay + # connected, receiving events, with nothing reading what it sends. + log(f"{kind} failed:\n" + traceback.format_exc().rstrip()) + emit({"type": "error", "message": f"{kind} failed"}) + + +def command(kind, cmd): + # Caller holds lock; arguments are already checked. + if kind == "ask": + text = cmd["text"].strip() + if text: interrupt() - elif kind == "new": + threading.Thread(target=reply, args=(turn, text), daemon=True).start() + elif kind == "stop": + interrupt() + elif kind == "new": + start_new() + emit({"type": "conversation", "id": None, "turns": []}) + elif kind == "list": + emit({"type": "conversations", "items": list_convs(DB)}) + elif kind == "open": + got = load_conv(DB, cmd["id"]) + if got is None: + emit({"type": "error", "message": "no such conversation"}) + return + interrupt() + cur.update(id=cmd["id"], history=got[0], turns=got[1]) + emit({"type": "conversation", "id": cur["id"], "turns": cur["turns"]}) + elif kind == "search": + emit({"type": "results", "items": search(DB, cmd["q"])}) + elif kind == "delete": + delete_conv(DB, cmd["id"]) + if cmd["id"] == cur["id"]: start_new() emit({"type": "conversation", "id": None, "turns": []}) - elif kind == "list": - emit({"type": "conversations", "items": list_convs(DB)}) - elif kind == "open": - got = load_conv(DB, cmd.get("id")) - if got is None: - emit({"type": "error", "message": "no such conversation"}) - return - interrupt() - cur.update(id=cmd["id"], history=got[0], turns=got[1]) - emit({"type": "conversation", "id": cur["id"], "turns": cur["turns"]}) - elif kind == "search": - emit({"type": "results", "items": search(DB, str(cmd.get("q", "")))}) - elif kind == "delete": - delete_conv(DB, cmd.get("id")) - if cmd.get("id") == cur["id"]: - start_new() - emit({"type": "conversation", "id": None, "turns": []}) - emit({"type": "conversations", "items": list_convs(DB)}) - else: - log(f"unknown command: {kind!r}") - emit({"type": "error", "message": f"unknown command {kind}"}) + emit({"type": "conversations", "items": list_convs(DB)}) + else: + log(f"unknown command: {kind!r}") + emit({"type": "error", "message": f"unknown command {kind}"}) def context_size(model): @@ -520,6 +564,8 @@ def answer(mine, text): + [{"role": "user", "content": f"[{now}] {text}"}]) first = len(msgs) # this turn's tool messages start here said, used = "", [] + # ponytail: a stopped reply notices at its next chunk, so llama-server's + # slot stays busy until then and the next question queues behind it. for rnd in range(CFG["tool_rounds"] + 1): calls, content = {}, "" # The last round offers no tools, which forces an answer. @@ -574,6 +620,7 @@ def main(): global DB # SIGTERM runs the finally below, so the socket is not left behind. signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) + served = False try: DB = db_open(DB_PATH) try: @@ -581,16 +628,17 @@ def main(): except OSError: pass # llama-server down: asked again per request serve() + served = True log(f"ready on {SOCK}") while True: signal.pause() except (KeyboardInterrupt, SystemExit) as e: - log(f"exit: {type(e).__name__}") + log(f"exit: {type(e).__name__} {e}".rstrip()) except BaseException: log("crashed:\n" + traceback.format_exc().rstrip()) raise finally: - if os.path.exists(SOCK): + if served and os.path.exists(SOCK): # never another daemon's socket os.remove(SOCK) @@ -763,6 +811,11 @@ def test(): script["delay"] = 0.02 send(type="ask", text="long one") until("token") + c3 = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + c3.connect(path) # mid-reply: hello carries the partial + h3 = json.loads(c3.makefile("rb").readline()) + assert h3["pending"]["user"] == "long one" and h3["pending"]["reply"].startswith("w0") + c3.close() send(type="stop") seen = until("done") assert seen[-1].get("stopped") is True @@ -785,6 +838,27 @@ def test(): assert until("error")[-1]["message"] == "no such conversation" cl.sendall(b"not json\n") assert until("error")[-1]["message"] == "bad command" + for cmd, why in (({"type": "open", "id": [1]}, "bad id"), + ({"type": "open", "id": str(conv)}, "bad id"), + ({"type": "delete"}, "bad id"), + ({"type": "ask", "text": None}, "bad text"), + ({"type": "search", "q": 3}, "bad query")): + send(**cmd) + assert until("error")[-1]["message"] == why, cmd + assert cur["id"] == conv # none of them touched the conversation + real, DB = DB, None # every query now raises + send(type="list") + assert until("error")[-1]["message"] == "list failed" + DB = real + send(type="list") # the reader survived it + assert [x["id"] for x in until("conversations")[-1]["items"]] == [conv] + try: + serve(path) # a second daemon on a live socket + assert False + except SystemExit as e: + assert "another assistant" in str(e) + send(type="list") # and the first one still serves + until("conversations") send(type="delete", id=conv) seen = until("conversations") assert {"type": "conversation", "id": None, "turns": []} in seen # was current |