feat: implement session-managed persistent KV cache architecture with slot persistence and management API

docs/SESSION_KV_CACHE.md

@@ -0,0 +1,88 @@
+# Session-Managed KV Cache for llama-server + ollama-proxy
+
+## Goal
+
+Make the multi-GPU llama.cpp backend (exposed via `ollama-proxy.py`) correctly manage KV cache **per session**, so that:
+
+1. Different users / agents / OpenWebUI chats do not pollute each other’s context.
+2. Large shared system prompts (2–30k tokens) used by Maki, OpenWebUI, and other tools are reused instead of being re-prefilled every time.
+3. Session state **survives reboots** (disk-backed).
+
+Frontends:
+- OpenWebUI
+- Maki (https://github.com/wmantly/maki)
+
+Backend stack:
+- `llama-server` (port 8080, `--parallel 1`, `--slot-save-path /var/cache/llama-slots`, `--cache-ram 16384`)
+- `ollama-proxy.py` (port 11434) – Ollama / OpenAI / Anthropic compatible surface with automatic session affinity
+
+---
+
+## Architecture
+
+```text
+OpenWebUI  ─┐
+            ├──► ollama-proxy.py (port 11434) ──► llama-server (port 8080)
+Maki       ─┘         │
+                      └── session_id → /var/cache/llama-slots/<id>.bin
+```
+
+- Single slot (`--parallel 1`) for maximum context and dedicated tensor parallelism across all 3 GPUs.
+- Proxy owns session affinity and decides when to save / restore / erase the slot.
+- Disk persistence via `--slot-save-path /var/cache/llama-slots`.
+- Hot prefix reuse via `--cache-ram 16384` + `cache_prompt: true`.
+
+---
+
+## 1. Session Identity Extraction
+
+The proxy extracts a stable `session_id` on every request according to this priority:
+
+1. **HTTP Headers**:
+   - `X-Session-Id`
+   - `X-Conversation-Id`
+   - `Session-Id`
+   - `Conversation-Id`
+2. **JSON Body Fields**:
+   - `session_id`
+   - `conversation_id`
+   - `chat_id`
+   - `id` (when structured as a chat identifier)
+3. **Fallback**:
+   - `sys-<sha256[:16]>` (hash of the system prompt so identical system prompts share a base KV slot)
+   - `"default"`
+
+---
+
+## 2. Proxy Session Lifecycle
+
+When a request arrives at `ollama-proxy.py`:
+
+```text
+with session_lock:
+    if session_id == current_session:
+        proceed
+
+    # 1. Persist previous session slot
+    if current_session and current_session != "default":
+        POST /slots/0?action=save {"filename": f"{current_session}.bin"}
+
+    # 2. Restore new session or start fresh
+    if os.path.exists(f"/var/cache/llama-slots/{session_id}.bin"):
+        POST /slots/0?action=restore {"filename": f"{session_id}.bin"}
+    else:
+        POST /slots/0?action=erase
+
+    current_session = session_id
+```
+
+---
+
+## 3. Session Management Endpoints
+
+Exposed on `ollama-proxy` (`port 11434`):
+
+* `GET /api/sessions/current`: Returns active `session_id` and slot file status.
+* `GET /api/sessions`: Lists all saved session `.bin` slot files with file sizes and timestamps.
+* `POST /api/sessions/clear`: Forces erase on slot 0 and resets active session to `"default"`.
+* `DELETE /api/sessions/<id>`: Deletes the persisted `.bin` cache file for a specific session.

scripts/ollama-proxy.py

@@ -15,6 +15,7 @@ and forwards to a llama.cpp `llama-server` OpenAI-compatible backend
 Pure stdlib — no third-party deps. Requires Python 3.8+.
 """
 
+import os
 import json
 import sys
 import time
@@ -25,6 +26,7 @@ import http.server
 import threading
 import re
 import subprocess
+import hashlib
 
 # ----------------------------------------------------------------------------
 # Config
@@ -32,6 +34,7 @@ import subprocess
 HOST        = "0.0.0.0"
 PORT        = 11434
 BACKEND     = "http://127.0.0.1:8080/v1"   # llama.cpp OpenAI endpoint
+LLAMA_BASE  = "http://127.0.0.1:8080"      # llama.cpp server root
 MODEL_NAME  = "Qwen3.8-Uncensored"          # primary name
 BACKEND_MODEL = MODEL_NAME                  # what we send llama.cpp
 VERSION     = "0.5.4"                       # fake ollama version
@@ -39,6 +42,90 @@ CTX_SIZE    = 262144
 MAX_OUTPUT  = 131072
 KEEP_ALIVE  = 300
 
+# Session Management & Slot Persistence
+SLOT_SAVE_PATH  = "/var/cache/llama-slots"
+SESSION_LOCK    = threading.Lock()
+CURRENT_SESSION = "default"
+
+def _slot_action(action, filename=None):
+    """Call llama-server /slots/0?action=save|restore|erase."""
+    url = f"{LLAMA_BASE}/slots/0?action={action}"
+    payload = {}
+    if filename:
+        payload["filename"] = filename
+    req = urllib.request.Request(
+        url,
+        data=json.dumps(payload).encode(),
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=30) as r:
+            return json.loads(r.read().decode())
+    except Exception as e:
+        print(f"[ollama-proxy] Slot action {action} (file: {filename}) error: {e}", flush=True)
+        return None
+
+def extract_session_id(headers, payload=None):
+    """Extract a stable session identifier from headers, body, or system prompt."""
+    if headers:
+        for h in ("X-Session-Id", "X-Conversation-Id", "Session-Id", "Conversation-Id", "x-session-id", "x-conversation-id"):
+            val = headers.get(h)
+            if val and str(val).strip():
+                return re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(val).strip())[:64]
+
+    if payload and isinstance(payload, dict):
+        for k in ("session_id", "conversation_id", "chat_id", "id"):
+            val = payload.get(k)
+            if val and isinstance(val, str) and val.strip():
+                return re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(val).strip())[:64]
+
+        # Fallback: hash of system message content
+        messages = payload.get("messages", [])
+        if isinstance(messages, list):
+            for m in messages:
+                if isinstance(m, dict) and m.get("role") == "system":
+                    content = m.get("content") or ""
+                    if content and len(str(content)) > 20:
+                        h = hashlib.sha256(str(content).encode("utf-8", errors="replace")).hexdigest()[:16]
+                        return f"sys-{h}"
+
+    return "default"
+
+def ensure_session(session_id):
+    """Ensure slot 0 contains the KV cache for session_id, saving/restoring as needed."""
+    global CURRENT_SESSION
+    if not session_id:
+        session_id = "default"
+
+    clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(session_id))[:64]
+
+    with SESSION_LOCK:
+        if clean_id == CURRENT_SESSION:
+            return
+
+        # 1. Save old session if not default
+        if CURRENT_SESSION and CURRENT_SESSION != "default":
+            save_file = f"{CURRENT_SESSION}.bin"
+            print(f"[ollama-proxy] Saving slot 0 for session '{CURRENT_SESSION}' -> {save_file}", flush=True)
+            _slot_action("save", save_file)
+
+        # 2. Restore new session or erase
+        target_file = f"{clean_id}.bin"
+        target_path = os.path.join(SLOT_SAVE_PATH, target_file)
+
+        if os.path.exists(target_path):
+            print(f"[ollama-proxy] Restoring slot 0 for session '{clean_id}' <- {target_file}", flush=True)
+            res = _slot_action("restore", target_file)
+            if not res:
+                print(f"[ollama-proxy] Restore failed for '{clean_id}', falling back to erase", flush=True)
+                _slot_action("erase")
+        else:
+            print(f"[ollama-proxy] Starting fresh slot for session '{clean_id}'", flush=True)
+            _slot_action("erase")
+
+        CURRENT_SESSION = clean_id
+
 MODEL_TAGS  = [
     "Qwen3.8-Uncensored:latest",
     "Qwen3.8-Uncensored",
@@ -87,7 +174,7 @@ def _post(path, payload, stream=False):
 
 def _backend_chat(messages, **kw):
     """Call llama.cpp /chat/completions, return parsed JSON (non-streaming)."""
-    body = {"model": BACKEND_MODEL, "messages": messages}
+    body = {"model": BACKEND_MODEL, "messages": messages, "cache_prompt": True}
     body.update(kw)
     body.setdefault("stream", False)
     with _post("/chat/completions", body) as r:
@@ -99,6 +186,7 @@ def _backend_stream(messages, **kw):
         "model": BACKEND_MODEL,
         "messages": messages,
         "stream": True,
+        "cache_prompt": True,
         "stream_options": {"include_usage": True},
     }
     body.update(kw)
@@ -1133,6 +1221,31 @@ class Handler(http.server.BaseHTTPRequestHandler):
             self._send(200, ollama_tags())
         elif path == "/api/ps":
             self._send(200, ollama_ps())
+        elif path == "/api/sessions/current":
+            active_file = os.path.join(SLOT_SAVE_PATH, f"{CURRENT_SESSION}.bin")
+            self._send(200, {
+                "session_id": CURRENT_SESSION,
+                "slot_save_path": SLOT_SAVE_PATH,
+                "filename": f"{CURRENT_SESSION}.bin",
+                "exists_on_disk": os.path.exists(active_file),
+                "file_size_bytes": os.path.getsize(active_file) if os.path.exists(active_file) else 0,
+            })
+        elif path == "/api/sessions":
+            sessions = []
+            if os.path.exists(SLOT_SAVE_PATH):
+                for fname in os.listdir(SLOT_SAVE_PATH):
+                    if fname.endswith(".bin"):
+                        fpath = os.path.join(SLOT_SAVE_PATH, fname)
+                        stat = os.stat(fpath)
+                        sessions.append({
+                            "session_id": fname[:-4],
+                            "filename": fname,
+                            "size_bytes": stat.st_size,
+                            "size_mb": round(stat.st_size / (1024 * 1024), 2),
+                            "modified_iso": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(stat.st_mtime)),
+                            "is_active": (fname[:-4] == CURRENT_SESSION),
+                        })
+            self._send(200, {"sessions": sorted(sessions, key=lambda s: s["size_bytes"], reverse=True)})
         elif path == "/v1/models":
             self._send(200, {
                 "object": "list",
@@ -1164,6 +1277,27 @@ class Handler(http.server.BaseHTTPRequestHandler):
         payload = self._read_json()
 
         try:
+            # Session-Aware Management Routes
+            if path == "/api/sessions/clear":
+                with SESSION_LOCK:
+                    _slot_action("erase")
+                    global CURRENT_SESSION
+                    CURRENT_SESSION = "default"
+                self._send(200, {"status": "cleared", "current_session": "default"})
+                return
+            elif path == "/api/sessions/save":
+                sid = payload.get("session_id") or CURRENT_SESSION
+                clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(sid))[:64]
+                with SESSION_LOCK:
+                    _slot_action("save", f"{clean_id}.bin")
+                self._send(200, {"status": "saved", "session_id": clean_id, "filename": f"{clean_id}.bin"})
+                return
+
+            # Extract & enforce session affinity for all generation and chat endpoints
+            if path in ("/api/chat", "/api/generate", "/v1/messages", "/v1/chat/completions", "/v1/completions", "/v1/responses"):
+                session_id = extract_session_id(self.headers, payload)
+                ensure_session(session_id)
+
             if path == "/api/chat":
                 result, is_stream = ollama_chat(payload)
                 if is_stream:
@@ -1204,6 +1338,7 @@ class Handler(http.server.BaseHTTPRequestHandler):
                 result, _ = anthropic_count_tokens(payload)
                 self._send(200, result)
             elif path in ("/v1/chat/completions", "/v1/completions", "/v1/embeddings", "/v1/responses"):
+                payload.setdefault("cache_prompt", True)
                 body = json.dumps(payload).encode()
                 backend_path = BACKEND + path[len("/v1"):]
                 req = urllib.request.Request(
@@ -1248,8 +1383,24 @@ class Handler(http.server.BaseHTTPRequestHandler):
         except Exception as e:
             self._send(500, {"error": f"{type(e).__name__}: {e}"})
 
+    def do_DELETE(self):
+        path = self.path.split("?")[0]
+        if path.startswith("/api/sessions/"):
+            sid = path[len("/api/sessions/"):].strip()
+            clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', sid)[:64]
+            target = os.path.join(SLOT_SAVE_PATH, f"{clean_id}.bin")
+            if os.path.exists(target):
+                try:
+                    os.remove(target)
+                    self._send(200, {"status": "deleted", "session_id": clean_id})
+                except Exception as e:
+                    self._send(500, {"error": f"Failed to delete session file: {e}"})
+            else:
+                self._send(404, {"error": f"Session '{clean_id}' not found on disk"})
+        else:
+            self.do_POST()
+
     do_PUT = do_POST
-    do_DELETE = do_POST
 
 
 def main():

scripts/start-server.sh

@@ -19,6 +19,8 @@ exec /opt/llama.cpp-xrip/build-nccl/bin/llama-server \
   -ngl 99 \
   -c 262144 \
   --parallel 1 \
+  --slot-save-path /var/cache/llama-slots \
+  --cache-ram 16384 \
   --split-mode tensor \
   --flash-attn on \
   --batch-size 1024 \

systemd/llama-server.service

@@ -19,7 +19,7 @@ TimeoutStopSec=60
 # Basic hardening (server only listens on :8080, reads models read-only)
 ProtectSystem=strict
 ProtectHome=read-only
-ReadWritePaths=
+ReadWritePaths=/var/cache/llama-slots /tmp
 NoNewPrivileges=true
 PrivateTmp=true
 # Do not cap memory/CPU; inference needs all of it