mission-control-v2/deploy/governor/governor.py

#!/usr/bin/env python3
"""Governor — duenner, zustandsloser Token-Waechter-Proxy (Phase 0).

Sitzt zwischen einem Coding-Agenten (Aider) und dem Modell-Endpoint (llama-swap
:8080). Reicht ALLES unveraendert durch — mit einer Ausnahme bei
/v1/chat/completions: er schaetzt die Token-Groesse der Anfrage (= Sessiongroesse,
weil die ganze Historie jede Runde mitkommt) und handelt nach zwei Schwellen:

  est >= SCHWELLE (soft): haengt eine Stopp-Anweisung als letzte User-Nachricht an
      ("SAVEPOINT.md finalisieren + stoppen") und leitet weiter. Das Modell schreibt
      EINEN ehrlichen Abschluss-Savepoint. Loggt FIRED.
  est >= HART-DECKEL (optional): antwortet SELBST mit einer kurzen Stopp-Nachricht,
      OHNE das Modell zu fragen. Verhindert, dass ueber die Grenze hinaus
      weitergearbeitet wird — genau das erzeugte in Tests eine Fassade. Loggt HARDSTOP.

Bewusst nur Standardbibliothek: kein pip, kein venv, laeuft mit System-python3.
Bewusst zustandslos: jede Anfrage wird fuer sich bewertet; keine Sitzungs-DB.

Konfiguration per Umgebungsvariablen (alle optional):
  GOV_PORT            Listen-Port                        (Default 8100)
  GOV_HOST            Listen-Adresse                     (Default 0.0.0.0)
  GOV_UPSTREAM        Modell-Endpoint                    (Default http://127.0.0.1:8080)
  GOV_THRESHOLD       Soft-Schwelle fuer den Einschub    (Default 25000)
  GOV_HARD_CEILING    Hart-Deckel; 0 = aus               (Default 0)
  GOV_CHARS_PER_TOKEN Heuristik Zeichen->Token           (Default 3.5, kalibriert)
  GOV_LOG             Logdatei (zusaetzlich zu stdout)   (Default ./governor.log)
  GOV_DIRECTIVE       Text des Soft-Einschubs (sonst Default unten)
  GOV_HARDSTOP_MSG    Text der Hart-Stopp-Antwort (sonst Default unten)
"""

import http.client
import json
import os
import re
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

# ---- Konfiguration ---------------------------------------------------------

PORT = int(os.environ.get("GOV_PORT", "8100"))
HOST = os.environ.get("GOV_HOST", "0.0.0.0")
UPSTREAM = os.environ.get("GOV_UPSTREAM", "http://127.0.0.1:8080")
THRESHOLD = int(os.environ.get("GOV_THRESHOLD", "25000"))
HARD_CEILING = int(os.environ.get("GOV_HARD_CEILING", "0"))  # 0 = deaktiviert
CHARS_PER_TOKEN = float(os.environ.get("GOV_CHARS_PER_TOKEN", "3.5"))
LOG_PATH = os.environ.get("GOV_LOG", os.path.join(os.getcwd(), "governor.log"))

DEFAULT_DIRECTIVE = (
    "[GOVERNOR — SITZUNGS-LIMIT ERREICHT] Der Kontext dieser Sitzung ist auf ~{est} "
    "Tokens gewachsen (Limit {threshold}). Beginne oder setze JETZT KEINE weiteren "
    "Code-Aenderungen fort. Stattdessen, in dieser Reihenfolge:\n"
    "1. Aktualisiere SAVEPOINT.md so, dass es den aktuellen Stand vollstaendig festhaelt: "
    "was WIRKLICH erledigt ist (nur was im Code steht — nichts aus Absicht oder git-"
    "Nachrichten ableiten), der genaue naechste Schritt, offene Fragen und alle "
    "Stolpersteine — genug, dass eine frische Sitzung ohne jede Erinnerung allein aus "
    "SAVEPOINT.md plus git-Historie sauber weitermachen kann.\n"
    "2. Halte dann an und sage dem Nutzer in einem Satz, dass er eine frische Sitzung "
    "starten soll. Gib ausser der SAVEPOINT.md-Aktualisierung und diesem Hinweis nichts aus."
)
DIRECTIVE = os.environ.get("GOV_DIRECTIVE", DEFAULT_DIRECTIVE)

DEFAULT_HARDSTOP = (
    "[GOVERNOR — HARTER STOPP] Das Sitzungs-Limit ist ueberschritten und der Savepoint "
    "sollte bereits finalisiert sein. Diese Sitzung nimmt keine weiteren Auftraege mehr an. "
    "Bitte starte eine FRISCHE Sitzung — sie liest SAVEPOINT.md und die git-Historie und "
    "macht sauber weiter. (Keine Code-Aenderung in dieser Antwort.)"
)
HARDSTOP_MSG = os.environ.get("GOV_HARDSTOP_MSG", DEFAULT_HARDSTOP)

