Take questions on the socket and stream replies back

Each client gets a reader thread; commands are ask, stop, new, list, open,
search and delete. A reply streams think, token and tool events through
say(), which checks the turn counter under the same lock that stop takes,
so a stopped reply cannot send a token after its done. A finished turn is
saved; a stopped one is not. hello carries the current conversation and
any reply in flight, so a reconnecting canvas redraws both. Errors from
llama-server become an error event and the daemon stays up; every exit
path of main() logs a line, which the voice version's silent death lacked.

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

1 files changed, 363 insertions, 49 deletions

diff --git a/assistant.py b/assistant.py
index f455d29..37e8ad8 100644
--- a/assistant.py
+++ b/assistant.py
@@ -208,11 +208,19 @@ def delete_conv(conn, conv):
         conn.execute("DELETE FROM conversations WHERE id = ?", (conv,))
 
 
-# The canvas's feed: JSON lines on a Unix socket. With no client connected
-# every event is dropped.
+# The canvas's feed: JSON lines on a Unix socket, both ways. Every client
+# gets every event. A client that stalls for 2 s is dropped rather than
+# holding up the others.
 clients, clients_lock = [], threading.Lock()
 now_state = {"state": "idle"}             # repeated in hello for a late client
-ctx = {"model": None, "max": 0}           # context size of the model in use
+ctx = {"model": None, "max": 0, "used": 0}
+# One lock for the conversation, the turn counter and the database. Emits
+# that depend on the turn happen under it, so a stopped reply cannot send a
+# token after its done.
+lock = threading.RLock()
+cur = {"id": None, "history": [], "turns": [], "pending": None}
+turn = 0          # bumped by every ask, stop, new and open: older replies are stale
+DB = None
 
 
 def emit(event):
@@ -223,14 +231,43 @@ def emit(event):
     with clients_lock:
         for c in clients[:]:
             try:
-                c.sendall(line)           # non-blocking: a stalled client errors
-            except OSError:
+                c.sendall(line)
+            except OSError:               # gone, or stalled past its timeout
                 clients.remove(c)
                 c.close()
 
 
 def hello():
-    return {"type": "hello", "ctx_max": ctx["max"], **now_state}
+    return {"type": "hello", **now_state, "model": ctx["model"], "ctx_used": ctx["used"],
+            "ctx_max": ctx["max"], "id": cur["id"], "turns": cur["turns"],
+            "pending": cur["pending"]}
+
+
+def drop(c):
+    with clients_lock:
+        if c in clients:
+            clients.remove(c)
+    c.close()
+
+
+def reader(c):
+    # A timed-out recv is only an idle client; the socket stays usable.
+    buf = b""
+    while True:
+        try:
+            data = c.recv(65536)
+        except TimeoutError:
+            continue
+        except OSError:
+            break
+        if not data:
+            break
+        buf += data
+        *lines, buf = buf.split(b"\n")
+        for line in lines:
+            if line.strip():
+                handle(line)
+    drop(c)
 
 
 def serve(path=SOCK):
@@ -239,7 +276,7 @@ def serve(path=SOCK):
     except FileNotFoundError:
         pass
     srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    old = os.umask(0o177)                 # 0600: only this user reads the stream
+    old = os.umask(0o177)                 # 0600: only this user talks to it
     try:
         srv.bind(path)
     finally:
@@ -249,17 +286,79 @@ def serve(path=SOCK):
     def accept():
         while True:
             c, _ = srv.accept()
-            try:
-                c.sendall((json.dumps(hello()) + "\n").encode())
-            except OSError:
-                c.close()
-                continue
-            c.setblocking(False)
-            with clients_lock:
-                clients.append(c)
+            c.settimeout(2)
+            # Under the lock: no event can fall between hello and the client
+            # joining the list.
+            with lock:
+                try:
+                    c.sendall((json.dumps(hello()) + "\n").encode())
+                except OSError:
+                    c.close()
+                    continue
+                with clients_lock:
+                    clients.append(c)
+            threading.Thread(target=reader, args=(c,), daemon=True).start()
     threading.Thread(target=accept, daemon=True).start()
 
 
