fix: resume Ollama chat jobs across sessions

dashboard/plugin_api.py

@@ -10,6 +10,7 @@ import json
 import mimetypes
 import os
 import re
+import select
 import shutil
 import socket
 import sqlite3
@@ -76,6 +77,8 @@ MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024
 MAX_ATTACHMENT_TEXT = 80_000
 MAX_URL_BYTES = 15 * 1024 * 1024
 CHAT_KEEP_ALIVE = -1
+CHAT_READ_TIMEOUT = 5.0
+CHAT_HEARTBEAT_INTERVAL = 5.0
 HARNESS_MIN_VALIDATORS = 1
 HARNESS_MAX_DRAFT_CHARS = 24_000
 HARNESS_MAX_VALIDATION_CHARS = 8_000
@@ -93,6 +96,7 @@ _chat_db_init_lock = threading.Lock()
 _chat_db_ready = False
 _chat_job_futures: dict[str, Any] = {}
 _chat_job_futures_lock = threading.Lock()
+_chat_job_heartbeat_at: dict[str, float] = {}
 _chat_job_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="ollama-chat-job")
 _DATABASE_URL = os.environ.get("OLLAMA_MANAGER_DATABASE_URL", "").strip()
 
@@ -501,6 +505,16 @@ def _update_chat_job(request_id: str, **values: Any) -> None:
         db.close()
 
 
+def _touch_chat_job(request_id: str, *, force: bool = False) -> None:
+    now = time.time()
+    with _chat_job_futures_lock:
+        previous = _chat_job_heartbeat_at.get(request_id, 0.0)
+        if not force and now - previous < CHAT_HEARTBEAT_INTERVAL:
+            return
+        _chat_job_heartbeat_at[request_id] = now
+    _update_chat_job(request_id, heartbeat_at=now)
+
+
 def _job_status_response(job: dict[str, Any]) -> dict[str, Any]:
     result = json.loads(job.get("result_json") or "{}")
     status = str(job.get("status") or "unknown")
@@ -515,12 +529,37 @@ def _job_status_response(job: dict[str, Any]) -> dict[str, Any]:
         "validator_models": json.loads(job.get("validator_models_json") or "[]"),
         "error": job.get("error") or "",
         "done": status in {"completed", "failed", "stopped", "canceled"},
+        "attempt": job.get("attempt"),
+        "started_at": job.get("started_at"),
+        "heartbeat_at": job.get("heartbeat_at"),
+        "finished_at": job.get("finished_at"),
         "updated_at": job.get("updated_at"),
     }
+    if job.get("heartbeat_at"):
+        response["heartbeat_age"] = round(max(0.0, time.time() - float(job["heartbeat_at"])), 1)
     response.update(result)
     return response
 
 
+def _list_chat_jobs(*, active_only: bool = True, conversation_id: str | None = None, limit: int = 50) -> list[dict[str, Any]]:
+    limit = max(1, min(int(limit), 100))
+    db = _chat_db()
+    try:
+        values: list[Any] = []
+        clauses: list[str] = []
+        if active_only:
+            clauses.append("status IN (?,?,?)")
+            values.extend(["queued", "running", "stopping"])
+        if conversation_id:
+            clauses.append("conversation_id=?")
+            values.append(conversation_id)
+        where = " WHERE " + " AND ".join(clauses) if clauses else ""
+        rows = db.execute(f"SELECT * FROM chat_jobs{where} ORDER BY updated_at DESC LIMIT ?", tuple(values) + (limit,)).fetchall()
+        return [_job_status_response(dict(row)) for row in rows]
+    finally:
+        db.close()
+
+
 def _chat_request_from_job(job: dict[str, Any]) -> ChatRequest:
     return ChatRequest.model_validate(json.loads(job["payload_json"]))
 
