server.py

#!/usr/bin/env python3
"""
server.py — a thin local HTTP/SSE wrapper around codegen.py for the VS Code extension.

PRINCIPLE (see BUILD-SPEC §0): the extension is UI only. ALL agent logic stays in
codegen.py. This server just:
  - starts a codegen run on POST /task,
  - streams every agent event to the webview over SSE (GET /events),
  - and blocks the run on file/run approvals until the webview answers (POST /approve).

It injects `emit` / `approver` callbacks into codegen (see codegen.agent_loop) instead
of codegen's print()/input(), so the CLI keeps working unchanged.

Endpoints:
  GET  /health                 -> {"ok": true}
  GET  /models                 -> {"models": [...]}  (from `ollama list`)
  POST /task   {task,dir,mode,model,think,autoApprove}  -> starts a run
  GET  /events (SSE)           -> stream of agent events (see EVENT PROTOCOL in the spec)
  POST /approve {id,approved,edited?}                   -> releases a pending approval
  POST /stop                   -> requests the current run to stop

Zero third-party deps: stdlib http.server only. Run:  python3 server.py
"""

import importlib
import json
import mimetypes
import os
import subprocess
import sys
import threading
import urllib.error
import urllib.request
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import Empty, Queue

import model_registry  # config-driven model registry (stdlib-only)

HOST = os.environ.get("CODEGEN_SERVER_HOST", "127.0.0.1")
PORT = int(os.environ.get("CODEGEN_SERVER_PORT", "8765"))
OLLAMA = os.environ.get("OLLAMA_HOST", "http://localhost:11434").rstrip("/")
GUI_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "gui")

# ── hot-reload of codegen.py ─────────────────────────────────────────────────
# The agent logic lives in codegen.py. We watch its mtime and reload the module
# when it changes, so edits take effect on the next /task WITHOUT restarting the
# server. (Changes to server.py itself still need a process restart.)
_codegen = None
_codegen_mtime = 0.0


def get_codegen():
    """Return the codegen module, hot-reloading it if codegen.py changed on disk.
    Called only between runs (start_run gates on no active run), so reload is safe."""
    global _codegen, _codegen_mtime
    import codegen as cg  # cached after first import; same object reload() updates in place
    try:
        mtime = os.path.getmtime(cg.__file__)
    except OSError:
        mtime = 0.0
    if _codegen is None:
        _codegen, _codegen_mtime = cg, mtime
    elif mtime > _codegen_mtime:
        importlib.reload(_codegen)
        _codegen_mtime = mtime
        print("🔄 codegen.py changed — reloaded.", file=sys.stderr)
        emit({"type": "assistant", "text": "🔄 codegen.py reloaded (auto)."})
    return _codegen


# ── event bus (pub/sub) ──────────────────────────────────────────────────────
# Every SSE client gets its own queue; publish() fans an event out to all of them.
_subscribers = []
_sub_lock = threading.Lock()


def subscribe():
    q = Queue()
    with _sub_lock:
        _subscribers.append(q)
    return q


def unsubscribe(q):
    with _sub_lock:
        if q in _subscribers:
            _subscribers.remove(q)


def publish(event):
    """Send an event dict to every connected SSE client."""
    with _sub_lock:
        subs = list(_subscribers)
    for q in subs:
        q.put(event)


# ── run state (one run at a time — matches a single small local model) ────────
class Run:
    def __init__(self):
        self.stop = threading.Event()       # set by POST /stop
        self.approvals = {}                 # id -> {"event": Event, "result": dict}
        self.lock = threading.Lock()
        self.thread = None

    def is_active(self):
        return self.thread is not None and self.thread.is_alive()


_run = Run()


def emit(event):
    """codegen calls this for every action/result/etc. Just fan it out to the UI."""
    publish(event)


def approver(action):
    """codegen calls this instead of input() when a file/run needs approval.

    `action` is the proposed event dict (already has a unique 'id'). We publish it,
    then block the worker thread until POST /approve answers (or the run is stopped).
    Returns {"approved": bool, "edited": str|None}.
    """
    aid = action.get("id") or uuid.uuid4().hex[:8]
    action["id"] = aid
    ev = threading.Event()
    slot = {"event": ev, "result": {"approved": False, "edited": None}}
    with _run.lock:
        _run.approvals[aid] = slot
    publish(action)
    # Wait for the answer, but stay responsive to /stop.
    while not ev.wait(timeout=0.25):
        if _run.stop.is_set():
            slot["result"] = {"approved": False, "edited": None, "stopped": True}
            break
    with _run.lock:
        _run.approvals.pop(aid, None)
    return slot["result"]