+def interrupt():
+    # Caller holds lock. A reply in flight ends now: done marks it stopped,
+    # its partial text is not saved, its thread notices the turn and returns.
+    global turn
+    turn += 1
+    if cur["pending"] is not None:
+        cur["pending"] = None
+        emit({"type": "done", "id": cur["id"], "stopped": True})
+        emit({"type": "state", "state": "idle"})
+
+
+def start_new():
+    interrupt()
+    cur.update(id=None, history=[], turns=[])
+
+
+def handle(raw):
+    try:
+        cmd = json.loads(raw)
+        kind = cmd["type"]
+    except (ValueError, TypeError, KeyError):
+        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":
+            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.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}"})
+
+
 def context_size(model):
     # The router can swap models, so ask again whenever the id changes.
     if ctx["model"] != model:
@@ -324,26 +423,27 @@ def collect_calls(calls, fragments):
 def llm_model():
     # The router keeps one model resident; asking for another would swap it
     # out under other clients, so use whatever is not unloaded.
-    with urllib.request.urlopen(CFG["llm"] + "/models") as r:
+    with urllib.request.urlopen(CFG["llm"] + "/models", timeout=10) as r:
         models = json.load(r)["data"]
     return next((m["id"] for m in models if m["status"]["value"] != "unloaded"),
                 models[0]["id"])
 
 
 def ask(messages, calls, tools=None):
-    # Yields reply text as it streams; tool-call fragments go into calls.
+    # Yields ("think", text) and ("token", text) as they stream; tool-call
+    # fragments go into calls.
     model = llm_model()
     context_size(model)
     body = {"model": model, "messages": messages, "stream": True,
             "stream_options": {"include_usage": True},
-            "chat_template_kwargs": {"enable_thinking": False}}
+            "chat_template_kwargs": {"enable_thinking": True}}
     if tools:
         body["tools"] = tools
     if CFG["temperature"] is not None:
         body["temperature"] = CFG["temperature"]
     req = urllib.request.Request(CFG["llm"] + "/chat/completions", json.dumps(body).encode(),
                                  {"Content-Type": "application/json"})
-    with urllib.request.urlopen(req) as r:
+    with urllib.request.urlopen(req, timeout=300) as r:
         for line in r:
             if not line.startswith(b"data: ") or line.strip() == b"data: [DONE]":
                 continue
@@ -351,39 +451,151 @@ def ask(messages, calls, tools=None):
             if chunk.get("usage"):
                 # How full the model's context is: history, tool results and
                 # tool definitions all count.
-                used = chunk["usage"]["prompt_tokens"]
-                log(f"context: {used} prompt tokens")
-                emit({"type": "ctx", "used": used, "max": ctx["max"]})
+                ctx["used"] = chunk["usage"]["prompt_tokens"]
+                emit({"type": "ctx", "used": ctx["used"], "max": ctx["max"]})
             choices = chunk["choices"]
             if not choices:
                 continue
             delta = choices[0]["delta"]
             collect_calls(calls, delta.get("tool_calls") or [])
+            if delta.get("reasoning_content"):
+                yield "think", delta["reasoning_content"]
             if delta.get("content"):
-                yield delta["content"]
+                yield "token", delta["content"]
+
+
+def say(mine, kind, value, done=False):
+    # Sends one piece of the reply, unless the reply went stale; False tells
+    # the reply thread to give up.
+    with lock:
+        if turn != mine:
+            return False
+        p = cur["pending"]
+        if kind == "think":
+            p["think"] += value
+            emit({"type": "think", "delta": value})
+        elif kind == "token":
+            p["reply"] += value
+            emit({"type": "token", "delta": value})
+        else:
+            if not done:
+                p["tools"].append(value)
+            emit({"type": "tool", "name": value, "done": done})
+            emit({"type": "state", "state": "thinking"} if done
+                 else {"type": "state", "state": "tool", "tool": value})
+        return True
+
 
