Protocol: check command arguments, keep readers alive, one daemon

From review of the protocol commit. An exception inside a command, such as
an id SQLite cannot bind, killed that client's reader and left a deaf
connection still receiving events; commands now run under a try that
answers error. Ids must be integers: delete with no id compared None to an
unsaved conversation's None and wiped it, and a string id deleted the row
through SQLite's affinity while the Python comparison missed, leaving the
current conversation pointing at nothing so every later answer failed to
save. ask, search and their arguments are type-checked instead of coerced.

serve() refuses to start over a socket that answers, since unlinking it
orphaned the live daemon's clients, and main() removes the socket only if
it served it. The accept loop logs and continues on errors instead of
dying silently, and a client line over a megabyte is dropped. The two
known ceilings, a stalled client delaying everyone once for 2 s and a
stopped reply holding the llama-server slot until its next chunk, are
marked. New checks cover each of these and a client connecting mid-reply.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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