def resolve_approval(aid, approved, edited=None):
    with _run.lock:
        slot = _run.approvals.get(aid)
    if not slot:
        return False
    slot["result"] = {"approved": bool(approved), "edited": edited}
    slot["event"].set()
    return True


# ── run drivers ──────────────────────────────────────────────────────────────
def _selftest_run():
    """Step-1 plumbing check — no ollama needed. Exercises every event type and the
    approval round-trip so the server can be validated with curl before codegen is wired."""
    emit({"type": "thinking"})
    emit({"type": "action", "kind": "search", "label": "raspberry pi 5 price"})
    emit({"type": "result", "text": "1. Raspberry Pi 5 — $60 (8GB)\n   https://example.com"})
    decision = approver({"type": "file_proposed", "id": "selftest1",
                         "path": "hello.py", "oldText": "", "newText": "print('hi')\n"})
    if decision.get("stopped"):
        emit({"type": "error", "text": "stopped by user"})
        return
    if decision["approved"]:
        content = decision.get("edited") or "print('hi')\n"
        emit({"type": "file_written", "path": "hello.py"})
        emit({"type": "assistant", "text": f"Wrote hello.py ({len(content)} bytes)."})
    else:
        emit({"type": "assistant", "text": "User rejected hello.py."})
    # Note: the {"type":"done"} event is emitted once by the worker's finally block.


def _codegen_run(params):
    """Drive the real agent. Loaded via get_codegen() so the server starts even if
    codegen's optional deps are missing, a bad import surfaces as a run error (not a
    crash), and on-disk edits to codegen.py are picked up automatically."""
    codegen = get_codegen()

    task = params.get("task", "").strip()
    base = os.path.abspath(params.get("dir") or ".")
    os.makedirs(base, exist_ok=True)
    mode = params.get("mode", "agent")
    auto = bool(params.get("autoApprove"))
    model = params.get("model")
    think = bool(params.get("think"))
    max_iters = int(params.get("max", 8))

    if model:
        codegen.MODEL = model
    # Per-run context-window selection from the UI slider (Ollama options.num_ctx —
    # overrides the Modelfile, no rebuild). None resets to the per-model registry default.
    nctx = params.get("num_ctx")
    codegen.NUM_CTX = int(nctx) if nctx else None
    codegen.THINK_PLAN = think
    codegen.PROJECT_MEMORY = codegen.load_memory(base)
    codegen.HOOKS = codegen.load_hooks(base)

    # auto-approve maps to codegen's -y; otherwise the webview approves each action.
    app = None if auto else approver
    should_stop = _run.stop.is_set  # checked between steps so Stop works even in auto mode

    if mode == "research":
        # Deep critical research: plan sub-questions -> grounded search each -> synthesize report.
        codegen.ASSISTANT = False
        codegen.research_pipeline(task, emit=emit, should_stop=should_stop)
    elif mode == "web":
        # Grounded web answer: harness searches + fetches pages, model only phrases. No file writes.
        codegen.ASSISTANT = False
        codegen.web_answer(task, emit=emit, should_stop=should_stop)
    elif mode == "build":
        codegen.ASSISTANT = False
        codegen.orchestrate(task, base, auto, emit=emit, approver=app, should_stop=should_stop)
    else:
        # webview modes: "agent" = autonomous coding loop, "chat" = conversational
        # assistant (prose answers, no forced file writes). Reset the global each run.
        codegen.ASSISTANT = (mode == "chat")
        base_prompt = codegen.ASSISTANT_PROMPT if codegen.ASSISTANT else codegen.SYSTEM_PROMPT
        messages = [{"role": "system", "content": codegen.with_memory(base_prompt)}]
        codegen.agent_loop(task, messages, base, auto, max_iters,
                           emit=emit, approver=app, should_stop=should_stop)


def start_run(params):
    if _run.is_active():
        return False, "a run is already active"
    _run.stop.clear()
    with _run.lock:
        _run.approvals.clear()

    def worker():
        try:
            if params.get("mode") == "selftest":
                _selftest_run()
            else:
                _codegen_run(params)
        except Exception as e:  # noqa: BLE001 — surface any failure to the UI
            import traceback
            emit({"type": "error", "text": f"{e}\n{traceback.format_exc()}"})
        finally:
            emit({"type": "done"})

    t = threading.Thread(target=worker, daemon=True)
    _run.thread = t
    t.start()
    return True, "started"


# Embedding models can't do chat (Ollama 400 "does not support chat"); hide them from the picker.
_EMBED_HINT = ("embed", "bge", "minilm", "gte", "e5-")


def _is_chat_model(name):
    low = name.lower()
    return not any(h in low for h in _EMBED_HINT)