@@ -1843,8 +1882,11 @@ def _stream_chat_request(payload: dict[str, Any], request_id: str, cancel_event:
     response_text: list[str] = []
     thinking_text: list[str] = []
     _chat_state(request_id, state="connecting", stage="Connecting to Ollama")
+    _touch_chat_job(parent_id or request_id, force=True)
     try:
         with urlopen(request, timeout=1800) as response:
+            socket_file = getattr(getattr(response, "fp", None), "raw", None)
+            socket_obj = getattr(socket_file, "_sock", None)
             _chat_state(request_id, response=response, state="generating", stage="Ollama is generating")
             while True:
                 with _chat_requests_lock:
@@ -1852,9 +1894,14 @@ def _stream_chat_request(payload: dict[str, Any], request_id: str, cancel_event:
                 if cancelled:
                     response.close()
                     raise _ChatStopped()
+                if socket_obj is not None:
+                    readable, _, _ = select.select([socket_obj], [], [], CHAT_READ_TIMEOUT)
+                    if not readable:
+                        _touch_chat_job(parent_id or request_id)
+                        continue
                 try:
                     raw_line = response.readline()
-                except (socket.timeout, TimeoutError, OSError, ValueError) as exc:
+                except (TimeoutError, OSError, ValueError) as exc:
                     with _chat_requests_lock:
                         cancelled = bool(_chat_requests.get(request_id, {}).get("cancel", threading.Event()).is_set())
                     if cancelled:
@@ -1891,6 +1938,7 @@ def _stream_chat_request(payload: dict[str, Any], request_id: str, cancel_event:
                 )
                 if event.get("done"):
                     break
+                _touch_chat_job(parent_id or request_id)
     except _ChatStopped:
         _chat_state(request_id, state="stopped", stage="Stopped by user", finished_at=time.time())
         raise
@@ -2079,20 +2127,26 @@ def _submit_chat_job(request_id: str) -> None:
 
 
 def _recover_chat_jobs() -> None:
-    try:
-        db = _chat_db()
+    while True:
         try:
-            rows = db.execute("SELECT request_id,status,conversation_id FROM chat_jobs WHERE status IN ('queued','running') ORDER BY created_at").fetchall()
-        finally:
-            db.close()
-        for row in rows:
-            request_id = str(row["request_id"])
-            if str(row["status"]) == "running":
-                _update_chat_job(request_id, status="queued", heartbeat_at=time.time())
-                _persist_chat_event(request_id, row["conversation_id"], "job-recovered", level="warning", stage="Recovered after dashboard restart")
-            _submit_chat_job(request_id)
-    except Exception:
-        return
+            db = _chat_db()
+            try:
+                rows = db.execute("SELECT request_id,status,conversation_id FROM chat_jobs WHERE status IN ('queued','running') ORDER BY created_at").fetchall()
+            finally:
+                db.close()
+            for row in rows:
+                request_id = str(row["request_id"])
+                with _chat_job_futures_lock:
+                    future = _chat_job_futures.get(request_id)
+                    future_missing = future is None or future.done()
+                if str(row["status"]) == "running" and future_missing:
+                    _update_chat_job(request_id, status="queued", heartbeat_at=time.time())
+                    _persist_chat_event(request_id, row["conversation_id"], "job-recovered", level="warning", stage="Recovered by server worker supervisor")
+                if str(row["status"]) == "queued" or future_missing:
+                    _submit_chat_job(request_id)
+        except Exception:
+            pass
+        time.sleep(10)
 
 
 @router.get("/chat/status/{request_id}")
@@ -2113,6 +2167,13 @@ def chat_status(request_id: str) -> dict[str, Any]:
     raise HTTPException(404, "Chat request not found")
 
 
+@router.get("/chat/jobs")
+def chat_jobs(active: bool = True, conversation_id: str | None = None, limit: int = 50) -> dict[str, Any]:
+    if conversation_id:
+        conversation_id = _conversation_id(conversation_id)
+    return {"jobs": _list_chat_jobs(active_only=bool(active), conversation_id=conversation_id, limit=limit)}
+
+
 @router.post("/chat/stop")
 def chat_stop(body: ChatStopRequest) -> dict[str, Any]:
     request_id = _valid_chat_request_id(body.request_id)