diff --git a/chat_logger.py b/chat_logger.py new file mode 100644 index 0000000..3ff3904 --- /dev/null +++ b/chat_logger.py @@ -0,0 +1,433 @@ +"""Chat logging and SSE chunk parsing. + +Extracted from proxy.py — code copied verbatim. +""" +from __future__ import annotations + +import asyncio +import json + +from pathlib import Path + +from log_paths import ( + DATE_FMT, + current_week_dir, + fmt_ts_full, + fmt_ts_short, + local_now, +) + + +RAW_BODY_CAP = 1024 * 1024 + + +class ChatLogger: + """Rotating chat logger — one file per day, bucketed by ISO week folder. + + Reopens when the local date rolls over (which also moves into a new week + folder when needed). Uses local Europe/Zurich timestamps with offset. + """ + + def __init__(self, log_dir: Path) -> None: + self.log_dir = log_dir + self.log_dir.mkdir(parents=True, exist_ok=True) + self._date: str | None = None + self._fh = None + self._raw_fh = None + self._lock = asyncio.Lock() + self._open_for_today() + + def _open_for_today(self) -> None: + date = local_now().strftime(DATE_FMT) + if self._date == date and self._fh is not None: + return + if self._fh is not None: + self._fh.close() + if self._raw_fh is not None: + self._raw_fh.close() + week_dir = current_week_dir(self.log_dir) + self.log_file = week_dir / f"chat-{date}.log" + self.raw_file = week_dir / f"chat-{date}.raw.jsonl" + self._fh = open(self.log_file, "a", encoding="utf-8") + self._raw_fh = open(self.raw_file, "a", encoding="utf-8") + self._date = date + + async def log_request(self, method: str, path: str, body: bytes | None, req_id: str) -> None: + async with self._lock: + self._open_for_today() + ts = fmt_ts_full() + self._fh.write(f"=== [{ts}] [req={req_id}] {method} {path} ===\n") + if body and path.rstrip("/") == "/v1/chat/completions": + self._write_latest_user_turn(body) + self._fh.flush() + self._write_raw(ts, method, path, body, req_id) + + def _write_raw(self, ts: str, method: str, path: str, body: bytes | None, req_id: str) -> None: + record: dict[str, object] = {"ts": ts, "req_id": req_id, "method": method, "path": path} + if body is None: + record["body"] = None + elif len(body) > RAW_BODY_CAP: + record["body"] = None + record["body_truncated"] = body[:RAW_BODY_CAP].decode("utf-8", errors="replace") + record["original_size"] = len(body) + else: + try: + record["body"] = json.loads(body) + except (ValueError, TypeError): + record["body_raw"] = body.decode("utf-8", errors="replace") + self._raw_fh.write(json.dumps(record, ensure_ascii=False) + "\n") + self._raw_fh.flush() + + def _write_latest_user_turn(self, body: bytes) -> None: + try: + payload = json.loads(body) + except (ValueError, TypeError): + return + messages = payload.get("messages") or [] + if not messages: + return + last = messages[-1] + if last.get("role") != "user": + return + text = _stringify_message_content(last.get("content")) + if text: + self._fh.write(f" [user] {text}\n") + + async def log_response(self, data: str, is_done: bool) -> None: + async with self._lock: + self._open_for_today() + ts = fmt_ts_short() + if is_done: + self._fh.write(f" [{ts}] [DONE]\n") + else: + self._fh.write(f" [{ts}] {data}\n") + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + self._fh = None + if self._raw_fh is not None: + self._raw_fh.close() + self._raw_fh = None + + +def _stringify_message_content(content: object) -> str: + if content is None: + return "" + if isinstance(content, str): + return content.replace("\n", " ").strip() + if isinstance(content, list): + parts: list[str] = [] + for item in content: + if not isinstance(item, dict): + continue + ctype = item.get("type") + if ctype == "text" and item.get("text"): + parts.append(str(item["text"]).replace("\n", " ").strip()) + elif ctype in ("image_url", "image"): + parts.append("[image]") + elif ctype == "input_audio": + parts.append("[audio]") + return " ".join(p for p in parts if p) + return str(content) + + +class SSEChunkLogger: + # Markers that can start a tool-call XML block + _START_MARKERS = ("", "", "function": ""} + + def __init__(self, wrapped, chat_logger: ChatLogger) -> None: + self._wrapped = wrapped + self._chat_logger = chat_logger + self._buffer = b"" + # --- existing logging state --- + self._current_kind: str | None = None + self._current_text = "" + self._tool_calls: dict[int, dict[str, str]] = {} + # --- rescue state machine --- + self._rescue_capturing: bool = False + self._rescue_buf: str = "" + self._rescue_kind: str | None = None # "tool_call" or "function" + self._reasoning_holdback: str = "" + self._rescued_any: bool = False + self._rescue_index: int = 0 + # cache upstream chunk id for synthesised events + self._last_chunk_id: str = "rescued" + + async def _flush_text(self) -> None: + if self._current_kind and self._current_text: + await self._chat_logger.log_response( + f"[{self._current_kind}] {self._current_text.strip()}", False + ) + self._current_kind = None + self._current_text = "" + + async def _flush_tool_calls(self) -> None: + if not self._tool_calls: + return + for idx in sorted(self._tool_calls): + tc = self._tool_calls[idx] + name = tc.get("name") or "?" + args = tc.get("arguments") or "" + await self._chat_logger.log_response(f"[tool_call] {name}({args})", False) + self._tool_calls = {} + + async def _flush_all(self) -> None: + await self._flush_text() + await self._flush_tool_calls() + + # ---- rescue helpers ---- + + @staticmethod + def _split_safe_prefix(text: str, markers: tuple[str, ...]) -> tuple[str, str]: + """Return (emit, holdback) where holdback is the longest suffix of *text* + that is a proper prefix of any *marker*.""" + for length in range(len(text), 0, -1): + suffix = text[len(text) - length:] + for marker in markers: + if len(suffix) < len(marker) and marker.startswith(suffix): + return text[: len(text) - length], suffix + return text, "" + + @staticmethod + def _parse_tool_call_xml(block: str) -> dict | None: + """Parse a tool-call XML block. Returns {name, arguments} or None.""" + # Find + import re as _re + m = _re.search(r"\s]+)", block) + if not m: + return None + name: str = m.group(1) + # Find all VALUE + args: dict[str, object] = {} + for pm in _re.finditer(r"(.*?)", block, _re.DOTALL): + key = pm.group(1) + value = pm.group(2).strip() + # Coerce: try JSON parse + try: + value = json.loads(value) + except (json.JSONDecodeError, ValueError): + pass + args[key] = value + return {"name": name, "arguments": args} + + def _build_synthesized_event(self, parsed: dict) -> bytes: + """Build a synthesised tool_calls SSE event from a parsed tool call.""" + import secrets as _secrets + event: dict = { + "id": self._last_chunk_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": self._rescue_index, + "id": f"call_{_secrets.token_hex(4)}", + "type": "function", + "function": { + "name": parsed["name"], + "arguments": json.dumps(parsed["arguments"]), + }, + } + ] + }, + "finish_reason": None, + } + ], + } + self._rescue_index += 1 + body = json.dumps(event, ensure_ascii=False) + return f"data: {body}\r\n\r\n".encode() + + # ---- main read loop ---- + + async def readany(self) -> bytes: + # Loop so we only ever return b"" at true EOF — the consumer treats an + # empty return as end-of-stream. A single upstream chunk may not complete + # an SSE event, in which case _readany_once returns None and we read more. + while True: + out = await self._readany_once() + if out is not None: + return out + + async def _readany_once(self) -> bytes | None: + data = await self._wrapped.content.readany() + if not data: + await self._flush_all() + if self._buffer: + leftover = self._buffer + self._buffer = b"" + return leftover + return b"" + self._buffer += data + outbound: list[bytes] = [] + while True: + crlf_idx = self._buffer.find(b"\r\n\r\n") + lf_idx = self._buffer.find(b"\n\n") + if crlf_idx == -1 and lf_idx == -1: + break + if crlf_idx != -1 and (lf_idx == -1 or crlf_idx <= lf_idx): + idx, sep_len = crlf_idx, 4 + else: + idx, sep_len = lf_idx, 2 + raw_event = self._buffer[: idx + sep_len] + self._buffer = self._buffer[idx + sep_len:] + event_text = raw_event.decode("utf-8", errors="replace").strip() + if not event_text: + outbound.append(raw_event) + continue + # Extract payload line + payload = "" + for line in event_text.splitlines(): + if line.startswith("data:"): + payload = line[5:].strip() + # Non-data lines, comments, [DONE] → pass through + if not payload: + outbound.append(raw_event) + continue + if payload == "[DONE]": + await self._flush_all() + await self._chat_logger.log_response("[DONE]", True) + outbound.append(raw_event) + continue + try: + obj = json.loads(payload) + except json.JSONDecodeError: + outbound.append(raw_event) + continue + # --- (a) existing logging on original delta --- + choices = obj.get("choices") or [] + if choices: + delta = choices[0].get("delta") or {} + reasoning = delta.get("reasoning_content") + content = delta.get("content") + tool_calls = delta.get("tool_calls") + if reasoning: + if self._current_kind != "thinking": + await self._flush_all() + self._current_kind = "thinking" + self._current_text += reasoning + if content: + if self._current_kind != "content": + await self._flush_all() + self._current_kind = "content" + self._current_text += content + if tool_calls: + await self._flush_text() + for tc in tool_calls: + i = tc.get("index", 0) + slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""}) + fn = tc.get("function") or {} + if fn.get("name"): + slot["name"] = fn["name"] + if fn.get("arguments"): + slot["arguments"] += fn["arguments"] + # Track chunk id for synthesised events + cid = obj.get("id") + if cid: + self._last_chunk_id = cid + # --- (b) build outbound bytes with rescue transform --- + outbound_event = self._transform_event(obj) + outbound.append(outbound_event) + return b"".join(outbound) if outbound else None + + def _transform_event(self, obj: dict) -> bytes: + """Transform a single parsed event dict into outbound SSE bytes, + applying the rescue state machine.""" + choices = obj.get("choices") or [] + if not choices: + # No choices — pass through + body = json.dumps(obj, ensure_ascii=False) + return f"data: {body}\r\n\r\n".encode() + + delta = choices[0].get("delta") or {} + reasoning = delta.get("reasoning_content") + + # If no reasoning_content, just rewrite finish_reason if needed + if not reasoning: + if self._rescued_any and choices[0].get("finish_reason") == "stop": + choices[0]["finish_reason"] = "tool_calls" + body = json.dumps(obj, ensure_ascii=False) + return f"data: {body}\r\n\r\n".encode() + + # Run rescue state machine on reasoning_content + work = self._reasoning_holdback + reasoning + self._reasoning_holdback = "" + prose_parts: list[str] = [] + synthesized: list[bytes] = [] + + while work: + if not self._rescue_capturing: + # Look for earliest start marker + earliest_pos = len(work) + earliest_marker: str | None = None + for marker in self._START_MARKERS: + pos = work.find(marker) + if pos != -1 and pos < earliest_pos: + earliest_pos = pos + earliest_marker = marker + + if earliest_marker is None: + # No marker found — apply split_safe_prefix + emit, holdback = self._split_safe_prefix(work, self._START_MARKERS) + prose_parts.append(emit) + self._reasoning_holdback = holdback + work = "" + else: + # A complete start marker is present, so everything before it + # is safe prose — no partial-marker holdback needed here. + prose_parts.append(work[:earliest_pos]) + # Start capturing + self._rescue_capturing = True + self._rescue_kind = ( + "tool_call" if earliest_marker == "" else "function" + ) + self._rescue_buf = earliest_marker + work = work[earliest_pos + len(earliest_marker):] + else: + # Capturing — look for end marker + end_marker = self._END_MARKERS[self._rescue_kind] + end_pos = work.find(end_marker) + if end_pos != -1: + self._rescue_buf += work[: end_pos + len(end_marker)] + # Parse the block + parsed = self._parse_tool_call_xml(self._rescue_buf) + if parsed: + synthesized.append(self._build_synthesized_event(parsed)) + self._rescued_any = True + self._rescue_capturing = False + self._rescue_buf = "" + self._rescue_kind = None + work = work[end_pos + len(end_marker):] + else: + # End marker not found — keep all of work in buffer + self._rescue_buf += work + work = "" + + # Build the outbound event + forwarded_reasoning = "".join(prose_parts) + if forwarded_reasoning: + delta["reasoning_content"] = forwarded_reasoning + else: + delta.pop("reasoning_content", None) + # If delta is now empty and has no other keys, we still emit the event + # (the caller handles skipping if needed) + + # Rewrite finish_reason if rescued + if self._rescued_any and choices[0].get("finish_reason") == "stop": + choices[0]["finish_reason"] = "tool_calls" + + body = json.dumps(obj, ensure_ascii=False) + result = f"data: {body}\r\n\r\n".encode() + # Append any synthesised events after the reasoning event + for syn in synthesized: + result += syn + return result + + def __getattr__(self, name: str) -> object: + return getattr(self._wrapped, name) diff --git a/embed_proxy.py b/embed_proxy.py index 610ae0b..71c9cc8 100644 --- a/embed_proxy.py +++ b/embed_proxy.py @@ -4,9 +4,7 @@ import asyncio import contextlib import logging -import os import secrets -import sys import time from dataclasses import dataclass from pathlib import Path @@ -15,16 +13,15 @@ from log_paths import ( DATE_FMT, - LocalTzFormatter, current_week_dir, - fmt_ts_full, local_now, ) -# Enable ANSI escape sequences on Windows 10+ -if sys.platform == "win32": - os.system("") - +from proxy_base import ( + API_KEY, ProxyConfig, auth_middleware, client_ip, configure_logging, + filter_request_headers, filter_response_headers, health_handler, idle_watchdog, +) +from router_manager import RouterManager ROOT = Path(__file__).resolve().parent SERVER_EXE = ROOT / "llama.cpp_latest" / "llama-server.exe" @@ -45,58 +42,12 @@ IDLE_CHECK_INTERVAL = 30 HEALTH_POLL_INTERVAL = 1.0 BOOT_TIMEOUT = 60 -LOAD_TIMEOUT = 120 RETRY_AFTER_SECONDS = 30 - -def _load_dotenv(path: Path) -> None: - """Load KEY=VALUE pairs from a .env file into os.environ (non-overwrite).""" - if not path.is_file(): - return - for line in path.read_text().splitlines(): - line = line.strip() - if not line or line.startswith("#") or "=" not in line: - continue - key, _, value = line.partition("=") - key, value = key.strip(), value.strip().strip("\"'") - os.environ.setdefault(key, value) - - -_load_dotenv(ROOT / ".env") - -API_KEY = os.environ.get("LLAMA_API_KEY") - -HOP_BY_HOP_HEADERS = { - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailer", - "transfer-encoding", - "upgrade", -} - LOAD_TRIGGER_PATHS = {"/v1/embeddings", "/embeddings", "/v1/rerank", "/rerank"} -@dataclass(frozen=True) -class ProxyConfig: - proxy_host: str - proxy_port: int - server_host: str - server_port: int - idle_timeout: int - idle_check_interval: int - health_poll_interval: float - boot_timeout: int - default_model: str - api_key: str - - @property - def backend_base_url(self) -> str: - return f"http://{self.server_host}:{self.server_port}" - +class EmbedProxyConfig(ProxyConfig): @property def server_command(self) -> list[str]: log_file = current_week_dir(ROOT / "logs") / f"embed-server-{local_now().strftime(DATE_FMT)}.log" @@ -114,158 +65,6 @@ def server_command(self) -> list[str]: ] -class ModelManager: - """Mirror of proxy.py's ModelManager, scoped to the embedding preset.""" - - def __init__(self, config: ProxyConfig, session: ClientSession) -> None: - self.config = config - self.session = session - self.process: asyncio.subprocess.Process | None = None - self._loaded: str | None = None - self._load_lock = asyncio.Lock() - self._active = 0 - self._last_activity = time.monotonic() - - @property - def server_running(self) -> bool: - return self.process is not None and self.process.returncode is None - - @property - def model_loaded(self) -> str | None: - return self._loaded - - @property - def active_requests(self) -> int: - return self._active - - def begin_request(self) -> None: - self._active += 1 - self._last_activity = time.monotonic() - - def end_request(self) -> None: - self._active = max(0, self._active - 1) - self._last_activity = time.monotonic() - - async def start_server(self) -> None: - if self.server_running: - return - logging.info( - "Starting embed router on %s:%s | boot_ts=%s", - self.config.server_host, - self.config.server_port, - fmt_ts_full(), - ) - self.process = await asyncio.create_subprocess_exec( - *self.config.server_command, cwd=str(ROOT), - # Pin to the 3090 Ti (GPU 0); hide the 2070 so layers aren't - # split onto its 8 GB and OOM/slow the embedder. - env={**os.environ, "CUDA_VISIBLE_DEVICES": "0"}, - ) - deadline = time.monotonic() + self.config.boot_timeout - while time.monotonic() < deadline: - if self.process.returncode is not None: - raise RuntimeError(f"router exited during boot: {self.process.returncode}") - try: - async with self.session.get( - f"{self.config.backend_base_url}/health", - timeout=ClientTimeout(total=5), - ) as r: - if r.status == 200: - logging.info("embed router is ready") - return - except (ClientError, asyncio.TimeoutError): - pass - await asyncio.sleep(self.config.health_poll_interval) - raise TimeoutError(f"embed router did not become healthy in {self.config.boot_timeout}s") - - async def stop_server(self) -> None: - if not self.server_running: - return - logging.info("Stopping embed router") - try: - self.process.terminate() - try: - await asyncio.wait_for(self.process.wait(), timeout=5) - except asyncio.TimeoutError: - self.process.kill() - await asyncio.wait_for(self.process.wait(), timeout=5) - except ProcessLookupError: - pass - self.process = None - self._loaded = None - - async def ensure_loaded(self, model: str) -> None: - async with self._load_lock: - if self._loaded == model: - return - current_status = await self._status(model) - if current_status == "loaded": - self._loaded = model - return - logging.info("Loading model: %s", model) - url = f"{self.config.backend_base_url}/models/load" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - async with self.session.post(url, headers=headers, json={"model": model}) as r: - if r.status >= 400: - body = await r.text() - if "already running" in body: - self._loaded = model - return - raise RuntimeError(f"load returned {r.status}: {body}") - deadline = time.monotonic() + LOAD_TIMEOUT - while time.monotonic() < deadline: - status = await self._status(model) - if status == "loaded": - self._loaded = model - logging.info("Model loaded: %s", model) - return - if status == "failed": - raise RuntimeError(f"model {model} failed to load") - await asyncio.sleep(0.5) - raise TimeoutError(f"model {model} did not load in {LOAD_TIMEOUT}s") - - async def unload(self, reason: str) -> None: - if self._loaded is None: - return - model = self._loaded - logging.info("Unloading %s (%s)", model, reason) - url = f"{self.config.backend_base_url}/models/unload" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - try: - async with self.session.post(url, headers=headers, json={"model": model}) as r: - if r.status >= 400: - logging.warning("unload returned %s", r.status) - except ClientError as e: - logging.warning("unload error: %s", e) - self._loaded = None - - async def _status(self, model: str) -> str: - url = f"{self.config.backend_base_url}/v1/models" - headers = {"Authorization": f"Bearer {self.config.api_key}"} - async with self.session.get(url, headers=headers) as r: - data = await r.json() - for entry in data.get("data", []): - if entry.get("id") == model: - status = entry.get("status") or {} - if isinstance(status, dict): - return status.get("value", "unknown") - return str(status) - return "unknown" - - async def unload_if_idle(self) -> None: - if self._active != 0 or self._loaded is None: - return - idle = time.monotonic() - self._last_activity - if idle >= self.config.idle_timeout: - await self.unload(f"idle for {int(idle)}s") - - def write_preset() -> None: content = ( f"[{MODEL_ID}]\n" @@ -280,21 +79,6 @@ def write_preset() -> None: PRESET_PATH.write_text(content, encoding="utf-8") -def configure_logging() -> None: - log_dir = ROOT / "logs" - week_dir = current_week_dir(log_dir) - log_file = week_dir / f"embed-proxy-{local_now().strftime(DATE_FMT)}.log" - fmt = LocalTzFormatter("%(asctime)s %(levelname)s %(message)s") - console = logging.StreamHandler() - console.setLevel(logging.INFO) - console.setFormatter(fmt) - file_handler = logging.FileHandler(log_file, encoding="utf-8") - file_handler.setLevel(logging.INFO) - file_handler.setFormatter(fmt) - logging.basicConfig(level=logging.INFO, handlers=[console, file_handler]) - logging.info("Log file: %s", log_file) - - def build_config() -> ProxyConfig: p = argparse.ArgumentParser(description="Router-mode proxy for the embedding model") p.add_argument("--proxy-host", default=PROXY_HOST) @@ -316,7 +100,7 @@ def build_config() -> ProxyConfig: write_preset() print(f"Embed proxy default: {MODEL_ID} @ {CTX_SIZE // 1024}k ctx") - return ProxyConfig( + return EmbedProxyConfig( proxy_host=args.proxy_host, proxy_port=args.proxy_port, server_host=args.server_host, @@ -330,63 +114,8 @@ def build_config() -> ProxyConfig: ) -def filter_request_headers(headers, api_key: str) -> dict[str, str]: - forwarded: dict[str, str] = {} - for name, value in headers.items(): - lowered = name.lower() - if lowered in ("host", "content-length") or lowered in HOP_BY_HOP_HEADERS: - continue - forwarded[name] = value - if api_key and "Authorization" not in forwarded: - forwarded["Authorization"] = f"Bearer {api_key}" - return forwarded - - -def filter_response_headers(headers) -> dict[str, str]: - forwarded: dict[str, str] = {} - for name, value in headers.items(): - lowered = name.lower() - if lowered == "content-length" or lowered in HOP_BY_HOP_HEADERS: - continue - forwarded[name] = value - return forwarded - - -def client_ip(request: web.Request) -> str: - """Real client IP — CF-Connecting-IP when behind the Cloudflare tunnel, - otherwise the peer address (which is just cloudflared's loopback).""" - return request.headers.get("CF-Connecting-IP") or request.remote or "-" - - -# Paths reachable without an API key. /health is the cloudflared/uptime probe. -PUBLIC_PATHS = {"/health"} - - -@web.middleware -async def auth_middleware(request: web.Request, handler): - """Reject any request that doesn't carry the configured API key. - - This proxy is the internet-facing origin for the Cloudflare tunnel, so - it — not the localhost-only llama-server behind it — is where client - authentication has to happen. filter_request_headers() still injects the - key on the *upstream* hop so the backend keeps trusting only this proxy. - """ - config: ProxyConfig = request.app["config"] - if not config.api_key or request.path.rstrip("/") in PUBLIC_PATHS: - return await handler(request) - header = request.headers.get("Authorization", "") - token = header[7:].strip() if header[:7].lower() == "bearer " else "" - if not token or not secrets.compare_digest(token, config.api_key): - logging.warning( - "401 unauthorized: %s %s from %s", - request.method, request.path, client_ip(request), - ) - return web.json_response({"error": "unauthorized"}, status=401) - return await handler(request) - - async def proxy_request(request: web.Request) -> web.StreamResponse: - manager: ModelManager = request.app["manager"] + manager: RouterManager = request.app["manager"] session: ClientSession = request.app["session"] req_id = secrets.token_hex(4) @@ -461,19 +190,9 @@ async def proxy_request(request: web.Request) -> web.StreamResponse: manager.end_request() -async def health_handler(request: web.Request) -> web.Response: - manager: ModelManager = request.app["manager"] - return web.json_response({ - "status": "ok", - "router": "running" if manager.server_running else "down", - "loaded": manager.model_loaded, - "active_requests": manager.active_requests, - }) - - async def lifecycle_context(app: web.Application): session = ClientSession(timeout=ClientTimeout(total=None)) - manager = ModelManager(app["config"], session) + manager = RouterManager(app["config"], session) app["session"] = session app["manager"] = manager @@ -492,15 +211,6 @@ async def lifecycle_context(app: web.Application): await session.close() -async def idle_watchdog(manager: ModelManager) -> None: - while True: - await asyncio.sleep(manager.config.idle_check_interval) - try: - await manager.unload_if_idle() - except Exception: - logging.exception("idle watchdog error") - - def build_app(config: ProxyConfig) -> web.Application: app = web.Application(client_max_size=32 * 1024 * 1024, middlewares=[auth_middleware]) app["config"] = config @@ -511,7 +221,7 @@ def build_app(config: ProxyConfig) -> web.Application: def main() -> int: - configure_logging() + configure_logging("embed-proxy") config = build_config() logging.info( "Embed proxy %s:%s -> router %s:%s | default=%s | idle=%ss", diff --git a/proxy.py b/proxy.py index 5d69bdf..8799e4f 100644 --- a/proxy.py +++ b/proxy.py @@ -5,28 +5,27 @@ import contextlib import json import logging -import os import secrets -import sys import time from dataclasses import dataclass from pathlib import Path from aiohttp import ClientError, ClientSession, ClientTimeout, web -from aiohttp.web_log import AccessLogger from log_paths import ( DATE_FMT, - LocalTzFormatter, current_week_dir, - fmt_ts_full, - fmt_ts_short, local_now, ) -# Enable ANSI escape sequences on Windows 10+ -if sys.platform == "win32": - os.system("") +from proxy_base import ( + API_KEY, ClientGone, DeadWorkerError, ForwardedAccessLogger, ProxyConfig, + auth_middleware, client_ip, configure_logging, filter_request_headers, + filter_response_headers, health_handler, idle_watchdog, + _is_dead_worker_response, +) +from router_manager import ChatRouterManager +from chat_logger import ChatLogger, SSEChunkLogger ROOT = Path(__file__).resolve().parent @@ -89,104 +88,10 @@ def preset_id(self, ctx: int) -> str: IDLE_CHECK_INTERVAL = 30 # check every 30s HEALTH_POLL_INTERVAL = 1.0 BOOT_TIMEOUT = 60 -LOAD_TIMEOUT = 300 RETRY_AFTER_SECONDS = 30 -MIN_RESIDENCY = 8.0 # minimum seconds a loaded model is kept before an eviction is allowed -# Post-cancel guard: how long to wait for a worker to answer a probe before -# treating it as wedged. Covers a worker still finishing an orphaned prefill; -# if it can't answer in this window we cycle it (the orphan was abandoned anyway). -GUARD_PROBE_TIMEOUT = 30.0 -# Lowercased substrings that identify a dead-worker 500 from the router. -# The router returns these when its HTTP client can't reach the worker child. -DEAD_WORKER_MARKERS: tuple[str, ...] = ( - "could not establish connection", - "failed to read connection", - "failed to write connection", - "http client error", -) - - -def _load_dotenv(path: Path) -> None: - """Load KEY=VALUE pairs from a .env file into os.environ (non-overwrite).""" - if not path.is_file(): - return - for line in path.read_text().splitlines(): - line = line.strip() - if not line or line.startswith("#") or "=" not in line: - continue - key, _, value = line.partition("=") - key, value = key.strip(), value.strip().strip("\"'") - os.environ.setdefault(key, value) - - -_load_dotenv(ROOT / ".env") - -API_KEY = os.environ.get("LLAMA_API_KEY") - -HOP_BY_HOP_HEADERS = { - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailer", - "transfer-encoding", - "upgrade", -} - - -class DeadWorkerError(Exception): - """Raised when the router returns a 500 whose body indicates the worker - child has died. Caught by proxy_request to trigger a recovery + retry.""" - - def __init__(self, status: int, body: bytes) -> None: - self.status = status - self.body = body - super().__init__(f"dead-worker {status}: {body[:200]!r}") - - -class ClientGone(Exception): - """Raised when the *downstream* client hangs up before we finish sending the - response (e.g. a connection reset during downstream.prepare()). Benign — it is - not a proxy fault, so it is logged calmly and returned as 499, unlike an - *upstream* reset from the router which is a real bad gateway (502).""" - - -def _is_dead_worker_response(status: int, body: bytes) -> bool: - """Return True when *status* ≥ 500 and *body* contains a dead-worker marker.""" - if status < 500: - return False - lowered = body.lower() - return any(m.encode() in lowered for m in DEAD_WORKER_MARKERS) -@dataclass(frozen=True) -class ProxyConfig: - proxy_host: str - proxy_port: int - server_host: str - server_port: int - idle_timeout: int - idle_check_interval: int - health_poll_interval: float - boot_timeout: int - default_model: str - api_key: str - embed_host: str - embed_port: int - chat_log: bool = True - - @property - def backend_base_url(self) -> str: - return f"http://{self.server_host}:{self.server_port}" - - @property - def embed_base_url(self) -> str: - host = self.embed_host - if ":" in host: - host = f"[{host}]" # IPv6 needs brackets in URLs - return f"http://{host}:{self.embed_port}" - +class ChatProxyConfig(ProxyConfig): @property def server_command(self) -> list[str]: log_file = current_week_dir(ROOT / "logs") / f"llama-server-{local_now().strftime(DATE_FMT)}.log" @@ -210,925 +115,6 @@ def server_command(self) -> list[str]: ] -class ModelManager: - """Owns the long-lived router process and on-demand model load/unload. - - The router process starts with the proxy and dies with it. Models are - loaded on first request and unloaded by the idle watchdog — the router - itself stays up so the cloudflared tunnel never breaks. - """ - - def __init__(self, config: ProxyConfig, session: ClientSession) -> None: - self.config = config - self.session = session - self.process: asyncio.subprocess.Process | None = None - self._loaded: str | None = None # alias of currently-loaded model, None if nothing - self._load_lock = asyncio.Lock() - self._active = 0 - self._last_activity = time.monotonic() - self._forwarding: int = 0 # count of requests currently streaming - self._idle_forward: asyncio.Event = asyncio.Event() - self._idle_forward.set() # set == 0 in-flight (idle) - self._loaded_at: float = 0.0 # monotonic time the current model finished loading - # Set when a request is aborted mid-generation (client cancel/disconnect). - # llama-server has an unfixed cancel→next-request desync that wedges/crashes - # the worker (ggml-org/llama.cpp#20921), so the NEXT model request probes the - # worker first instead of detonating on it. Guarded by _guard_lock. - self._worker_suspect: bool = False - self._guard_lock = asyncio.Lock() - - @property - def server_running(self) -> bool: - return self.process is not None and self.process.returncode is None - - @property - def model_loaded(self) -> str | None: - return self._loaded - - @property - def active_requests(self) -> int: - return self._active - - def begin_request(self) -> None: - self._active += 1 - self._last_activity = time.monotonic() - - def end_request(self) -> None: - self._active = max(0, self._active - 1) - self._last_activity = time.monotonic() - - def _begin_forward(self) -> None: - """Register a new in-flight streaming request.""" - self._forwarding += 1 - self._idle_forward.clear() # not idle while a request is streaming - - def _end_forward(self) -> None: - """Deregister a completed streaming request.""" - self._forwarding = max(0, self._forwarding - 1) - if self._forwarding == 0: - self._idle_forward.set() # signal idle to any waiting switcher - - async def start_server(self) -> None: - if self.server_running: - return - logging.info( - "Starting router on %s:%s | boot_ts=%s", - self.config.server_host, - self.config.server_port, - fmt_ts_full(), - ) - self.process = await asyncio.create_subprocess_exec( - *self.config.server_command, cwd=str(ROOT), - # Pin to the 3090 Ti (GPU 0); hide the 2070 so layers aren't - # split onto its 8 GB and OOM/slow the big models. - env={**os.environ, "CUDA_VISIBLE_DEVICES": "0"}, - ) - deadline = time.monotonic() + self.config.boot_timeout - while time.monotonic() < deadline: - if self.process.returncode is not None: - raise RuntimeError(f"router exited during boot: {self.process.returncode}") - try: - async with self.session.get( - f"{self.config.backend_base_url}/health", - timeout=ClientTimeout(total=5), - ) as r: - if r.status == 200: - logging.info("router is ready") - return - except (ClientError, asyncio.TimeoutError): - pass - await asyncio.sleep(self.config.health_poll_interval) - raise TimeoutError(f"router did not become healthy in {self.config.boot_timeout}s") - - async def stop_server(self) -> None: - if not self.server_running: - return - logging.info("Stopping router") - try: - self.process.terminate() - try: - await asyncio.wait_for(self.process.wait(), timeout=5) - except asyncio.TimeoutError: - self.process.kill() - await asyncio.wait_for(self.process.wait(), timeout=5) - except ProcessLookupError: - pass - self.process = None - self._loaded = None - - async def _switch_and_load_locked(self, model: str) -> None: - """Load *model* into the router. Caller MUST already hold ``_load_lock``. - - Syncs cached ``_loaded`` state from the router before issuing a load - to avoid spurious "already running" 400s, then polls until the model - reports ``loaded`` or a timeout expires. Sets ``_loaded_at`` on every - successful load so the residency window starts fresh. - """ - if self._loaded == model: - return - # Sync state: the router may already have this model loaded (e.g. a - # direct /models/load call through the proxy). Check before issuing - # another load — otherwise the router returns 400 "already running". - current_status = await self._status(model) - if current_status == "loaded": - self._loaded = model - self._loaded_at = time.monotonic() - return - logging.info("Loading model: %s", model) - url = f"{self.config.backend_base_url}/models/load" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - async with self.session.post(url, headers=headers, json={"model": model}) as r: - if r.status >= 400: - body = await r.text() - if "already running" in body: - self._loaded = model - self._loaded_at = time.monotonic() - return - raise RuntimeError(f"load returned {r.status}: {body}") - deadline = time.monotonic() + LOAD_TIMEOUT - while time.monotonic() < deadline: - status = await self._status(model) - if status == "loaded": - self._loaded = model - self._loaded_at = time.monotonic() - logging.info("Model loaded: %s", model) - return - if status == "failed": - raise RuntimeError(f"model {model} failed to load") - await asyncio.sleep(0.5) - raise TimeoutError(f"model {model} did not load in {LOAD_TIMEOUT}s") - - async def ensure_loaded(self, model: str) -> None: - """Ensure *model* is loaded. Acquires ``_load_lock`` internally. - - Used by explicit ``/models/load`` proxy pass-through so that a direct - client load request goes through the same serialisation path. - """ - async with self._load_lock: - await self._switch_and_load_locked(model) - - @contextlib.asynccontextmanager - async def use_model(self, model: str): - """Async context manager that ensures *model* is loaded for the duration - of a streaming forward, preventing mid-stream eviction. - - Acquiring ``_load_lock`` for a switch: - 1. DRAIN — waits for all currently-streaming requests to finish before - evicting the old model (``_idle_forward`` is only set when - ``_forwarding == 0``). - 2. MIN-RESIDENCY — after a fresh load, waits out the remainder of - ``MIN_RESIDENCY`` seconds before allowing another eviction, preventing - rapid ping-pong reloads by concurrent agents on different models. - - ``_begin_forward()`` is called while the lock is still held so the - counter increment is atomic with respect to any concurrent switcher. - """ - async with self._load_lock: - if self._loaded != model: - # DRAIN: never evict a model that has an in-flight request. - # _begin_forward() only runs under _load_lock, so no new - # forward can sneak in while we are deciding to switch. - if self._loaded is not None: - await self._idle_forward.wait() - # MIN-RESIDENCY: don't evict a model loaded less than - # MIN_RESIDENCY seconds ago; wait out the remainder. - wait = MIN_RESIDENCY - (time.monotonic() - self._loaded_at) - if wait > 0: - await asyncio.sleep(wait) - await self._switch_and_load_locked(model) - self._begin_forward() # register UNDER the lock — atomic vs a switch - try: - yield - finally: - self._end_forward() - - async def unload(self, reason: str) -> None: - if self._loaded is None: - return - model = self._loaded - logging.info("Unloading %s (%s)", model, reason) - url = f"{self.config.backend_base_url}/models/unload" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - try: - async with self.session.post(url, headers=headers, json={"model": model}) as r: - if r.status >= 400: - logging.warning("unload returned %s", r.status) - except ClientError as e: - logging.warning("unload error: %s", e) - self._loaded = None - - async def recover_worker(self, model: str, dead_detected_at: float) -> bool: - """Force-cycle the router worker for *model* after a dead-worker 500. - - Acquires ``_load_lock`` so concurrent callers serialise. Only the - first caller actually cycles unload/load; subsequent callers whose - ``dead_detected_at`` is before a recent ``_loaded_at`` know a peer - already completed recovery and return True immediately. - - NOTE: We do NOT check the router's ``_status()`` to skip the cycle — - the router can report ``loaded`` while the worker child is already - dead. We always force an unload+reload on the first caller. - - Returns True when the worker is confirmed loaded, False on timeout/ - error so the caller can return 503. - """ - async with self._load_lock: - # Re-check: if _loaded_at was updated AFTER the dead-worker was - # detected, a peer coroutine already completed recovery. - if self._loaded == model and self._loaded_at > dead_detected_at: - logging.info( - "[recover_worker] %s already recovered by peer (loaded_at=%.3f > detected=%.3f)", - model, self._loaded_at, dead_detected_at, - ) - return True - - logging.warning( - "[recover_worker] cycling unload/load for dead worker %s", model - ) - self._loaded = None # invalidate proxy cache immediately - - # Unload — force the router to tear down the dead worker entry. - # Best-effort: the router may error if it can't contact the child, - # but we continue to the load step regardless. - url_unload = f"{self.config.backend_base_url}/models/unload" - auth_headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - try: - async with self.session.post( - url_unload, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400: - logging.warning("[recover_worker] unload returned %s: %s", r.status, body_text) - else: - logging.info("[recover_worker] unload OK for %s", model) - except (ClientError, asyncio.TimeoutError) as e: - logging.warning("[recover_worker] unload error (ignored): %s", e) - - # Wait for the router to reflect the unloaded state so that - # _switch_and_load_locked doesn't see "loaded" and short-circuit. - unload_deadline = time.monotonic() + 15 - while time.monotonic() < unload_deadline: - try: - status = await self._status(model) - if status != "loaded": - logging.info("[recover_worker] router confirms %s is %s", model, status) - break - except Exception: - pass - await asyncio.sleep(0.5) - else: - logging.warning( - "[recover_worker] router still shows loaded after 15s — proceeding anyway" - ) - - # Reload — spawn fresh worker(s), draining any buffered exit - # signals the router queued during unload. Each exit signal - # consumes exactly one fresh worker; we loop until the worker - # stays alive for a full second after reporting "loaded". - # - # We do NOT use _switch_and_load_locked here because its poll - # loops until "loaded" is observed — if a queued exit fires - # between the worker reporting ready and our next poll (< 500 ms), - # the status flips to "unloaded" and the loop spins for 300 s. - # Instead we implement a custom load + rapid-poll that detects - # the loaded→unloaded flash and retries the load immediately. - load_url = f"{self.config.backend_base_url}/models/load" - max_load_attempts = 4 - per_load_timeout = 120 # seconds to wait for status != "loading" - - for load_attempt in range(1, max_load_attempts + 1): - logging.info( - "[recover_worker] load attempt %d/%d for %s", - load_attempt, max_load_attempts, model, - ) - try: - async with self.session.post( - load_url, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400 and "already running" not in body_text: - logging.warning("[recover_worker] load returned %s: %s", r.status, body_text) - except (ClientError, asyncio.TimeoutError) as e: - logging.error("[recover_worker] load POST failed: %s", e) - return False - - # Poll until status leaves "loading" (either loaded or unloaded) - poll_deadline = time.monotonic() + per_load_timeout - prev_status = "" - seen_loading = False # have we observed this attempt actually start? - while time.monotonic() < poll_deadline: - try: - status = await self._status(model) - except Exception: - await asyncio.sleep(0.5) - continue - if status != prev_status: - logging.info("[recover_worker] %s status: %s", model, status) - prev_status = status - if status == "loading": - seen_loading = True - if status == "loaded": - # Give the router 1.5s to process any pending exit signal - # before declaring victory. - await asyncio.sleep(1.5) - status2 = await self._status(model) - if status2 == "loaded": - self._loaded = model - self._loaded_at = time.monotonic() - logging.info( - "[recover_worker] worker stable for %s (attempt %d)", - model, load_attempt, - ) - return True - logging.warning( - "[recover_worker] fresh worker for %s exited immediately " - "(status after 1.5s: %s, attempt %d/%d)", - model, status2, load_attempt, max_load_attempts, - ) - break # queued exit consumed — try next load attempt - if status == "failed": - logging.error("[recover_worker] model %s failed to load", model) - return False - if status == "unloaded" and seen_loading: - # The worker started loading but was then stopped before ever - # reaching "loaded" — a queued exit signal consumed it (the - # router force-kills it after its ~10s stop timeout) or it - # crashed mid-load. For a slow-loading model the kill lands - # before "loaded" is ever observed, so the loaded→flash case - # above never fires. Retry the next load immediately instead - # of spinning here for the full per_load_timeout. - logging.warning( - "[recover_worker] fresh worker for %s died during load " - "(status: unloaded, attempt %d/%d) — retrying", - model, load_attempt, max_load_attempts, - ) - break # queued exit consumed — try next load attempt - await asyncio.sleep(0.25) - else: - logging.error( - "[recover_worker] timed out waiting for %s to load (attempt %d/%d)", - model, load_attempt, max_load_attempts, - ) - # Don't return False here — let the for-loop try the next attempt. - # The router may need a fresh /models/load POST to clear a stuck - # "loading" state from a previous crash. - # Retry the unload to give the router a chance to tear down - # any zombie worker entry from the timed-out attempt. - if load_attempt < max_load_attempts: - logging.info( - "[recover_worker] retrying unload before next load attempt" - ) - try: - async with self.session.post( - url_unload, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400: - logging.warning( - "[recover_worker] retry-unload returned %s: %s", - r.status, body_text, - ) - else: - logging.info( - "[recover_worker] retry-unload OK for %s", model - ) - except (ClientError, asyncio.TimeoutError) as e: - logging.warning( - "[recover_worker] retry-unload error (ignored): %s", e - ) - # Brief pause to let the router settle before the next load. - await asyncio.sleep(2) - continue - - logging.error("[recover_worker] all %d load attempts failed for %s", max_load_attempts, model) - return False - - def mark_worker_suspect(self) -> None: - """Flag the worker as possibly desynced after a mid-generation client - abort (cancel/disconnect). The next model request will probe before use. - See _worker_suspect for the upstream bug this guards against.""" - self._worker_suspect = True - - async def _probe_worker(self, model: str) -> bool: - """Send a minimal generation to detect a wedged/crashed worker after a - cancel. Returns True if the worker answers normally; False if it returns - a dead-worker error, errors out, or times out (a busy-finishing-an-orphan - or genuinely-wedged worker both warrant a recovery cycle).""" - url = f"{self.config.backend_base_url}/v1/chat/completions" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - payload = { - "model": model, - "messages": [{"role": "user", "content": "ping"}], - "max_tokens": 1, - "stream": False, - } - try: - async with self.session.post( - url, headers=headers, json=payload, - timeout=ClientTimeout(total=GUARD_PROBE_TIMEOUT), - ) as r: - body = await r.read() - if _is_dead_worker_response(r.status, body): - return False - return r.status < 500 - except (ClientError, asyncio.TimeoutError) as e: - logging.warning("[guard] probe error for %s: %s", model, e) - return False - - async def guard_after_cancel(self, model: str) -> None: - """If a prior request was aborted mid-generation, probe the worker once - before the next request uses it. Recover proactively if the probe shows - it wedged/dead — so the next real request lands on a clean worker instead - of triggering the upstream cancel→next-request crash.""" - if not self._worker_suspect: - return - async with self._guard_lock: - if not self._worker_suspect: - return # a concurrent caller already handled this episode - # Only the currently-loaded worker can be the wedged one. If a different - # model (or nothing) is loaded, use_model will spawn a fresh worker - # anyway, so there is nothing to probe — just clear the flag. - if self._loaded != model: - self._worker_suspect = False - return - logging.info("[guard] worker suspect after cancel — probing %s", model) - healthy = await self._probe_worker(model) - if healthy: - logging.info("[guard] worker %s healthy after cancel", model) - else: - logging.warning( - "[guard] worker %s wedged/dead after cancel — recovering", model - ) - await self.recover_worker(model, time.monotonic()) - self._worker_suspect = False - - async def _status(self, model: str) -> str: - url = f"{self.config.backend_base_url}/v1/models" - headers = {"Authorization": f"Bearer {self.config.api_key}"} - async with self.session.get(url, headers=headers) as r: - data = await r.json() - for entry in data.get("data", []): - if entry.get("id") == model: - status = entry.get("status") or {} - if isinstance(status, dict): - return status.get("value", "unknown") - return str(status) - return "unknown" - - async def unload_if_idle(self) -> None: - # Never unload while a streaming forward is in progress. - if self._active != 0 or self._forwarding != 0 or self._loaded is None: - return - idle = time.monotonic() - self._last_activity - if idle >= self.config.idle_timeout: - await self.unload(f"idle for {int(idle)}s") - - -RAW_BODY_CAP = 1024 * 1024 - - -class ChatLogger: - """Rotating chat logger — one file per day, bucketed by ISO week folder. - - Reopens when the local date rolls over (which also moves into a new week - folder when needed). Uses local Europe/Zurich timestamps with offset. - """ - - def __init__(self, log_dir: Path) -> None: - self.log_dir = log_dir - self.log_dir.mkdir(parents=True, exist_ok=True) - self._date: str | None = None - self._fh = None - self._raw_fh = None - self._lock = asyncio.Lock() - self._open_for_today() - - def _open_for_today(self) -> None: - date = local_now().strftime(DATE_FMT) - if self._date == date and self._fh is not None: - return - if self._fh is not None: - self._fh.close() - if self._raw_fh is not None: - self._raw_fh.close() - week_dir = current_week_dir(self.log_dir) - self.log_file = week_dir / f"chat-{date}.log" - self.raw_file = week_dir / f"chat-{date}.raw.jsonl" - self._fh = open(self.log_file, "a", encoding="utf-8") - self._raw_fh = open(self.raw_file, "a", encoding="utf-8") - self._date = date - - async def log_request(self, method: str, path: str, body: bytes | None, req_id: str) -> None: - async with self._lock: - self._open_for_today() - ts = fmt_ts_full() - self._fh.write(f"=== [{ts}] [req={req_id}] {method} {path} ===\n") - if body and path.rstrip("/") == "/v1/chat/completions": - self._write_latest_user_turn(body) - self._fh.flush() - self._write_raw(ts, method, path, body, req_id) - - def _write_raw(self, ts: str, method: str, path: str, body: bytes | None, req_id: str) -> None: - record: dict[str, object] = {"ts": ts, "req_id": req_id, "method": method, "path": path} - if body is None: - record["body"] = None - elif len(body) > RAW_BODY_CAP: - record["body"] = None - record["body_truncated"] = body[:RAW_BODY_CAP].decode("utf-8", errors="replace") - record["original_size"] = len(body) - else: - try: - record["body"] = json.loads(body) - except (ValueError, TypeError): - record["body_raw"] = body.decode("utf-8", errors="replace") - self._raw_fh.write(json.dumps(record, ensure_ascii=False) + "\n") - self._raw_fh.flush() - - def _write_latest_user_turn(self, body: bytes) -> None: - try: - payload = json.loads(body) - except (ValueError, TypeError): - return - messages = payload.get("messages") or [] - if not messages: - return - last = messages[-1] - if last.get("role") != "user": - return - text = _stringify_message_content(last.get("content")) - if text: - self._fh.write(f" [user] {text}\n") - - async def log_response(self, data: str, is_done: bool) -> None: - async with self._lock: - self._open_for_today() - ts = fmt_ts_short() - if is_done: - self._fh.write(f" [{ts}] [DONE]\n") - else: - self._fh.write(f" [{ts}] {data}\n") - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - self._fh = None - if self._raw_fh is not None: - self._raw_fh.close() - self._raw_fh = None - - -def _stringify_message_content(content: object) -> str: - if content is None: - return "" - if isinstance(content, str): - return content.replace("\n", " ").strip() - if isinstance(content, list): - parts: list[str] = [] - for item in content: - if not isinstance(item, dict): - continue - ctype = item.get("type") - if ctype == "text" and item.get("text"): - parts.append(str(item["text"]).replace("\n", " ").strip()) - elif ctype in ("image_url", "image"): - parts.append("[image]") - elif ctype == "input_audio": - parts.append("[audio]") - return " ".join(p for p in parts if p) - return str(content) - - -class SSEChunkLogger: - # Markers that can start a tool-call XML block - _START_MARKERS = ("", "", "function": ""} - - def __init__(self, wrapped, chat_logger: ChatLogger) -> None: - self._wrapped = wrapped - self._chat_logger = chat_logger - self._buffer = b"" - # --- existing logging state --- - self._current_kind: str | None = None - self._current_text = "" - self._tool_calls: dict[int, dict[str, str]] = {} - # --- rescue state machine --- - self._rescue_capturing: bool = False - self._rescue_buf: str = "" - self._rescue_kind: str | None = None # "tool_call" or "function" - self._reasoning_holdback: str = "" - self._rescued_any: bool = False - self._rescue_index: int = 0 - # cache upstream chunk id for synthesised events - self._last_chunk_id: str = "rescued" - - async def _flush_text(self) -> None: - if self._current_kind and self._current_text: - await self._chat_logger.log_response( - f"[{self._current_kind}] {self._current_text.strip()}", False - ) - self._current_kind = None - self._current_text = "" - - async def _flush_tool_calls(self) -> None: - if not self._tool_calls: - return - for idx in sorted(self._tool_calls): - tc = self._tool_calls[idx] - name = tc.get("name") or "?" - args = tc.get("arguments") or "" - await self._chat_logger.log_response(f"[tool_call] {name}({args})", False) - self._tool_calls = {} - - async def _flush_all(self) -> None: - await self._flush_text() - await self._flush_tool_calls() - - # ---- rescue helpers ---- - - @staticmethod - def _split_safe_prefix(text: str, markers: tuple[str, ...]) -> tuple[str, str]: - """Return (emit, holdback) where holdback is the longest suffix of *text* - that is a proper prefix of any *marker*.""" - for length in range(len(text), 0, -1): - suffix = text[len(text) - length:] - for marker in markers: - if len(suffix) < len(marker) and marker.startswith(suffix): - return text[: len(text) - length], suffix - return text, "" - - @staticmethod - def _parse_tool_call_xml(block: str) -> dict | None: - """Parse a tool-call XML block. Returns {name, arguments} or None.""" - # Find - import re as _re - m = _re.search(r"\s]+)", block) - if not m: - return None - name: str = m.group(1) - # Find all VALUE - args: dict[str, object] = {} - for pm in _re.finditer(r"(.*?)", block, _re.DOTALL): - key = pm.group(1) - value = pm.group(2).strip() - # Coerce: try JSON parse - try: - value = json.loads(value) - except (json.JSONDecodeError, ValueError): - pass - args[key] = value - return {"name": name, "arguments": args} - - def _build_synthesized_event(self, parsed: dict) -> bytes: - """Build a synthesised tool_calls SSE event from a parsed tool call.""" - import secrets as _secrets - event: dict = { - "id": self._last_chunk_id, - "object": "chat.completion.chunk", - "choices": [ - { - "index": 0, - "delta": { - "tool_calls": [ - { - "index": self._rescue_index, - "id": f"call_{_secrets.token_hex(4)}", - "type": "function", - "function": { - "name": parsed["name"], - "arguments": json.dumps(parsed["arguments"]), - }, - } - ] - }, - "finish_reason": None, - } - ], - } - self._rescue_index += 1 - body = json.dumps(event, ensure_ascii=False) - return f"data: {body}\r\n\r\n".encode() - - # ---- main read loop ---- - - async def readany(self) -> bytes: - # Loop so we only ever return b"" at true EOF — the consumer treats an - # empty return as end-of-stream. A single upstream chunk may not complete - # an SSE event, in which case _readany_once returns None and we read more. - while True: - out = await self._readany_once() - if out is not None: - return out - - async def _readany_once(self) -> bytes | None: - data = await self._wrapped.content.readany() - if not data: - await self._flush_all() - if self._buffer: - leftover = self._buffer - self._buffer = b"" - return leftover - return b"" - self._buffer += data - outbound: list[bytes] = [] - while True: - crlf_idx = self._buffer.find(b"\r\n\r\n") - lf_idx = self._buffer.find(b"\n\n") - if crlf_idx == -1 and lf_idx == -1: - break - if crlf_idx != -1 and (lf_idx == -1 or crlf_idx <= lf_idx): - idx, sep_len = crlf_idx, 4 - else: - idx, sep_len = lf_idx, 2 - raw_event = self._buffer[: idx + sep_len] - self._buffer = self._buffer[idx + sep_len:] - event_text = raw_event.decode("utf-8", errors="replace").strip() - if not event_text: - outbound.append(raw_event) - continue - # Extract payload line - payload = "" - for line in event_text.splitlines(): - if line.startswith("data:"): - payload = line[5:].strip() - # Non-data lines, comments, [DONE] → pass through - if not payload: - outbound.append(raw_event) - continue - if payload == "[DONE]": - await self._flush_all() - await self._chat_logger.log_response("[DONE]", True) - outbound.append(raw_event) - continue - try: - obj = json.loads(payload) - except json.JSONDecodeError: - outbound.append(raw_event) - continue - # --- (a) existing logging on original delta --- - choices = obj.get("choices") or [] - if choices: - delta = choices[0].get("delta") or {} - reasoning = delta.get("reasoning_content") - content = delta.get("content") - tool_calls = delta.get("tool_calls") - if reasoning: - if self._current_kind != "thinking": - await self._flush_all() - self._current_kind = "thinking" - self._current_text += reasoning - if content: - if self._current_kind != "content": - await self._flush_all() - self._current_kind = "content" - self._current_text += content - if tool_calls: - await self._flush_text() - for tc in tool_calls: - i = tc.get("index", 0) - slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""}) - fn = tc.get("function") or {} - if fn.get("name"): - slot["name"] = fn["name"] - if fn.get("arguments"): - slot["arguments"] += fn["arguments"] - # Track chunk id for synthesised events - cid = obj.get("id") - if cid: - self._last_chunk_id = cid - # --- (b) build outbound bytes with rescue transform --- - outbound_event = self._transform_event(obj) - outbound.append(outbound_event) - return b"".join(outbound) if outbound else None - - def _transform_event(self, obj: dict) -> bytes: - """Transform a single parsed event dict into outbound SSE bytes, - applying the rescue state machine.""" - choices = obj.get("choices") or [] - if not choices: - # No choices — pass through - body = json.dumps(obj, ensure_ascii=False) - return f"data: {body}\r\n\r\n".encode() - - delta = choices[0].get("delta") or {} - reasoning = delta.get("reasoning_content") - - # If no reasoning_content, just rewrite finish_reason if needed - if not reasoning: - if self._rescued_any and choices[0].get("finish_reason") == "stop": - choices[0]["finish_reason"] = "tool_calls" - body = json.dumps(obj, ensure_ascii=False) - return f"data: {body}\r\n\r\n".encode() - - # Run rescue state machine on reasoning_content - work = self._reasoning_holdback + reasoning - self._reasoning_holdback = "" - prose_parts: list[str] = [] - synthesized: list[bytes] = [] - - while work: - if not self._rescue_capturing: - # Look for earliest start marker - earliest_pos = len(work) - earliest_marker: str | None = None - for marker in self._START_MARKERS: - pos = work.find(marker) - if pos != -1 and pos < earliest_pos: - earliest_pos = pos - earliest_marker = marker - - if earliest_marker is None: - # No marker found — apply split_safe_prefix - emit, holdback = self._split_safe_prefix(work, self._START_MARKERS) - prose_parts.append(emit) - self._reasoning_holdback = holdback - work = "" - else: - # A complete start marker is present, so everything before it - # is safe prose — no partial-marker holdback needed here. - prose_parts.append(work[:earliest_pos]) - # Start capturing - self._rescue_capturing = True - self._rescue_kind = ( - "tool_call" if earliest_marker == "" else "function" - ) - self._rescue_buf = earliest_marker - work = work[earliest_pos + len(earliest_marker):] - else: - # Capturing — look for end marker - end_marker = self._END_MARKERS[self._rescue_kind] - end_pos = work.find(end_marker) - if end_pos != -1: - self._rescue_buf += work[: end_pos + len(end_marker)] - # Parse the block - parsed = self._parse_tool_call_xml(self._rescue_buf) - if parsed: - synthesized.append(self._build_synthesized_event(parsed)) - self._rescued_any = True - self._rescue_capturing = False - self._rescue_buf = "" - self._rescue_kind = None - work = work[end_pos + len(end_marker):] - else: - # End marker not found — keep all of work in buffer - self._rescue_buf += work - work = "" - - # Build the outbound event - forwarded_reasoning = "".join(prose_parts) - if forwarded_reasoning: - delta["reasoning_content"] = forwarded_reasoning - else: - delta.pop("reasoning_content", None) - # If delta is now empty and has no other keys, we still emit the event - # (the caller handles skipping if needed) - - # Rewrite finish_reason if rescued - if self._rescued_any and choices[0].get("finish_reason") == "stop": - choices[0]["finish_reason"] = "tool_calls" - - body = json.dumps(obj, ensure_ascii=False) - result = f"data: {body}\r\n\r\n".encode() - # Append any synthesised events after the reasoning event - for syn in synthesized: - result += syn - return result - - def __getattr__(self, name: str) -> object: - return getattr(self._wrapped, name) - - -def configure_logging() -> None: - log_dir = ROOT / "logs" - week_dir = current_week_dir(log_dir) - log_file = week_dir / f"proxy-{local_now().strftime(DATE_FMT)}.log" - fmt = LocalTzFormatter("%(asctime)s %(levelname)s %(message)s") - console = logging.StreamHandler() - console.setLevel(logging.INFO) - console.setFormatter(fmt) - file_handler = logging.FileHandler(log_file, encoding="utf-8") - file_handler.setLevel(logging.INFO) - file_handler.setFormatter(fmt) - logging.basicConfig(level=logging.INFO, handlers=[console, file_handler]) - logging.info("Log file: %s", log_file) - - DEFAULT_MODEL: ModelChoice = MODELS[0] DEFAULT_CTX: int = CTX_CHOICES[-1] @@ -1233,7 +219,7 @@ def build_config() -> ProxyConfig: print(f"Default: {model.label} @ {ctx // 1024}k ctx (id: {default_id})") print(f"Exposed presets: {len(MODELS) * len(CTX_CHOICES)} (one per model×ctx combo)") - return ProxyConfig( + return ChatProxyConfig( proxy_host=args.proxy_host, proxy_port=args.proxy_port, server_host=args.server_host, @@ -1250,28 +236,6 @@ def build_config() -> ProxyConfig: ) -def filter_request_headers(headers, api_key: str) -> dict[str, str]: - forwarded: dict[str, str] = {} - for name, value in headers.items(): - lowered = name.lower() - if lowered in ("host", "content-length") or lowered in HOP_BY_HOP_HEADERS: - continue - forwarded[name] = value - if api_key and "Authorization" not in forwarded: - forwarded["Authorization"] = f"Bearer {api_key}" - return forwarded - - -def filter_response_headers(headers) -> dict[str, str]: - forwarded: dict[str, str] = {} - for name, value in headers.items(): - lowered = name.lower() - if lowered == "content-length" or lowered in HOP_BY_HOP_HEADERS: - continue - forwarded[name] = value - return forwarded - - def _strip_chat_prefix(path: str) -> str: """Strip the /chat alias prefix so backend sees plain /v1/... paths.""" if path == "/chat": @@ -1281,67 +245,6 @@ def _strip_chat_prefix(path: str) -> str: return path -# Reverse-proxy hops allowed to set forwarding headers. We listen on -# 0.0.0.0:8001, so a direct LAN client could spoof X-Forwarded-For / -# CF-Connecting-IP; only trust those headers when the TCP peer is Caddy (LAN -# front) or cloudflared (loopback). Override the Caddy IP via $TRUSTED_PROXY. -TRUSTED_PROXIES = {"127.0.0.1", "::1", os.environ.get("TRUSTED_PROXY", "192.168.178.43")} - - -def client_ip(request: web.Request) -> str: - """Real client IP. When the request arrives from a trusted reverse proxy - (Caddy on the LAN, or cloudflared on loopback) we read the forwarded client - out of CF-Connecting-IP / X-Forwarded-For; otherwise we report the raw TCP - peer. Untrusted peers can't spoof their way to a fake IP.""" - peer = request.remote or "-" - if peer in TRUSTED_PROXIES: - forwarded = ( - request.headers.get("CF-Connecting-IP") - or request.headers.get("X-Forwarded-For", "").split(",")[0].strip() - ) - if forwarded: - return forwarded - return peer - - -class ForwardedAccessLogger(AccessLogger): - """aiohttp access logger that resolves the ``%a`` atom through client_ip(), - so the access line shows the real client instead of the reverse-proxy hop.""" - - @staticmethod - def _format_a(request, response, time): - if request is None: - return "-" - return client_ip(request) - - -# Paths reachable without an API key. /health is the cloudflared/uptime probe. -PUBLIC_PATHS = {"/health"} - - -@web.middleware -async def auth_middleware(request: web.Request, handler): - """Reject any request that doesn't carry the configured API key. - - This proxy is the internet-facing origin for the Cloudflare tunnel, so - it — not the localhost-only llama-server behind it — is where client - authentication has to happen. filter_request_headers() still injects the - key on the *upstream* hop so the backend keeps trusting only this proxy. - """ - config: ProxyConfig = request.app["config"] - if not config.api_key or request.path.rstrip("/") in PUBLIC_PATHS: - return await handler(request) - header = request.headers.get("Authorization", "") - token = header[7:].strip() if header[:7].lower() == "bearer " else "" - if not token or not secrets.compare_digest(token, config.api_key): - logging.warning( - "401 unauthorized: %s %s from %s", - request.method, request.path, client_ip(request), - ) - return web.json_response({"error": "unauthorized"}, status=401) - return await handler(request) - - def _inject_cache_prompt(body: bytes | None, method: str, path: str) -> bytes | None: if method != "POST" or path.rstrip("/") != "/v1/chat/completions" or not body: return body @@ -1369,7 +272,7 @@ def _model_from_body(body: bytes | None, fallback: str) -> str: async def proxy_request(request: web.Request) -> web.StreamResponse: - manager: ModelManager = request.app["manager"] + manager: ChatRouterManager = request.app["manager"] session: ClientSession = request.app["session"] chat_logger: ChatLogger | None = request.app.get("chat_logger") @@ -1619,7 +522,7 @@ async def models_handler(request: web.Request) -> web.Response: "Retry" in pi. We rewrite the flag to false whenever `value` says the preset is simply unloaded — `value` is the source of truth. """ - manager: ModelManager = request.app["manager"] + manager: ChatRouterManager = request.app["manager"] session: ClientSession = request.app["session"] effective_path = _strip_chat_prefix(request.path) query = request.rel_url.query_string @@ -1653,7 +556,7 @@ async def props_handler(request: web.Request) -> web.Response: instead of "Load & switch"). We rewrite to a clean 200 JSON whose shape matches the exact equality checks in baseModel.getStatus(). """ - manager: ModelManager = request.app["manager"] + manager: ChatRouterManager = request.app["manager"] session: ClientSession = request.app["session"] effective_path = _strip_chat_prefix(request.path) query = request.rel_url.query_string @@ -1742,19 +645,9 @@ async def embed_forward(request: web.Request) -> web.StreamResponse: ) -async def health_handler(request: web.Request) -> web.Response: - manager: ModelManager = request.app["manager"] - return web.json_response({ - "status": "ok", - "router": "running" if manager.server_running else "down", - "loaded": manager.model_loaded, - "active_requests": manager.active_requests, - }) - - async def lifecycle_context(app: web.Application): session = ClientSession(timeout=ClientTimeout(total=None)) - manager = ModelManager(app["config"], session) + manager = ChatRouterManager(app["config"], session) chat_logger = ChatLogger(ROOT / "logs") if app["config"].chat_log else None app["session"] = session @@ -1777,15 +670,6 @@ async def lifecycle_context(app: web.Application): chat_logger.close() -async def idle_watchdog(manager: ModelManager) -> None: - while True: - await asyncio.sleep(manager.config.idle_check_interval) - try: - await manager.unload_if_idle() - except Exception: - logging.exception("idle watchdog error") - - def build_app(config: ProxyConfig) -> web.Application: app = web.Application(client_max_size=128 * 1024 * 1024, middlewares=[auth_middleware]) app["config"] = config diff --git a/proxy_base.py b/proxy_base.py new file mode 100644 index 0000000..eff63fd --- /dev/null +++ b/proxy_base.py @@ -0,0 +1,246 @@ +"""Shared leaf code for proxy.py and embed_proxy.py. + +Mechanical extraction — code copied verbatim from proxy.py (canonical source). +A later phase rewrites proxy.py and embed_proxy.py to import from here. +""" +from __future__ import annotations + +import asyncio +import logging +import os +import secrets +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from aiohttp import ClientError, ClientSession, ClientTimeout, web +from aiohttp.web_log import AccessLogger + +from log_paths import ( + DATE_FMT, + LocalTzFormatter, + current_week_dir, + local_now, +) + +if TYPE_CHECKING: + from router_manager import RouterManager + +# Enable ANSI escape sequences on Windows 10+ +if sys.platform == "win32": + os.system("") + + +ROOT = Path(__file__).resolve().parent + + +def _load_dotenv(path: Path) -> None: + """Load KEY=VALUE pairs from a .env file into os.environ (non-overwrite).""" + if not path.is_file(): + return + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + key, value = key.strip(), value.strip().strip("\"'") + os.environ.setdefault(key, value) + + +_load_dotenv(ROOT / ".env") + +API_KEY = os.environ.get("LLAMA_API_KEY") + +HOP_BY_HOP_HEADERS = { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + + +class DeadWorkerError(Exception): + """Raised when the router returns a 500 whose body indicates the worker + child has died. Caught by proxy_request to trigger a recovery + retry.""" + + def __init__(self, status: int, body: bytes) -> None: + self.status = status + self.body = body + super().__init__(f"dead-worker {status}: {body[:200]!r}") + + +class ClientGone(Exception): + """Raised when the *downstream* client hangs up before we finish sending the + response (e.g. a connection reset during downstream.prepare()). Benign — it is + not a proxy fault, so it is logged calmly and returned as 499, unlike an + *upstream* reset from the router which is a real bad gateway (502).""" + + +# Lowercased substrings that identify a dead-worker 500 from the router. +# The router returns these when its HTTP client can't reach the worker child. +DEAD_WORKER_MARKERS: tuple[str, ...] = ( + "could not establish connection", + "failed to read connection", + "failed to write connection", + "http client error", +) + + +def _is_dead_worker_response(status: int, body: bytes) -> bool: + """Return True when *status* ≥ 500 and *body* contains a dead-worker marker.""" + if status < 500: + return False + lowered = body.lower() + return any(m.encode() in lowered for m in DEAD_WORKER_MARKERS) + + +@dataclass(frozen=True) +class ProxyConfig: + proxy_host: str + proxy_port: int + server_host: str + server_port: int + idle_timeout: int + idle_check_interval: int + health_poll_interval: float + boot_timeout: int + default_model: str + api_key: str + embed_host: str | None = None + embed_port: int | None = None + chat_log: bool = True + + @property + def backend_base_url(self) -> str: + return f"http://{self.server_host}:{self.server_port}" + + @property + def embed_base_url(self) -> str: + if self.embed_host is None or self.embed_port is None: + raise AttributeError("embed_host/embed_port not set on this config") + host = self.embed_host + if ":" in host: + host = f"[{host}]" # IPv6 needs brackets in URLs + return f"http://{host}:{self.embed_port}" + +def filter_request_headers(headers, api_key: str) -> dict[str, str]: + forwarded: dict[str, str] = {} + for name, value in headers.items(): + lowered = name.lower() + if lowered in ("host", "content-length") or lowered in HOP_BY_HOP_HEADERS: + continue + forwarded[name] = value + if api_key and "Authorization" not in forwarded: + forwarded["Authorization"] = f"Bearer {api_key}" + return forwarded + + +def filter_response_headers(headers) -> dict[str, str]: + forwarded: dict[str, str] = {} + for name, value in headers.items(): + lowered = name.lower() + if lowered == "content-length" or lowered in HOP_BY_HOP_HEADERS: + continue + forwarded[name] = value + return forwarded + + +# Reverse-proxy hops allowed to set forwarding headers. We listen on +# 0.0.0.0:8001, so a direct LAN client could spoof X-Forwarded-For / +# CF-Connecting-IP; only trust those headers when the TCP peer is Caddy (LAN +# front) or cloudflared (loopback). Override the Caddy IP via $TRUSTED_PROXY. +TRUSTED_PROXIES = {"127.0.0.1", "::1", os.environ.get("TRUSTED_PROXY", "192.168.178.43")} + + +def client_ip(request: web.Request) -> str: + """Real client IP. When the request arrives from a trusted reverse proxy + (Caddy on the LAN, or cloudflared on loopback) we read the forwarded client + out of CF-Connecting-IP / X-Forwarded-For; otherwise we report the raw TCP + peer. Untrusted peers can't spoof their way to a fake IP.""" + peer = request.remote or "-" + if peer in TRUSTED_PROXIES: + forwarded = ( + request.headers.get("CF-Connecting-IP") + or request.headers.get("X-Forwarded-For", "").split(",")[0].strip() + ) + if forwarded: + return forwarded + return peer + + +class ForwardedAccessLogger(AccessLogger): + """aiohttp access logger that resolves the ``%a`` atom through client_ip(), + so the access line shows the real client instead of the reverse-proxy hop.""" + + @staticmethod + def _format_a(request, response, time): + if request is None: + return "-" + return client_ip(request) + + +# Paths reachable without an API key. /health is the cloudflared/uptime probe. +PUBLIC_PATHS = {"/health"} + + +@web.middleware +async def auth_middleware(request: web.Request, handler): + """Reject any request that doesn't carry the configured API key. + + This proxy is the internet-facing origin for the Cloudflare tunnel, so + it — not the localhost-only llama-server behind it — is where client + authentication has to happen. filter_request_headers() still injects the + key on the *upstream* hop so the backend keeps trusting only this proxy. + """ + config: ProxyConfig = request.app["config"] + if not config.api_key or request.path.rstrip("/") in PUBLIC_PATHS: + return await handler(request) + header = request.headers.get("Authorization", "") + token = header[7:].strip() if header[:7].lower() == "bearer " else "" + if not token or not secrets.compare_digest(token, config.api_key): + logging.warning( + "401 unauthorized: %s %s from %s", + request.method, request.path, client_ip(request), + ) + return web.json_response({"error": "unauthorized"}, status=401) + return await handler(request) + + +def configure_logging(log_stem: str = "proxy") -> None: + log_dir = ROOT / "logs" + week_dir = current_week_dir(log_dir) + log_file = week_dir / f"{log_stem}-{local_now().strftime(DATE_FMT)}.log" + fmt = LocalTzFormatter("%(asctime)s %(levelname)s %(message)s") + console = logging.StreamHandler() + console.setLevel(logging.INFO) + console.setFormatter(fmt) + file_handler = logging.FileHandler(log_file, encoding="utf-8") + file_handler.setLevel(logging.INFO) + file_handler.setFormatter(fmt) + logging.basicConfig(level=logging.INFO, handlers=[console, file_handler]) + logging.info("Log file: %s", log_file) + + +async def health_handler(request: web.Request) -> web.Response: + manager: RouterManager = request.app["manager"] + return web.json_response({ + "status": "ok", + "router": "running" if manager.server_running else "down", + "loaded": manager.model_loaded, + "active_requests": manager.active_requests, + }) + + +async def idle_watchdog(manager: RouterManager) -> None: + while True: + await asyncio.sleep(manager.config.idle_check_interval) + try: + await manager.unload_if_idle() + except Exception: + logging.exception("idle watchdog error") diff --git a/router_manager.py b/router_manager.py new file mode 100644 index 0000000..000eb4c --- /dev/null +++ b/router_manager.py @@ -0,0 +1,570 @@ +"""Shared router lifecycle management for proxy.py and embed_proxy.py. + +Mechanical extraction from proxy.py (ChatRouterManager) and embed_proxy.py +(RouterManager base). Behaviour-preserving — no logic changes. +""" +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import time + +from aiohttp import ClientError, ClientSession, ClientTimeout + +from log_paths import fmt_ts_full +from proxy_base import ( + ProxyConfig, + ROOT, + _is_dead_worker_response, +) + +# ── Module constants (chat-only — used only by ChatRouterManager methods) ── + +MIN_RESIDENCY = 8.0 # minimum seconds a loaded model is kept before an eviction is allowed +# Post-cancel guard: how long to wait for a worker to answer a probe before +# treating it as wedged. Covers a worker still finishing an orphaned prefill; +# if it can't answer in this window we cycle it (the orphan was abandoned anyway). +GUARD_PROBE_TIMEOUT = 30.0 + + +# ── Base class ────────────────────────────────────────────────────────────── + +class RouterManager: + """Mirror of proxy.py's ModelManager, scoped to the embedding preset.""" + + LOAD_TIMEOUT = 120 + ROUTER_LABEL = "embed router" # base IS the embed manager; ChatRouterManager overrides to "router" + + def __init__(self, config: ProxyConfig, session: ClientSession) -> None: + self.config = config + self.session = session + self.process: asyncio.subprocess.Process | None = None + self._loaded: str | None = None + self._load_lock = asyncio.Lock() + self._active = 0 + self._last_activity = time.monotonic() + + @property + def server_running(self) -> bool: + return self.process is not None and self.process.returncode is None + + @property + def model_loaded(self) -> str | None: + return self._loaded + + @property + def active_requests(self) -> int: + return self._active + + def begin_request(self) -> None: + self._active += 1 + self._last_activity = time.monotonic() + + def end_request(self) -> None: + self._active = max(0, self._active - 1) + self._last_activity = time.monotonic() + + async def start_server(self) -> None: + if self.server_running: + return + logging.info( + "Starting %s on %s:%s | boot_ts=%s", + self.ROUTER_LABEL, + self.config.server_host, + self.config.server_port, + fmt_ts_full(), + ) + self.process = await asyncio.create_subprocess_exec( + *self.config.server_command, cwd=str(ROOT), + # Pin to the 3090 Ti (GPU 0); hide the 2070 so layers aren't + # split onto its 8 GB and OOM/slow the embedder. + env={**os.environ, "CUDA_VISIBLE_DEVICES": "0"}, + ) + deadline = time.monotonic() + self.config.boot_timeout + while time.monotonic() < deadline: + if self.process.returncode is not None: + raise RuntimeError(f"router exited during boot: {self.process.returncode}") + try: + async with self.session.get( + f"{self.config.backend_base_url}/health", + timeout=ClientTimeout(total=5), + ) as r: + if r.status == 200: + logging.info("%s is ready", self.ROUTER_LABEL) + return + except (ClientError, asyncio.TimeoutError): + pass + await asyncio.sleep(self.config.health_poll_interval) + raise TimeoutError(f"{self.ROUTER_LABEL} did not become healthy in {self.config.boot_timeout}s") + + async def stop_server(self) -> None: + if not self.server_running: + return + logging.info("Stopping %s", self.ROUTER_LABEL) + try: + self.process.terminate() + try: + await asyncio.wait_for(self.process.wait(), timeout=5) + except asyncio.TimeoutError: + self.process.kill() + await asyncio.wait_for(self.process.wait(), timeout=5) + except ProcessLookupError: + pass + self.process = None + self._loaded = None + + async def ensure_loaded(self, model: str) -> None: + async with self._load_lock: + await self._load_locked(model) + + async def _load_locked(self, model: str) -> None: + if self._loaded == model: + return + current_status = await self._status(model) + if current_status == "loaded": + self._loaded = model + return + logging.info("Loading model: %s", model) + url = f"{self.config.backend_base_url}/models/load" + headers = { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + } + async with self.session.post(url, headers=headers, json={"model": model}) as r: + if r.status >= 400: + body = await r.text() + if "already running" in body: + self._loaded = model + return + raise RuntimeError(f"load returned {r.status}: {body}") + deadline = time.monotonic() + self.LOAD_TIMEOUT + while time.monotonic() < deadline: + status = await self._status(model) + if status == "loaded": + self._loaded = model + logging.info("Model loaded: %s", model) + return + if status == "failed": + raise RuntimeError(f"model {model} failed to load") + await asyncio.sleep(0.5) + raise TimeoutError(f"model {model} did not load in {self.LOAD_TIMEOUT}s") + + async def unload(self, reason: str) -> None: + if self._loaded is None: + return + model = self._loaded + logging.info("Unloading %s (%s)", model, reason) + url = f"{self.config.backend_base_url}/models/unload" + headers = { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + } + try: + async with self.session.post(url, headers=headers, json={"model": model}) as r: + if r.status >= 400: + logging.warning("unload returned %s", r.status) + except ClientError as e: + logging.warning("unload error: %s", e) + self._loaded = None + + async def _status(self, model: str) -> str: + url = f"{self.config.backend_base_url}/v1/models" + headers = {"Authorization": f"Bearer {self.config.api_key}"} + async with self.session.get(url, headers=headers) as r: + data = await r.json() + for entry in data.get("data", []): + if entry.get("id") == model: + status = entry.get("status") or {} + if isinstance(status, dict): + return status.get("value", "unknown") + return str(status) + return "unknown" + + async def unload_if_idle(self) -> None: + if self._active != 0 or self._loaded is None: + return + idle = time.monotonic() - self._last_activity + if idle >= self.config.idle_timeout: + await self.unload(f"idle for {int(idle)}s") + + +# ── Chat subclass ─────────────────────────────────────────────────────────── + +class ChatRouterManager(RouterManager): + """Owns the long-lived router process and on-demand model load/unload. + + The router process starts with the proxy and dies with it. Models are + loaded on first request and unloaded by the idle watchdog — the router + itself stays up so the cloudflared tunnel never breaks. + """ + + LOAD_TIMEOUT = 300 + ROUTER_LABEL = "router" + + def __init__(self, config: ProxyConfig, session: ClientSession) -> None: + super().__init__(config, session) + self._forwarding: int = 0 # count of requests currently streaming + self._idle_forward: asyncio.Event = asyncio.Event() + self._idle_forward.set() # set == 0 in-flight (idle) + self._loaded_at: float = 0.0 # monotonic time the current model finished loading + # Set when a request is aborted mid-generation (client cancel/disconnect). + # llama-server has an unfixed cancel→next-request desync that wedges/crashes + # the worker (ggml-org/llama.cpp#20921), so the NEXT model request probes the + # worker first instead of detonating on it. Guarded by _guard_lock. + self._worker_suspect: bool = False + self._guard_lock = asyncio.Lock() + + def _begin_forward(self) -> None: + """Register a new in-flight streaming request.""" + self._forwarding += 1 + self._idle_forward.clear() # not idle while a request is streaming + + def _end_forward(self) -> None: + """Deregister a completed streaming request.""" + self._forwarding = max(0, self._forwarding - 1) + if self._forwarding == 0: + self._idle_forward.set() # signal idle to any waiting switcher + + async def _load_locked(self, model: str) -> None: + """Load *model* into the router. Caller MUST already hold ``_load_lock``. + + Syncs cached ``_loaded`` state from the router before issuing a load + to avoid spurious "already running" 400s, then polls until the model + reports ``loaded`` or a timeout expires. Sets ``_loaded_at`` on every + successful load so the residency window starts fresh. + """ + if self._loaded == model: + return + # Sync state: the router may already have this model loaded (e.g. a + # direct /models/load call through the proxy). Check before issuing + # another load — otherwise the router returns 400 "already running". + current_status = await self._status(model) + if current_status == "loaded": + self._loaded = model + self._loaded_at = time.monotonic() + return + logging.info("Loading model: %s", model) + url = f"{self.config.backend_base_url}/models/load" + headers = { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + } + async with self.session.post(url, headers=headers, json={"model": model}) as r: + if r.status >= 400: + body = await r.text() + if "already running" in body: + self._loaded = model + self._loaded_at = time.monotonic() + return + raise RuntimeError(f"load returned {r.status}: {body}") + deadline = time.monotonic() + self.LOAD_TIMEOUT + while time.monotonic() < deadline: + status = await self._status(model) + if status == "loaded": + self._loaded = model + self._loaded_at = time.monotonic() + logging.info("Model loaded: %s", model) + return + if status == "failed": + raise RuntimeError(f"model {model} failed to load") + await asyncio.sleep(0.5) + raise TimeoutError(f"model {model} did not load in {self.LOAD_TIMEOUT}s") + + @contextlib.asynccontextmanager + async def use_model(self, model: str): + """Async context manager that ensures *model* is loaded for the duration + of a streaming forward, preventing mid-stream eviction. + + Acquiring ``_load_lock`` for a switch: + 1. DRAIN — waits for all currently-streaming requests to finish before + evicting the old model (``_idle_forward`` is only set when + ``_forwarding == 0``). + 2. MIN-RESIDENCY — after a fresh load, waits out the remainder of + ``MIN_RESIDENCY`` seconds before allowing another eviction, preventing + rapid ping-pong reloads by concurrent agents on different models. + + ``_begin_forward()`` is called while the lock is still held so the + counter increment is atomic with respect to any concurrent switcher. + """ + async with self._load_lock: + if self._loaded != model: + # DRAIN: never evict a model that has an in-flight request. + # _begin_forward() only runs under _load_lock, so no new + # forward can sneak in while we are deciding to switch. + if self._loaded is not None: + await self._idle_forward.wait() + # MIN-RESIDENCY: don't evict a model loaded less than + # MIN_RESIDENCY seconds ago; wait out the remainder. + wait = MIN_RESIDENCY - (time.monotonic() - self._loaded_at) + if wait > 0: + await asyncio.sleep(wait) + await self._load_locked(model) + self._begin_forward() # register UNDER the lock — atomic vs a switch + try: + yield + finally: + self._end_forward() + + async def recover_worker(self, model: str, dead_detected_at: float) -> bool: + """Force-cycle the router worker for *model* after a dead-worker 500. + + Acquires ``_load_lock`` so concurrent callers serialise. Only the + first caller actually cycles unload/load; subsequent callers whose + ``dead_detected_at`` is before a recent ``_loaded_at`` know a peer + already completed recovery and return True immediately. + + NOTE: We do NOT check the router's ``_status()`` to skip the cycle — + the router can report ``loaded`` while the worker child is already + dead. We always force an unload+reload on the first caller. + + Returns True when the worker is confirmed loaded, False on timeout/ + error so the caller can return 503. + """ + async with self._load_lock: + # Re-check: if _loaded_at was updated AFTER the dead-worker was + # detected, a peer coroutine already completed recovery. + if self._loaded == model and self._loaded_at > dead_detected_at: + logging.info( + "[recover_worker] %s already recovered by peer (loaded_at=%.3f > detected=%.3f)", + model, self._loaded_at, dead_detected_at, + ) + return True + + logging.warning( + "[recover_worker] cycling unload/load for dead worker %s", model + ) + self._loaded = None # invalidate proxy cache immediately + + # Unload — force the router to tear down the dead worker entry. + # Best-effort: the router may error if it can't contact the child, + # but we continue to the load step regardless. + url_unload = f"{self.config.backend_base_url}/models/unload" + auth_headers = { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + } + try: + async with self.session.post( + url_unload, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400: + logging.warning("[recover_worker] unload returned %s: %s", r.status, body_text) + else: + logging.info("[recover_worker] unload OK for %s", model) + except (ClientError, asyncio.TimeoutError) as e: + logging.warning("[recover_worker] unload error (ignored): %s", e) + + # Wait for the router to reflect the unloaded state so that + # _load_locked doesn't see "loaded" and short-circuit. + unload_deadline = time.monotonic() + 15 + while time.monotonic() < unload_deadline: + try: + status = await self._status(model) + if status != "loaded": + logging.info("[recover_worker] router confirms %s is %s", model, status) + break + except Exception: + pass + await asyncio.sleep(0.5) + else: + logging.warning( + "[recover_worker] router still shows loaded after 15s — proceeding anyway" + ) + + # Reload — spawn fresh worker(s), draining any buffered exit + # signals the router queued during unload. Each exit signal + # consumes exactly one fresh worker; we loop until the worker + # stays alive for a full second after reporting "loaded". + # + # We do NOT use _load_locked here because its poll + # loops until "loaded" is observed — if a queued exit fires + # between the worker reporting ready and our next poll (< 500 ms), + # the status flips to "unloaded" and the loop spins for 300 s. + # Instead we implement a custom load + rapid-poll that detects + # the loaded→unloaded flash and retries the load immediately. + load_url = f"{self.config.backend_base_url}/models/load" + max_load_attempts = 4 + per_load_timeout = 120 # seconds to wait for status != "loading" + + for load_attempt in range(1, max_load_attempts + 1): + logging.info( + "[recover_worker] load attempt %d/%d for %s", + load_attempt, max_load_attempts, model, + ) + try: + async with self.session.post( + load_url, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400 and "already running" not in body_text: + logging.warning("[recover_worker] load returned %s: %s", r.status, body_text) + except (ClientError, asyncio.TimeoutError) as e: + logging.error("[recover_worker] load POST failed: %s", e) + return False + + # Poll until status leaves "loading" (either loaded or unloaded) + poll_deadline = time.monotonic() + per_load_timeout + prev_status = "" + seen_loading = False # have we observed this attempt actually start? + while time.monotonic() < poll_deadline: + try: + status = await self._status(model) + except Exception: + await asyncio.sleep(0.5) + continue + if status != prev_status: + logging.info("[recover_worker] %s status: %s", model, status) + prev_status = status + if status == "loading": + seen_loading = True + if status == "loaded": + # Give the router 1.5s to process any pending exit signal + # before declaring victory. + await asyncio.sleep(1.5) + status2 = await self._status(model) + if status2 == "loaded": + self._loaded = model + self._loaded_at = time.monotonic() + logging.info( + "[recover_worker] worker stable for %s (attempt %d)", + model, load_attempt, + ) + return True + logging.warning( + "[recover_worker] fresh worker for %s exited immediately " + "(status after 1.5s: %s, attempt %d/%d)", + model, status2, load_attempt, max_load_attempts, + ) + break # queued exit consumed — try next load attempt + if status == "failed": + logging.error("[recover_worker] model %s failed to load", model) + return False + if status == "unloaded" and seen_loading: + # The worker started loading but was then stopped before ever + # reaching "loaded" — a queued exit signal consumed it (the + # router force-kills it after its ~10s stop timeout) or it + # crashed mid-load. For a slow-loading model the kill lands + # before "loaded" is ever observed, so the loaded→flash case + # above never fires. Retry the next load immediately instead + # of spinning here for the full per_load_timeout. + logging.warning( + "[recover_worker] fresh worker for %s died during load " + "(status: unloaded, attempt %d/%d) — retrying", + model, load_attempt, max_load_attempts, + ) + break # queued exit consumed — try next load attempt + await asyncio.sleep(0.25) + else: + logging.error( + "[recover_worker] timed out waiting for %s to load (attempt %d/%d)", + model, load_attempt, max_load_attempts, + ) + # Don't return False here — let the for-loop try the next attempt. + # The router may need a fresh /models/load POST to clear a stuck + # "loading" state from a previous crash. + # Retry the unload to give the router a chance to tear down + # any zombie worker entry from the timed-out attempt. + if load_attempt < max_load_attempts: + logging.info( + "[recover_worker] retrying unload before next load attempt" + ) + try: + async with self.session.post( + url_unload, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400: + logging.warning( + "[recover_worker] retry-unload returned %s: %s", + r.status, body_text, + ) + else: + logging.info( + "[recover_worker] retry-unload OK for %s", model + ) + except (ClientError, asyncio.TimeoutError) as e: + logging.warning( + "[recover_worker] retry-unload error (ignored): %s", e + ) + # Brief pause to let the router settle before the next load. + await asyncio.sleep(2) + continue + + logging.error("[recover_worker] all %d load attempts failed for %s", max_load_attempts, model) + return False + + def mark_worker_suspect(self) -> None: + """Flag the worker as possibly desynced after a mid-generation client + abort (cancel/disconnect). The next model request will probe before use. + See _worker_suspect for the upstream bug this guards against.""" + self._worker_suspect = True + + async def _probe_worker(self, model: str) -> bool: + """Send a minimal generation to detect a wedged/crashed worker after a + cancel. Returns True if the worker answers normally; False if it returns + a dead-worker error, errors out, or times out (a busy-finishing-an-orphan + or genuinely-wedged worker both warrant a recovery cycle).""" + url = f"{self.config.backend_base_url}/v1/chat/completions" + headers = { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + } + payload = { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 1, + "stream": False, + } + try: + async with self.session.post( + url, headers=headers, json=payload, + timeout=ClientTimeout(total=GUARD_PROBE_TIMEOUT), + ) as r: + body = await r.read() + if _is_dead_worker_response(r.status, body): + return False + return r.status < 500 + except (ClientError, asyncio.TimeoutError) as e: + logging.warning("[guard] probe error for %s: %s", model, e) + return False + + async def guard_after_cancel(self, model: str) -> None: + """If a prior request was aborted mid-generation, probe the worker once + before the next request uses it. Recover proactively if the probe shows + it wedged/dead — so the next real request lands on a clean worker instead + of triggering the upstream cancel→next-request crash.""" + if not self._worker_suspect: + return + async with self._guard_lock: + if not self._worker_suspect: + return # a concurrent caller already handled this episode + # Only the currently-loaded worker can be the wedged one. If a different + # model (or nothing) is loaded, use_model will spawn a fresh worker + # anyway, so there is nothing to probe — just clear the flag. + if self._loaded != model: + self._worker_suspect = False + return + logging.info("[guard] worker suspect after cancel — probing %s", model) + healthy = await self._probe_worker(model) + if healthy: + logging.info("[guard] worker %s healthy after cancel", model) + else: + logging.warning( + "[guard] worker %s wedged/dead after cancel — recovering", model + ) + await self.recover_worker(model, time.monotonic()) + self._worker_suspect = False + + async def unload_if_idle(self) -> None: + # Never unload while a streaming forward is in progress. + if self._active != 0 or self._forwarding != 0 or self._loaded is None: + return + idle = time.monotonic() - self._last_activity + if idle >= self.config.idle_timeout: + await self.unload(f"idle for {int(idle)}s")