Hang up on llama-server at once when a reply is stopped

A stop was noticed at the reply's next chunk, and llama-server stops a
task only when its client hangs up. Through a long prompt pass, measured
at 12232 tokens for a file with page images, no chunk comes, so the GPU
kept working after Stop until the first one. interrupt() now shuts the
reply's socket from the stop handler; the request goes through
http.client, which exposes it. A fake server that stays silent records the
hang-up within a second of the stop.

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

1 files changed, 68 insertions, 11 deletions

diff --git a/assistant.py b/assistant.py
index ebc8ba9..8cdf232 100644
--- a/assistant.py
+++ b/assistant.py
@@ -24,6 +24,7 @@ Run: python3 assistant.py        (--test runs the self-check)
 import base64
 import glob
 import hashlib
+import http.client
 import json
 import mimetypes
 import os
@@ -404,6 +405,7 @@ ctx = {"model": None, "max": 0, "used": 0}
 # token after its done.
 lock = threading.RLock()
 cur = {"id": None, "history": [], "turns": [], "pending": None}
+inflight = {"conn": None}   # the HTTP connection the current reply streams over
 turn = 0          # bumped by every ask, stop, new and open: older replies are stale
 DB = None
 
@@ -512,8 +514,17 @@ def serve(path=SOCK):
 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.
+    # Its connection is shut too: llama-server stops only when the client
+    # hangs up, and a reply waiting for its first chunk through a long
+    # prompt pass would otherwise keep the GPU busy until that chunk.
     global turn
     turn += 1
+    conn = inflight["conn"]
+    if conn is not None and conn.sock is not None:
+        try:
+            conn.sock.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
     if cur["pending"] is not None:
         cur["pending"] = None
         emit({"type": "done", "id": cur["id"], "stopped": True})
@@ -667,9 +678,10 @@ def llm_model():
     return model_inputs()[0]
 
 
-def ask(messages, calls, tools=None):
+def ask(messages, calls, tools=None, mine=None):
     # Yields ("think", text) and ("token", text) as they stream; tool-call
-    # fragments go into calls.
+    # fragments go into calls. mine is the reply's turn: a stop shuts the
+    # connection from the other thread (interrupt), which ends the reading.
     model = llm_model()
     context_size(model)
     body = {"model": model, "messages": messages, "stream": True,
@@ -679,14 +691,33 @@ def ask(messages, calls, tools=None):
         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"})
+    # http.client rather than urllib: it exposes the socket a stop must shut.
+    u = urllib.parse.urlsplit(CFG["llm"] + "/chat/completions")
+    conn = (http.client.HTTPSConnection if u.scheme == "https" else http.client.HTTPConnection)(
+        u.hostname, u.port, timeout=300)
+    with lock:
+        if mine is not None and turn != mine:
+            return                        # stopped before the request went out
+        inflight["conn"] = conn
     try:
-        r = urllib.request.urlopen(req, timeout=300)
-    except urllib.error.HTTPError as e:
+        yield from stream(conn, u.path, body, calls)
+    except OSError:
+        if mine is not None and turn != mine:
+            return                        # the stop hung up under the read
+        raise
+    finally:
+        with lock:
+            if inflight["conn"] is conn:
+                inflight["conn"] = None
+        conn.close()
+
+
+def stream(conn, path, body, calls):
+    conn.request("POST", path, json.dumps(body).encode(), {"Content-Type": "application/json"})
+    r = conn.getresponse()
+    if r.status >= 400:
         # The body says why (a full context, a bad request); the status alone does not.
-        raise RuntimeError(f"llama-server {e.code}: "
-                           f"{e.read().decode(errors='replace')[:300]}") from None
+        raise RuntimeError(f"llama-server {r.status}: {r.read().decode(errors='replace')[:300]}")
     with r:
         streamed, other = False, b""
         for line in r:
@@ -796,13 +827,13 @@ def answer(mine, text, files=()):
     # mid-answer would flip from image to gone.
     sent = wire(msgs, images)
     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.
+    # A stop shuts this reply's connection (interrupt), so llama-server's slot
+    # frees even mid prompt pass. ponytail: a running tool still runs out.
     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(sent, calls,
-                               TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
+                               TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None, mine):
             if not say(mine, kind, piece):
                 log("interrupted")
                 return
@@ -1107,6 +1138,16 @@ def test():
 
         def do_POST(self):
             script["body"] = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+            if script.get("hold"):
+                # A long prompt pass: nothing is sent, so only the client
+                # hanging up ends it. Records how soon that happened.
+                import select
+                t0 = time.time()
+                while time.time() - t0 < 5:
+                    if select.select([self.connection], [], [], 0.05)[0]:
+                        if not self.connection.recv(1, socket.MSG_PEEK):
+                            script["hung_up"] = time.time()
+                            return
             self.send_response(200)
             self.end_headers()
             try:
@@ -1188,6 +1229,22 @@ def test():
         assert e["message"] == "context full" and e["turn"] is True
         assert script["body"]["chat_template_kwargs"]["enable_thinking"] is False
         CFG["thinking"] = True
+        # Stop while the server is still reading the prompt: the daemon hangs
+        # up at once rather than at the first chunk, which frees the GPU.
+        script.update(hold=True, hung_up=None)
+        script["replies"] = [[]]
+        send(type="ask", text="a long prompt")
+        until("state")                               # thinking: the request is out
+        time.sleep(0.3)
+        stopped = time.time()
+        send(type="stop")
+        assert until("done")[-1].get("stopped") is True
+        for _ in range(40):
+            if script["hung_up"]:
+                break
+            time.sleep(0.05)
+        assert script["hung_up"] and script["hung_up"] - stopped < 1, script["hung_up"]
+        script["hold"] = False
         # Stop mid-stream: no token after done, nothing saved.
         script["replies"] = [[{"content": f"w{i} "} for i in range(50)]]
         script["delay"] = 0.02