def list_models():
    try:
        r = subprocess.run(["ollama", "list"], capture_output=True, text=True, timeout=10)
        names = []
        for line in r.stdout.splitlines()[1:]:  # skip header
            parts = line.split()
            if parts and _is_chat_model(parts[0]):
                names.append(parts[0])
        return names
    except Exception:
        return []


# ── Ollama model management (proxy the official HTTP API) ─────────────────────
# Exact request/response shapes per github.com/ollama/ollama/blob/main/docs/api.md.
def _ollama_get(path, timeout=15):
    req = urllib.request.Request(OLLAMA + path)
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8", "ignore") or "{}")


def installed_models():
    """Merge /api/tags (installed) with /api/ps (loaded) into rich model cards."""
    try:
        tags = _ollama_get("/api/tags").get("models", [])
    except Exception as e:  # noqa: BLE001
        return {"models": [], "error": f"Ollama unreachable at {OLLAMA}: {e}"}
    try:
        loaded = {m.get("name") for m in _ollama_get("/api/ps").get("models", [])}
    except Exception:
        loaded = set()
    out = []
    for m in tags:
        d = m.get("details", {}) or {}
        out.append({
            "name": m.get("name"),
            "size": m.get("size"),
            "family": d.get("family"),
            "parameter_size": d.get("parameter_size"),
            "quantization": d.get("quantization_level"),
            "modified_at": m.get("modified_at"),
            "loaded": m.get("name") in loaded,
            "chat": _is_chat_model(m.get("name", "")),
        })
    out.sort(key=lambda x: (x.get("name") or "").lower())
    return {"models": out}


def _pull_worker(name):
    """Stream /api/pull NDJSON and re-publish aggregated progress to the SSE bus.
    Progress is aggregated per layer digest: percent = sum(completed)/sum(total)."""
    publish({"type": "pull_progress", "model": name, "status": "starting", "percent": 0})
    digests = {}
    try:
        data = json.dumps({"model": name, "stream": True}).encode()
        req = urllib.request.Request(OLLAMA + "/api/pull", data=data,
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=3600) as resp:
            for line in resp:  # NDJSON — one object per line
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except ValueError:
                    continue
                if obj.get("error"):
                    publish({"type": "pull_progress", "model": name, "status": "error",
                             "error": obj["error"], "done": True})
                    return
                if obj.get("digest") and obj.get("total"):
                    digests[obj["digest"]] = (obj.get("completed", 0) or 0, obj.get("total", 0) or 0)
                tot = sum(t for _, t in digests.values())
                comp = sum(c for c, _ in digests.values())
                pct = round(comp / tot * 100, 1) if tot else 0
                publish({"type": "pull_progress", "model": name, "status": obj.get("status", ""),
                         "completed": comp, "total": tot, "percent": pct})
        publish({"type": "pull_progress", "model": name, "status": "success",
                 "percent": 100, "done": True})
    except Exception as e:  # noqa: BLE001
        publish({"type": "pull_progress", "model": name, "status": "error",
                 "error": str(e), "done": True})


def _create_worker(name, params):
    """Stream /api/create (create a model from a registry base + parameters)."""
    publish({"type": "create_progress", "model": name, "status": "starting"})
    payload = {"model": name, "stream": True}
    if params.get("from"):
        payload["from"] = params["from"]
    p = {}
    if params.get("num_ctx"):
        p["num_ctx"] = int(params["num_ctx"])
    if params.get("temperature") is not None:
        p["temperature"] = float(params["temperature"])
    if p:
        payload["parameters"] = p
    if params.get("system"):
        payload["system"] = params["system"]
    try:
        data = json.dumps(payload).encode()
        req = urllib.request.Request(OLLAMA + "/api/create", data=data,
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=3600) as resp:
            for line in resp:
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except ValueError:
                    continue
                if obj.get("error"):
                    publish({"type": "create_progress", "model": name, "status": "error",
                             "error": obj["error"], "done": True})
                    return
                publish({"type": "create_progress", "model": name, "status": obj.get("status", "")})
        publish({"type": "create_progress", "model": name, "status": "success", "done": True})
    except Exception as e:  # noqa: BLE001
        publish({"type": "create_progress", "model": name, "status": "error",
                 "error": str(e), "done": True})


def delete_model(name):
    """DELETE /api/delete. Returns (ok, error)."""
    try:
        data = json.dumps({"model": name}).encode()
        req = urllib.request.Request(OLLAMA + "/api/delete", data=data,
                                     headers={"Content-Type": "application/json"}, method="DELETE")
        with urllib.request.urlopen(req, timeout=30) as r:
            r.read()
        return True, ""
    except urllib.error.HTTPError as e:
        return False, f"HTTP {e.code}"
    except Exception as e:  # noqa: BLE001
        return False, str(e)


def start_bg(target, *a):
    threading.Thread(target=target, args=a, daemon=True).start()


