Restore source files from Google Drive backup

llm/server.py

@@ -74,6 +74,39 @@ def warm_model() -> None:
         logger.warning("warm_model failed (non-fatal): {}", e)
 
 
+def clear_kv_cache(hot: bool = True) -> None:
+    """Wipe oMLX's SSD prefix/KV cache (and, by default, the in-memory hot cache).
+
+    oMLX prefix-caches every prompt into an SSD-backed paged KV store so repeated prefixes
+    skip prefill. Within a single agent that's a big win, but across personas the prefixes
+    diverge, so a finished agent's blocks are dead weight that just keeps eating SSD (and the
+    hot cache keeps eating RAM) until the disk fills or the prefill memory guard trips. Call
+    this at each agent boundary to keep peak SSD/RAM use to roughly one agent's context.
+
+    Auth: oMLX's admin endpoints want a session cookie, minted by POSTing the main API key to
+    /admin/api/login; httpx.Client keeps that cookie for the follow-up clears. Fully non-fatal
+    — a cache-clear hiccup must never abort a run.
+    """
+    s = get_settings()
+    base = f"http://{s.llm_host}:{s.llm_port}"
+    try:
+        with httpx.Client(base_url=base, timeout=30) as c:
+            c.post(
+                "/admin/api/login",
+                json={"api_key": s.llm_api_key, "remember": False},
+            ).raise_for_status()
+            ssd = c.post("/admin/api/ssd-cache/clear")
+            ssd.raise_for_status()
+            msg = f"cleared oMLX SSD KV cache ({ssd.json().get('total_deleted', '?')} blocks)"
+            if hot:
+                h = c.post("/admin/api/hot-cache/clear")
+                h.raise_for_status()
+                msg += f", hot cache reclaimed {h.json().get('bytes_reclaimed', 0) / 1e6:.0f}MB"
+            logger.info(msg)
+    except Exception as e:
+        logger.warning("clear_kv_cache failed (non-fatal): {}", e)
+
+
 def ensure_ready(warm: bool = True) -> bool:
     start_server()
     ok = wait_healthy(120)