From 8b33e0d213ccfcb2a8d082dac9e6e608e1664a1a Mon Sep 17 00:00:00 2001
From: Pagwin <dev@pagwin.xyz>
Date: Wed, 17 Jun 2026 21:49:13 -0400
Subject: [PATCH] handled early disconnect

---
 main.py  | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 solve.py |  6 +++++-
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/main.py b/main.py
index 9e19851..432931d 100644
--- a/main.py
+++ b/main.py
@@ -14,6 +14,8 @@ from __future__ import annotations
 
 import json
 import os
+import select
+import socket
 import threading
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
@@ -50,6 +52,18 @@ class Handler(BaseHTTPRequestHandler):
         else:
             self._send(404, json.dumps({"error": "not found"}))
 
+    def _client_gone(self):
+        # True once the peer has closed its end. We've already consumed the
+        # request body, so any readable data here is EOF (b"") for a closed tab.
+        sock = self.connection
+        try:
+            r, _, _ = select.select([sock], [], [], 0)
+            if r:
+                return sock.recv(1, socket.MSG_PEEK) == b""
+        except OSError:
+            return True
+        return False
+
     def do_POST(self):
         if self.path != "/solve":
             self._send(404, json.dumps({"error": "not found"}))
@@ -58,17 +72,50 @@ class Handler(BaseHTTPRequestHandler):
         if not _solve_lock.acquire(blocking=False):
             self._send(429, json.dumps({"error": "Another solve is already in progress"}))
             return
+        released = False
         try:
             length = int(self.headers.get("Content-Length", 0))
             payload = json.loads(self.rfile.read(length) or b"{}")
             problem = problem_from_dict(payload.get("problem", {}))
             max_time = float(payload.get("max_time_seconds", 30.0))
-            sol = solve(problem, max_time_seconds=max_time)
-            self._send(200, json.dumps(solution_to_dict(sol)))
+
+            # Run the (blocking) solve in a worker thread so this thread is free
+            # to watch the socket. solver_sink hands us the CpSolver so we can
+            # abort the search if the client disconnects.
+            result = {}
+            solver_box = {}
+
+            def run():
+                try:
+                    result["sol"] = solve(
+                        problem, max_time_seconds=max_time,
+                        solver_sink=lambda s: solver_box.setdefault("s", s),
+                    )
+                except Exception as exc:
+                    result["err"] = exc
+
+            worker = threading.Thread(target=run, daemon=True)
+            worker.start()
+            while worker.is_alive():
+                worker.join(0.25)
+                if worker.is_alive() and self._client_gone():
+                    # Tab closed: stop the search and free the lock immediately
+                    # so the next viewer can start. The worker unwinds on its own.
+                    s = solver_box.get("s")
+                    if s is not None:
+                        s.StopSearch()
+                    _solve_lock.release()
+                    released = True
+                    return
+
+            if "err" in result:
+                raise result["err"]
+            self._send(200, json.dumps(solution_to_dict(result["sol"])))
         except Exception as exc:  # surface errors to the browser
             self._send(400, json.dumps({"error": f"{type(exc).__name__}: {exc}"}))
         finally:
-            _solve_lock.release()
+            if not released:
+                _solve_lock.release()
 
     def log_message(self, fmt, *args):  # quieter console
         pass
diff --git a/solve.py b/solve.py
index 80fb71e..f58a67b 100644
--- a/solve.py
+++ b/solve.py
@@ -1328,12 +1328,16 @@ class _Builder:
 # --------------------------------------------------------------------------- #
 
 def solve(problem: Problem, max_time_seconds: float = 30.0,
-          workers: int = 8) -> Solution:
+          workers: int = 8, solver_sink=None) -> Solution:
     builder = _Builder(problem)
     model = builder.build()
     solver = cp_model.CpSolver()
     solver.parameters.max_time_in_seconds = max_time_seconds
     solver.parameters.num_search_workers = workers
+    # Hand the solver to the caller (if any) so it can call StopSearch() to
+    # abort the search early, e.g. when the requesting client disconnects.
+    if solver_sink is not None:
+        solver_sink(solver)
     status = solver.Solve(model)
     status_name = solver.StatusName(status)