+def reply(mine, text):
+    try:
+        answer(mine, text)
+    except Exception as e:
+        log("reply failed:\n" + traceback.format_exc().rstrip())
+        with lock:
+            if turn == mine:
+                cur["pending"] = None
+                emit({"type": "error", "message": str(e) or type(e).__name__})
+                emit({"type": "state", "state": "idle"})
+
+
+def answer(mine, text):
+    with lock:
+        if turn != mine:
+            return
+        history = list(cur["history"])
+        cur["pending"] = {"user": text, "reply": "", "think": "", "tools": []}
+        emit({"type": "user", "text": text})
+        emit({"type": "state", "state": "thinking"})
+    log(f"you: {text}")
+    question = {"role": "user", "content": text}
+    # The model has no clock: without the date it makes one up for tools.
+    # Also stamped on the question itself, which history does not keep: with
+    # the time only in the system prompt, "what time is it?" asked again an
+    # hour later copied the old answer from history. Stamping every message
+    # in history fixed that but lost tool arguments (account filter 0/6).
+    now = time.strftime("%A %d %B %Y, %H:%M")
+    # ponytail: last 10 turns, summarise if long sessions need memory
+    msgs = ([{"role": "system", "content": f"{CFG['system']} It is {now}."}]
+            + [m for past in history[-10:] for m in past]
+            + [{"role": "user", "content": f"[{now}] {text}"}])
+    first = len(msgs)                            # this turn's tool messages start here
+    said, used = "", []
+    for rnd in range(CFG["tool_rounds"] + 1):
+        calls, content = {}, ""
+        # The last round offers no tools, which forces an answer.
+        for kind, piece in ask(msgs, calls, TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
+            if not say(mine, kind, piece):
+                log("interrupted")
+                return
+            if kind == "token":
+                content += piece
+        said += content
+        if not calls:
+            break
+        msgs.append({"role": "assistant", "content": content, "tool_calls": [
+            {"id": c["id"], "type": "function",
+             "function": {"name": c["name"], "arguments": c["arguments"]}}
+            for c in calls.values()]})
+        # ponytail: a stop does not kill a running tool, its result is dropped
+        for c in calls.values():
+            if not say(mine, "tool", c["name"]):
+                log("interrupted")
+                return
+            used.append(c["name"])
+            result = run_tool(TOOLS, c["name"], c["arguments"])
+            log(f"tool {c['name']}({c['arguments']}): {result[:120]!r}")
+            msgs.append({"role": "tool", "tool_call_id": c["id"], "content": result})
+            if not say(mine, "tool", c["name"], done=True):
+                log("interrupted")
+                return
+    # History keeps the tool calls: with only what was said, the model saw
+    # facts answered without tools and invented the next ones (no tool call
+    # 9/10 on a replayed session, 0/10 with the calls kept). Results are
+    # replaced by a placeholder, for context: same 10/10, and a follow-up on
+    # an old result calls the tool again (10/10). A stopped exchange is
+    # dropped whole, so user and assistant still alternate.
+    past = [{**m, "content": "(shown to the user at the time)"} if m["role"] == "tool" else m
+            for m in msgs[first:]]
+    entry = [question] + past + [{"role": "assistant", "content": said}]
+    with lock:
+        if turn != mine:
+            log("interrupted")
+            return
+        cur["id"] = save_turn(DB, cur["id"], text, said, used, entry)
+        cur["history"].append(entry)
+        cur["turns"].append({"user": text, "reply": said, "tools": used})
+        cur["pending"] = None
+        emit({"type": "done", "id": cur["id"]})
+        emit({"type": "state", "state": "idle"})
+    log("assistant:", said)
 
 
 def main():
+    global DB
     # SIGTERM runs the finally below, so the socket is not left behind.
     signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
     try:
-        context_size(llm_model())         # so the first hello has it
-    except OSError:
-        pass                              # llama-server down: asked again per request
-    serve()
-    log(f"ready on {SOCK}")
-    try:
+        DB = db_open(DB_PATH)
+        try:
+            context_size(llm_model())     # so the first hello has it
+        except OSError:
+            pass                          # llama-server down: asked again per request
+        serve()
+        log(f"ready on {SOCK}")
         while True:
             signal.pause()
-    except KeyboardInterrupt:
-        pass
+    except (KeyboardInterrupt, SystemExit) as e:
+        log(f"exit: {type(e).__name__}")
+    except BaseException:
+        log("crashed:\n" + traceback.format_exc().rstrip())
+        raise
     finally:
         if os.path.exists(SOCK):
             os.remove(SOCK)
 
 
 def test():
+    import http.server
     import tempfile
 
     c = merge(DEFAULTS, {"tool_rounds": 1, "env": {"K": "v"}})
@@ -466,28 +678,130 @@ def test():
         assert os.stat(os.path.dirname(path)).st_mode & 0o777 == 0o700
         assert os.stat(path).st_mode & 0o777 == 0o600
 
-    a, b = socket.socketpair()
-    a.setblocking(False)
-    clients.append(a)
-    emit({"type": "ctx", "used": 1, "max": 2})
-    assert json.loads(b.recv(4096)) == {"type": "ctx", "used": 1, "max": 2}
-    emit({"type": "state", "state": "tool", "tool": "search"})
-    assert now_state == {"state": "tool", "tool": "search"}
-    b.recv(4096)
-    b.close()
-    emit({"type": "state", "state": "idle"})             # the dead client is dropped
-    emit({"type": "state", "state": "idle"})
-    assert a not in clients and now_state == {"state": "idle"}
+    # Protocol, against a fake llama-server.
+    global DB, TOOLS, TOOL_SCHEMA
+    script = {"replies": [], "delay": 0}
+
+    class Fake(http.server.BaseHTTPRequestHandler):
+        def log_message(self, *a):
+            pass
+
+        def do_GET(self):
+            body = ({"data": [{"id": "m", "status": {"value": "loaded"}}]}
+                    if self.path.endswith("/models")
+                    else {"default_generation_settings": {"n_ctx": 4096}})
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(json.dumps(body).encode())
+
+        def do_POST(self):
+            self.rfile.read(int(self.headers["Content-Length"]))
+            self.send_response(200)
+            self.end_headers()
+            try:
+                for delta in script["replies"].pop(0):
+                    self.wfile.write(b"data: " + json.dumps({"choices": [{"delta": delta}]})
+                                     .encode() + b"\n\n")
+                    self.wfile.flush()
+                    time.sleep(script["delay"])
+                self.wfile.write(b'data: {"choices": [], "usage": {"prompt_tokens": 42}}\n\n'
+                                 b"data: [DONE]\n\n")
+            except BrokenPipeError:              # a stopped reply hangs up mid-stream
+                pass
+
+    srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Fake)
+    threading.Thread(target=srv.serve_forever, daemon=True).start()
+    CFG["llm"] = f"http://127.0.0.1:{srv.server_address[1]}/v1"
+    DB = db_open(":memory:")
+    TOOLS = {"now": {"description": "d", "command": ["sh", "-c", "echo noon"]}}
+    TOOL_SCHEMA = schema(TOOLS)
     with tempfile.TemporaryDirectory() as d:
         path = os.path.join(d, "s.sock")
         serve(path)
         assert os.stat(path).st_mode & 0o777 == 0o600
