Broadcast assistant events on a Unix socket for the overlay

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

1 files changed, 107 insertions, 2 deletions

diff --git a/assistant.py b/assistant.py
index 7052d69..fa03f2f 100644
--- a/assistant.py
+++ b/assistant.py
@@ -30,12 +30,14 @@ import queue
 import random
 import re
 import signal
+import socket
 import subprocess
 import sys
 import threading
 import time
 import tomllib
 import unicodedata
+import urllib.parse
 import urllib.request
 import wave
 
@@ -60,6 +62,8 @@ DEFAULTS = {
     "voice": {"it": {"lang": "it", "blend": {"if_sara": 0.8, "af_bella": 0.2}},
               "en": {"lang": "en-us", "blend": {"af_heart": 1.0}}},
     "env": {},
+    "overlay": {"name": "Assistant", "position": "center", "hide_delay": 3.0,
+                "user_color": "accent", "assistant_color": "green"},
 }
 OPEN = {"env", "blend"}                   # free-form tables: replaced whole, not merged
 
@@ -113,6 +117,7 @@ TOOLS_DIR = os.path.join(CONFIG_DIR, "tools")
 TOOLS = load_tools(TOOLS_DIR)
 TOOL_SCHEMA = schema(TOOLS) or None
 FRAME = 1 / 30                            # one audio level per frame, seconds
+SOCK = os.path.join(os.environ.get("XDG_RUNTIME_DIR", "/tmp"), "desktop-assistant.sock")
 
 
 def log(*a):
@@ -158,6 +163,72 @@ def notify_ignored(reason):
     notify_id = r.stdout.strip() or notify_id
 
 
+# The overlay's feed: JSON lines on a Unix socket, one event per line. With
+# no client connected every event is dropped, so the voice loop never waits
+# on a UI.
+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
+
+
+def emit(event):
+    if event["type"] == "state":
+        now_state.clear()
+        now_state.update({k: v for k, v in event.items() if k != "type"})
+    line = (json.dumps(event) + "\n").encode()
+    with clients_lock:
+        for c in clients[:]:
+            try:
+                c.sendall(line)           # non-blocking: a stalled client errors
+            except OSError:
+                clients.remove(c)
+                c.close()
+
+
+def hello():
+    return {"type": "hello", **CFG["overlay"], "ctx_max": ctx["max"], **now_state}
+
+
+def serve(path=SOCK):
+    try:
+        os.unlink(path)                   # left over from a crash
+    except FileNotFoundError:
+        pass
+    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    old = os.umask(0o177)                 # 0600: only this user reads the stream
+    try:
+        srv.bind(path)
+    finally:
+        os.umask(old)
+    srv.listen()
+
+    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)
+    threading.Thread(target=accept, daemon=True).start()
+
+
+def context_size(model):
+    # The router can swap models, so ask again whenever the id changes.
+    if ctx["model"] != model:
+        url = f"{CFG['llm'].removesuffix('/v1')}/props?model={urllib.parse.quote(model)}"
+        try:
+            with urllib.request.urlopen(url, timeout=5) as r:
+                ctx["max"] = json.load(r)["default_generation_settings"]["n_ctx"]
+        except (OSError, KeyError, ValueError):
+            ctx["max"] = 0
+        ctx["model"] = model
+    return ctx["max"]
+
+
 def latin(text):
     return all(unicodedata.name(c, "").startswith("LATIN") for c in text if c.isalpha())
 
@@ -296,7 +367,9 @@ def llm_model():
 
 def ask(messages, calls, tools=None):
     # Yields reply text as it streams; tool-call fragments go into calls.
-    body = {"model": llm_model(), "messages": messages, "stream": True,
+    model = llm_model()
+    context_size(model)
+    body = {"model": model, "messages": messages, "stream": True,
             "stream_options": {"include_usage": True},
             "chat_template_kwargs": {"enable_thinking": False}}
     if tools:
@@ -311,7 +384,9 @@ 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.
-                log(f"context: {chunk['usage']['prompt_tokens']} prompt tokens")
+                used = chunk["usage"]["prompt_tokens"]
+                log(f"context: {used} prompt tokens")
+                emit({"type": "ctx", "used": used, "max": ctx["max"]})
             choices = chunk["choices"]
             if not choices:
                 continue
@@ -593,6 +668,36 @@ def test():
     assert len(env) == 30 and max(env) == 0.0             # 1 s at 30 frames/s
     env = envelope(np.full(24000, 0.5, dtype=np.float32), 24000)
     assert all(abs(v - scale(16384)) < 1e-6 for v in env)
+    import socket
+    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"}
+    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" and first["name"] == CFG["overlay"]["name"]
+        assert first["state"] == "idle" and "ctx_max" in first
+        for _ in range(100):                             # accept() appends just after hello
+            if clients:
+                break
+            time.sleep(0.01)
+        emit({"type": "level", "who": "user", "value": 0.5})
+        assert json.loads(f.readline()) == {"type": "level", "who": "user", "value": 0.5}
+        c.close()
     print("ok")