up = urlparse(UPSTREAM)
UP_HOST = up.hostname or "127.0.0.1"
UP_PORT = up.port or 80

HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade",
}

_log_lock = threading.Lock()
_PROMPT_TOKENS_RE = re.compile(r'"prompt_tokens"\s*:\s*(\d+)')


def log(line: str) -> None:
    """Eine Zeile nach stdout UND in die Logdatei (thread-sicher)."""
    stamp = time.strftime("%Y-%m-%dT%H:%M:%S")
    msg = f"{stamp} {line}"
    with _log_lock:
        print(msg, flush=True)
        try:
            with open(LOG_PATH, "a", encoding="utf-8") as fh:
                fh.write(msg + "\n")
        except OSError:
            pass


def estimate_tokens(messages) -> int:
    """Grobe, aber stabile Heuristik: Zeichen aller Nachrichteninhalte / CPT.

    Gegen die echten prompt_tokens aus der Antwort kalibriert (CPT=3.5 traf am
    24.07. auf ~1-3 % genau). Zaehlt Text in String- und Multimodal-Listen-Inhalten;
    kleiner Aufschlag je Nachricht fuer Rollen-/Template-Overhead.
    """
    chars = 0
    for m in messages or []:
        chars += 4
        content = m.get("content") if isinstance(m, dict) else None
        if isinstance(content, str):
            chars += len(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and isinstance(part.get("text"), str):
                    chars += len(part["text"])
    return int(chars / CHARS_PER_TOKEN)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "Governor/0.2"

    def log_message(self, *args):
        pass

    def do_GET(self):
        self._proxy()

    def do_POST(self):
        self._proxy()

    def do_PUT(self):
        self._proxy()

    def do_DELETE(self):
        self._proxy()

    def do_OPTIONS(self):
        self._proxy()

    # -- Kern ---------------------------------------------------------------
    def _read_body(self) -> bytes:
        length = self.headers.get("Content-Length")
        if length is None:
            return b""
        try:
            return self.rfile.read(int(length))
        except (ValueError, OSError):
            return b""

    def _proxy(self) -> None:
        body = self._read_body()
        path = self.path
        # Query-String vor der Endpunkt-Erkennung abschneiden (sonst umgeht
        # z. B. ?api-version=... den Governor).
        clean_path = path.split("?", 1)[0]
        is_chat = clean_path.rstrip("/").endswith("/chat/completions")

        action = "passthrough"
        est = None
        streaming = False
        model = ""
        if is_chat and body:
            action, body, est, streaming, model = self._decide(body)

        # HARTER STOPP: selbst antworten, Upstream nie fragen.
        if action == "hard":
            self._send_canned_stop(model, streaming, est)
            log(f"chat est={est} thr={THRESHOLD} hard={HARD_CEILING} HARDSTOP "
                f"stream={streaming} status=200")
            return

        # Header fuer Upstream aufbereiten.
        out_headers = {}
        for k, v in self.headers.items():
            kl = k.lower()
            if kl in HOP_BY_HOP or kl in ("host", "content-length", "accept-encoding"):
                continue
            out_headers[k] = v
        out_headers["Host"] = f"{UP_HOST}:{UP_PORT}"
        out_headers["Accept-Encoding"] = "identity"
        if body:
            out_headers["Content-Length"] = str(len(body))
        out_headers["Connection"] = "close"

        conn = None
        try:
            conn = http.client.HTTPConnection(UP_HOST, UP_PORT, timeout=600)
            conn.request(self.command, path, body=body or None, headers=out_headers)
            resp = conn.getresponse()
        except (OSError, http.client.HTTPException) as exc:
            log(f"ERROR upstream {self.command} {path}: {exc!r}")
            if conn is not None:
                conn.close()
            self._safe_error(502, f"governor upstream: {exc}")
            return

        self.send_response(resp.status)
        for k, v in resp.getheaders():
            kl = k.lower()
            # hop-by-hop + Laenge raus; Date/Server setzt send_response schon selbst.
            if kl in HOP_BY_HOP or kl in ("content-length", "date", "server"):
                continue
            self.send_header(k, v)
        self.send_header("Connection", "close")
        self.end_headers()

        tail = bytearray()
        try:
            while True:
                # read1() gibt jedes Upstream-Stueck sofort zurueck (echtes SSE-
                # Durchreichen). read() wuerde bis 64 KB oder Stream-Ende puffern
                # und streamendes Aider die ganze Generierung haengen lassen.
                chunk = resp.read1(65536)
                if not chunk:
                    break
                self.wfile.write(chunk)
                self.wfile.flush()
                tail.extend(chunk)
                if len(tail) > 16384:
                    del tail[:-16384]
        except OSError:
            pass
        finally:
            conn.close()

        exact = self._scan_prompt_tokens(tail)
        if is_chat:
            exact_s = str(exact) if exact is not None else "-"
            flag = "FIRED" if action == "soft" else "ok"
            log(f"chat est={est} exact={exact_s} thr={THRESHOLD} {flag} "
                f"stream={streaming} status={resp.status}")

    def _decide(self, body: bytes):
        """Aktion bestimmen: passthrough | soft (Einschub) | hard (Selbstantwort).

        Rueckgabe: (action, body, est, streaming, model).
        """
        try:
            data = json.loads(body)
        except (ValueError, UnicodeDecodeError):
            return "passthrough", body, None, False, ""
        if not isinstance(data, dict):
            return "passthrough", body, None, False, ""

        messages = data.get("messages")
        streaming = bool(data.get("stream"))
        model = data.get("model", "") or ""
        est = estimate_tokens(messages if isinstance(messages, list) else [])

        if HARD_CEILING > 0 and est >= HARD_CEILING:
            return "hard", body, est, streaming, model

        if est >= THRESHOLD and isinstance(messages, list):
            # Sichere Substitution statt str.format: ein operator-gesetzter
            # GOV_DIRECTIVE mit { } (JSON/Code-Beispiel) darf nicht crashen.
            directive = (DIRECTIVE.replace("{est}", str(est))
                                  .replace("{threshold}", str(THRESHOLD)))
            messages.append({"role": "user", "content": directive})
            data["messages"] = messages
            return "soft", json.dumps(data).encode("utf-8"), est, streaming, model

        return "passthrough", body, est, streaming, model

    def _send_canned_stop(self, model: str, streaming: bool, est) -> None:
        """OpenAI-kompatible Stopp-Antwort selbst erzeugen (kein Upstream-Call)."""
        created = int(time.time())
        usage = {"prompt_tokens": est or 0, "completion_tokens": 0,
                 "total_tokens": est or 0}
        try:
            if streaming:
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.send_header("Connection", "close")
                self.end_headers()

                def sse(obj):
                    self.wfile.write(b"data: " + json.dumps(obj).encode() + b"\n\n")
                    self.wfile.flush()

                base = {"id": "governor-hardstop", "object": "chat.completion.chunk",
                        "created": created, "model": model}
                sse({**base, "choices": [{"index": 0, "delta": {"role": "assistant"},
                                          "finish_reason": None}]})
                sse({**base, "choices": [{"index": 0, "delta": {"content": HARDSTOP_MSG},
                                          "finish_reason": None}]})
                sse({**base, "choices": [{"index": 0, "delta": {},
                                          "finish_reason": "stop"}], "usage": usage})
                self.wfile.write(b"data: [DONE]\n\n")
                self.wfile.flush()
            else:
                payload = {
                    "id": "governor-hardstop", "object": "chat.completion",
                    "created": created, "model": model,
                    "choices": [{"index": 0, "finish_reason": "stop",
                                 "message": {"role": "assistant", "content": HARDSTOP_MSG}}],
                    "usage": usage,
                }
                data = json.dumps(payload).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(data)))
                self.send_header("Connection", "close")
                self.end_headers()
                self.wfile.write(data)
                self.wfile.flush()
        except OSError:
            pass

    def _safe_error(self, status: int, msg: str) -> None:
        try:
            data = json.dumps({"error": msg}).encode()
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(data)
        except OSError:
            pass

    @staticmethod
    def _scan_prompt_tokens(tail: bytearray):
        if not tail:
            return None
        try:
            text = tail.decode("utf-8", errors="ignore")
        except Exception:  # noqa: BLE001
            return None
        matches = _PROMPT_TOKENS_RE.findall(text)
        if not matches:
            return None
        try:
            return int(matches[-1])
        except ValueError:
            return None


def main() -> int:
    log_dir = os.path.dirname(LOG_PATH)
    if log_dir and not os.path.isdir(log_dir):
        try:
            os.makedirs(log_dir, exist_ok=True)
        except OSError:
            pass

    server = ThreadingHTTPServer((HOST, PORT), Handler)
    server.daemon_threads = True
    hard = HARD_CEILING if HARD_CEILING > 0 else "aus"
    log(f"Governor startet auf {HOST}:{PORT} -> {UPSTREAM} | Soft={THRESHOLD} "
        f"Hart={hard} | CPT={CHARS_PER_TOKEN} | Log={LOG_PATH}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        log("Governor beendet (SIGINT).")
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())