-        c = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        c.connect(path)
-        f = c.makefile()
-        first = json.loads(f.readline())
-        assert first == {"type": "hello", "ctx_max": ctx["max"], "state": "idle"}
-        c.close()
+        cl = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        cl.connect(path)
+        cl.settimeout(5)
+        f = cl.makefile("rb")
+
+        def send(**cmd):
+            cl.sendall((json.dumps(cmd) + "\n").encode())
+
+        def until(kind):
+            seen = []
+            while True:
+                e = json.loads(f.readline())
+                seen.append(e)
+                if e["type"] == kind:
+                    return seen
+
+        h = json.loads(f.readline())
+        assert h["type"] == "hello" and h["state"] == "idle" and h["turns"] == []
+        assert h["pending"] is None and "ctx_max" in h
+        # Thinking, a tool round, then the answer.
+        script["replies"] = [
+            [{"reasoning_content": "hmm"},
+             {"tool_calls": [{"index": 0, "id": "c1", "function": {"name": "now",
+                                                                   "arguments": "{}"}}]}],
+            [{"content": "It is "}, {"content": "noon."}]]
+        send(type="ask", text="what time is it")
+        got = [(e["type"], e.get("delta") or e.get("name") or e.get("state") or e.get("text"))
+               for e in until("done") if e["type"] != "ctx"]
+        assert got == [("user", "what time is it"), ("state", "thinking"), ("think", "hmm"),
+                       ("tool", "now"), ("state", "tool"), ("tool", "now"),
+                       ("state", "thinking"), ("token", "It is "), ("token", "noon."),
+                       ("done", None)], got
+        assert until("state")[-1]["state"] == "idle"
+        conv = cur["id"]
+        assert load_conv(DB, conv)[1] == [{"user": "what time is it", "reply": "It is noon.",
+                                           "tools": ["now"]}]
+        # The tool result went to the model but history keeps a placeholder.
+        assert load_conv(DB, conv)[0][0][2] == {"role": "tool", "tool_call_id": "c1",
+                                                "content": "(shown to the user at the time)"}
+        # Stop mid-stream: no token after done, nothing saved.
+        script["replies"] = [[{"content": f"w{i} "} for i in range(50)]]
+        script["delay"] = 0.02
+        send(type="ask", text="long one")
+        until("token")
+        send(type="stop")
+        seen = until("done")
+        assert seen[-1].get("stopped") is True
+        time.sleep(0.3)                              # the reply thread notices and quits
+        send(type="list")
+        late = until("conversations")
+        assert not [e for e in late if e["type"] == "token"], late
+        assert len(load_conv(DB, conv)[1]) == 1
+        script["delay"] = 0
+        # History browsing.
+        send(type="new")
+        assert until("conversation")[-1] == {"type": "conversation", "id": None, "turns": []}
+        send(type="list")
+        assert [x["id"] for x in until("conversations")[-1]["items"]] == [conv]
+        send(type="search", q="tim")
+        assert until("results")[-1]["items"][0]["id"] == conv
+        send(type="open", id=conv)
+        assert until("conversation")[-1]["turns"][0]["reply"] == "It is noon."
+        send(type="open", id=999)
+        assert until("error")[-1]["message"] == "no such conversation"
+        cl.sendall(b"not json\n")
+        assert until("error")[-1]["message"] == "bad command"
+        send(type="delete", id=conv)
+        seen = until("conversations")
+        assert {"type": "conversation", "id": None, "turns": []} in seen   # was current
+        assert seen[-1]["items"] == []
+        # llama-server down: an error, not a dead thread.
+        CFG["llm"] = "http://127.0.0.1:9/v1"
+        send(type="ask", text="anyone?")
+        assert until("error")[-1]["type"] == "error"
+        assert until("state")[-1]["state"] == "idle" and cur["pending"] is None
+        # A second client gets hello with the conversation so far.
+        c2 = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        c2.connect(path)
+        assert json.loads(c2.makefile("rb").readline())["type"] == "hello"
+        c2.close()
+        cl.close()
+    srv.shutdown()
+
     print("ok")