# ── static GUI serving ───────────────────────────────────────────────────────
def _serve_static(handler, path):
    """Serve the standalone GUI from gui/. '/' -> index.html. Path-confined."""
    rel = "index.html" if path in ("/", "") else path.lstrip("/")
    full = os.path.normpath(os.path.join(GUI_DIR, rel))
    if not full.startswith(GUI_DIR) or not os.path.isfile(full):
        handler._send_json({"error": "not found"}, 404)
        return
    ctype = mimetypes.guess_type(full)[0] or "application/octet-stream"
    with open(full, "rb") as f:
        body = f.read()
    handler.send_response(200)
    handler.send_header("Content-Type", ctype)
    handler.send_header("Content-Length", str(len(body)))
    handler.send_header("Cache-Control", "no-cache")
    handler.end_headers()
    handler.wfile.write(body)


# ── HTTP handler ─────────────────────────────────────────────────────────────
class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *args):
        pass  # quiet; the agent's own logs go to its stderr

    def _send_json(self, obj, status=200):
        body = json.dumps(obj).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(body)

    def _read_json(self):
        n = int(self.headers.get("Content-Length", 0))
        if not n:
            return {}
        try:
            return json.loads(self.rfile.read(n).decode() or "{}")
        except ValueError:
            return {}

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    def do_GET(self):
        path = self.path.split("?", 1)[0]
        if path == "/health":
            return self._send_json({"ok": True})
        if path == "/models":
            return self._send_json({"models": list_models()})
        if path == "/models/installed":
            return self._send_json(installed_models())
        if path == "/registry":
            return self._send_json(model_registry.load())
        if path.startswith("/events"):
            return self._stream_events()
        # everything else: serve the standalone GUI's static assets.
        return _serve_static(self, path)

    def do_POST(self):
        path = self.path.split("?", 1)[0]
        if path == "/task":
            params = self._read_json()
            ok, msg = start_run(params)
            return self._send_json({"ok": ok, "message": msg}, 200 if ok else 409)
        if path == "/approve":
            data = self._read_json()
            ok = resolve_approval(data.get("id"), data.get("approved"), data.get("edited"))
            return self._send_json({"ok": ok})
        if path == "/stop":
            _run.stop.set()
            return self._send_json({"ok": True})
        if path == "/models/pull":
            name = (self._read_json().get("name") or "").strip()
            if not name:
                return self._send_json({"ok": False, "error": "name required"}, 400)
            start_bg(_pull_worker, name)
            return self._send_json({"ok": True, "message": f"pulling {name}"})
        if path == "/models/create":
            data = self._read_json()
            name = (data.get("name") or "").strip()
            if not name or not data.get("from"):
                return self._send_json({"ok": False, "error": "name and from required"}, 400)
            start_bg(_create_worker, name, data)
            return self._send_json({"ok": True, "message": f"creating {name}"})
        if path == "/models/delete":
            name = (self._read_json().get("name") or "").strip()
            ok, err = delete_model(name)
            return self._send_json({"ok": ok, "error": err}, 200 if ok else 400)
        self._send_json({"error": "not found"}, 404)

    def do_PUT(self):
        path = self.path.split("?", 1)[0]
        if path == "/registry":
            data = self._read_json()
            ok, errors = model_registry.validate(data)
            if not ok:
                return self._send_json({"ok": False, "errors": errors}, 400)
            try:
                model_registry.save(data)
            except (OSError, ValueError) as e:
                return self._send_json({"ok": False, "errors": [str(e)]}, 500)
            return self._send_json({"ok": True})
        self._send_json({"error": "not found"}, 404)

    def _stream_events(self):
        # With protocol_version HTTP/1.1 an infinite stream must NOT advertise a
        # Content-Length; mark the connection to close so the framework doesn't try.
        self.close_connection = True
        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.send_header("X-Accel-Buffering", "no")  # defeat any proxy buffering
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        q = subscribe()
        try:
            self.wfile.write(b": connected\n\n")
            self.wfile.flush()
            while True:
                try:
                    event = q.get(timeout=15)
                except Empty:
                    self.wfile.write(b": ping\n\n")  # keep the connection alive
                    self.wfile.flush()
                    continue
                payload = json.dumps(event)
                self.wfile.write(f"data: {payload}\n\n".encode())
                self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            pass  # client disconnected
        finally:
            unsubscribe(q)


def main():
    httpd = ThreadingHTTPServer((HOST, PORT), Handler)
    httpd.daemon_threads = True  # open SSE threads must not block interpreter exit on Ctrl-C
    print(f"🌐 codegen server on http://{HOST}:{PORT}  (Ctrl-C to stop)", file=sys.stderr)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nBye 👋", file=sys.stderr)


if __name__ == "__main__":
    main()