m64_fd6b2982ea67 / agent_gateway.py
381 lines, 15.0 KBraw
1"""Agent gateway for the opencode-basalt-5090 capsule (adapted from opencode-dsflash-openrouter-gb10's).23Adapted from an earlier agent gateway. There the gateway was4also the sandbox's only door out (LLM forward + web proxy). Here m64 owns that: the5capsule sits on an internal network whose only egress is the m64 gate, which lets out only6what the user granted at serve time (nothing by default): the model (Basalt) runs on7127.0.0.1:9000 in the same container. So this file keeps just the harness-neutral control API and8drives the local `opencode serve` on 127.0.0.1:4096.910Routes (the agent.session kind; the m64 front proxies only these and /health):1112 GET /health 200 when opencode answers13 GET /v1/agent what this is14 POST /v1/sessions new session, fresh empty directory15 POST /v1/sessions/{id}/messages {"text": ...}; wait=true (default) blocks for the reply16 GET /v1/sessions/{id}/events SSE stream of this session's events17 GET /v1/sessions/{id}/files[?path=] list a dir / download a file18 PUT /v1/sessions/{id}/files?path= upload a file (raw body)19 POST /v1/sessions/{id}/abort stop the running turn20"""21import asyncio22import json23import logging24import mimetypes25import os26import stat27import time28import uuid29from contextlib import asynccontextmanager30from pathlib import Path31from urllib.parse import quote3233import httpx34import uvicorn35from fastapi import FastAPI, HTTPException, Request36from fastapi.responses import JSONResponse, StreamingResponse3738PORT = int(os.environ.get("M64_PORT", "9500"))39OPENCODE = os.environ.get("OPENCODE_URL", "http://127.0.0.1:4096")40PROVIDER = os.environ.get("AGENT_PROVIDER", "basalt")41MODEL = os.environ.get("AGENT_MODEL", "qwen3.8-flash-next-iq3_s-basalt-384k")4243# What the user granted this run when serving it (m64 serve --allow-external-network /44# --allow-internal-network); the m64 gate enforces it.45_REACH = {46 "": "nothing: no egress at all",47 "external": "the public internet only (never the LAN or tailnet)",48 "internal": "the LAN and tailnet only",49 "external,internal": "the public internet, the LAN and the tailnet",50}51NETWORK = _REACH.get(os.environ.get("M64_NETWORK", ""), os.environ.get("M64_NETWORK", ""))52WORK = Path(os.environ.get("WORK_DIR", "/scratch/work"))53SESSIONS_DIR = WORK / "sessions"54REGISTRY = WORK / ".gateway" / "sessions.json"5556logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")57log = logging.getLogger("agent_gateway")585960# ── Session registry (lives in /scratch, so it goes with the run) ────────────────6162def _load_registry() -> dict:63 try:64 return json.loads(REGISTRY.read_text())65 except Exception:66 return {}676869def _save_registry(reg: dict) -> None:70 REGISTRY.parent.mkdir(parents=True, exist_ok=True)71 tmp = REGISTRY.with_suffix(".tmp")72 tmp.write_text(json.dumps(reg, indent=2))73 tmp.replace(REGISTRY)747576def _session(reg: dict, sid: str) -> dict:77 if sid not in reg:78 raise HTTPException(404, f"unknown session '{sid}'")79 return reg[sid]808182def _safe_path(root: Path, rel: str) -> Path:83 if "\x00" in rel:84 raise HTTPException(400, "invalid path")85 try:86 p = (root / rel.lstrip("/")).resolve()87 except (OSError, ValueError, RuntimeError): # e.g. a symlink loop88 raise HTTPException(400, "invalid path")89 if p != root.resolve() and root.resolve() not in p.parents:90 raise HTTPException(400, "path escapes the session directory")91 return p929394def _kind(p: Path) -> str | None:95 """"file", "dir", "other" (FIFO, socket, device...), or None if it does not exist."""96 try:97 mode = p.stat().st_mode98 except (FileNotFoundError, NotADirectoryError):99 return None100 except OSError:101 return "other"102 return "file" if stat.S_ISREG(mode) else "dir" if stat.S_ISDIR(mode) else "other"103104105# Files are written by the agent, so their bytes are untrusted. They are only ever served106# as downloads, sandboxed, never sniffed, and typed as one of these inert types or bytes.107_INERT_TYPES = {"text/plain", "image/png", "image/jpeg", "image/gif", "image/webp"}108_DOWNLOAD_HEADERS = {"X-Content-Type-Options": "nosniff", "Content-Security-Policy": "sandbox"}109110111def _open_regular(p: Path, flags: int):112 """Open p as a regular file, never blocking on a FIFO or following a final symlink."""113 try:114 fd = os.open(p, flags | os.O_NONBLOCK | os.O_NOFOLLOW, 0o644)115 except OSError:116 raise HTTPException(400, "not a regular file")117 if not stat.S_ISREG(os.fstat(fd).st_mode):118 os.close(fd)119 raise HTTPException(400, "not a regular file")120 os.set_blocking(fd, True)121 return os.fdopen(fd, "rb" if flags == os.O_RDONLY else "wb")122123124def _download_type(name: str) -> str:125 guessed = mimetypes.guess_type(name)[0]126 return guessed if guessed in _INERT_TYPES else "application/octet-stream"127128129# ── opencode translator ─────────────────────────────────────────────────────────130131class Opencode:132 """`opencode serve` HTTP API. Every call carries ?directory= so each session works in133 its own directory under $WORK/sessions."""134135 def __init__(self, base: str):136 self.base = base.rstrip("/")137 # trust_env=False: this is loopback; never route it through HTTPS_PROXY.138 self.http = httpx.AsyncClient(base_url=self.base, timeout=httpx.Timeout(30, read=None),139 trust_env=False)140141 async def healthy(self) -> bool:142 try:143 return (await self.http.get("/global/health", timeout=3)).status_code == 200144 except Exception:145 return False146147 async def create(self, directory: str, title: str | None) -> str:148 r = await self.http.post("/session", params={"directory": directory},149 json={"title": title} if title else {})150 r.raise_for_status()151 return r.json()["id"]152153 @staticmethod154 def _prompt_body(text: str) -> dict:155 return {"parts": [{"type": "text", "text": text}],156 "model": {"providerID": PROVIDER, "modelID": MODEL}}157158 async def send(self, sid: str, directory: str, text: str) -> dict:159 r = await self.http.post(f"/session/{sid}/message", params={"directory": directory},160 json=self._prompt_body(text))161 r.raise_for_status()162 reply = self._normalize(r.json())163 # One turn = several assistant messages (one per tool step); the response is only164 # the last. Gather the tool calls of the whole turn from the history.165 turn = []166 for m in reversed(await self.history(sid, directory)):167 if m["role"] == "user":168 break169 turn.append(m)170 reply["tools"] = [t for m in reversed(turn) for t in m["tools"]]171 reply["steps"] = len(turn)172 return reply173174 async def send_async(self, sid: str, directory: str, text: str) -> None:175 r = await self.http.post(f"/session/{sid}/prompt_async", params={"directory": directory},176 json=self._prompt_body(text))177 r.raise_for_status()178179 async def history(self, sid: str, directory: str) -> list:180 r = await self.http.get(f"/session/{sid}/message", params={"directory": directory})181 r.raise_for_status()182 return [self._normalize(m) for m in r.json()]183184 async def abort(self, sid: str, directory: str) -> None:185 await self.http.post(f"/session/{sid}/abort", params={"directory": directory})186187 async def events(self, sid: str, directory: str):188 """Yield SSE frames for events that belong to `sid`."""189 async with self.http.stream("GET", "/event", params={"directory": directory}) as r:190 async for line in r.aiter_lines():191 if not line.startswith("data:"):192 continue193 try:194 ev = json.loads(line[5:])195 except ValueError:196 continue197 props = ev.get("properties") or {}198 owner = (props.get("sessionID")199 or (props.get("info") or {}).get("sessionID")200 or (props.get("part") or {}).get("sessionID")201 or (props.get("info") or {}).get("id"))202 if owner == sid or ev.get("type") == "server.connected":203 yield f"data: {json.dumps(ev)}\n\n"204205 @staticmethod206 def _normalize(msg: dict) -> dict:207 info = msg.get("info") or {}208 parts = msg.get("parts") or []209 text = "".join(p.get("text", "") for p in parts if p.get("type") == "text")210 tools = [{"tool": p.get("tool"),211 "status": (p.get("state") or {}).get("status"),212 "input": (p.get("state") or {}).get("input"),213 "output": ((p.get("state") or {}).get("output") or "")[:4000]}214 for p in parts if p.get("type") == "tool"]215 out = {"id": info.get("id"), "role": info.get("role"), "text": text, "tools": tools}216 if info.get("role") == "assistant":217 out["tokens"] = info.get("tokens")218 out["finish"] = info.get("finish")219 if info.get("error"):220 out["error"] = info["error"]221 return out222223224harness = Opencode(OPENCODE)225226227def _report(phase: str, detail: str) -> None:228 """Append a phase line to m64's status file (if running under m64)."""229 if path := os.environ.get("M64_STATUS"):230 with open(path, "a") as f:231 f.write(json.dumps({"phase": phase, "detail": detail}) + "\n")232233234async def _report_ready_when_up() -> None:235 while not await harness.healthy():236 await asyncio.sleep(1)237 _report("ready", f"opencode up, model {MODEL} on Basalt")238239240@asynccontextmanager241async def _lifespan(app: FastAPI):242 SESSIONS_DIR.mkdir(parents=True, exist_ok=True)243 task = asyncio.create_task(_report_ready_when_up())244 yield245 task.cancel()246247248app = FastAPI(title="opencode-basalt agent gateway", docs_url=None, redoc_url=None,249 openapi_url=None, lifespan=_lifespan)250251252@app.get("/health")253async def health():254 ok = await harness.healthy()255 return JSONResponse({"status": "ok" if ok else "starting", "harness": "opencode"},256 status_code=200 if ok else 503)257258259@app.get("/v1/agent")260def describe():261 return {262 "harness": "opencode", "llm": MODEL, "llm_provider": "Basalt (local, RTX 5090)",263 "sandbox": {264 "network": f"internal docker network; its only way out is the m64 gate, which lets out "265 f"{NETWORK} and logs every connection. The model runs in the same "266 "container (Basalt on 127.0.0.1:9000)",267 "filesystem": "each session gets a fresh empty directory under the run's scratch",268 "privileges": "non-root (1000:1000), all capabilities dropped, no-new-privileges",269 },270 }271272273@app.post("/v1/sessions", status_code=201)274async def create_session(request: Request):275 body = await _json(request)276 reg = _load_registry()277 d = SESSIONS_DIR / uuid.uuid4().hex[:12]278 d.mkdir(parents=True)279 sid = await harness.create(str(d), body.get("title"))280 reg[sid] = {"id": sid, "title": body.get("title"), "created": time.time(),281 "directory": str(d)}282 _save_registry(reg)283 log.info("SESSION_CREATE %s dir=%s", sid, d)284 return {k: v for k, v in reg[sid].items() if k != "directory"}285286287@app.post("/v1/sessions/{sid}/messages")288async def send_message(sid: str, request: Request):289 body = await _json(request)290 text = body.get("text") or body.get("prompt")291 if not text or not isinstance(text, str):292 raise HTTPException(400, "body needs 'text'")293 s = _session(_load_registry(), sid)294 if body.get("wait", True):295 t0 = time.time()296 reply = await harness.send(sid, s["directory"], text)297 return {**reply, "elapsed_s": round(time.time() - t0, 1)}298 await harness.send_async(sid, s["directory"], text)299 return JSONResponse({"session_id": sid, "status": "busy",300 "events": f"GET /v1/sessions/{sid}/events"}, status_code=202)301302303@app.post("/v1/sessions/{sid}/abort")304async def abort(sid: str):305 s = _session(_load_registry(), sid)306 await harness.abort(sid, s["directory"])307 return {"aborted": sid}308309310@app.get("/v1/sessions/{sid}/events")311async def events(sid: str):312 s = _session(_load_registry(), sid)313 return StreamingResponse(harness.events(sid, s["directory"]), media_type="text/event-stream",314 headers={"Cache-Control": "no-cache"})315316317@app.get("/v1/sessions/{sid}/files")318async def get_files(sid: str, path: str = "."):319 s = _session(_load_registry(), sid)320 root = Path(s["directory"])321 p = _safe_path(root, path)322 kind = _kind(p)323 if kind is None:324 raise HTTPException(404, f"{path} not found")325 if kind == "other":326 raise HTTPException(400, f"{path} is not a regular file or directory")327 if kind == "dir":328 entries = []329 for c in sorted(p.iterdir()):330 k = _kind(c) or "other"331 entries.append({"name": c.name, "type": k,332 "size": c.stat().st_size if k == "file" else None})333 return {"path": str(p.relative_to(root)) if p != root else ".", "entries": entries}334 # The agent may swap the file for a FIFO after the check above: open without335 # blocking and check what was actually opened.336 f = _open_regular(p, os.O_RDONLY)337 size = os.fstat(f.fileno()).st_size338339 def chunks():340 with f:341 while data := f.read(65536):342 yield data343344 headers = {**_DOWNLOAD_HEADERS, "Content-Length": str(size),345 "Content-Disposition": "attachment; filename*=UTF-8''" + quote(p.name)}346 return StreamingResponse(chunks(), media_type=_download_type(p.name), headers=headers)347348349@app.put("/v1/sessions/{sid}/files", status_code=201)350async def put_file(sid: str, request: Request, path: str = ""):351 if not path.strip("/."):352 raise HTTPException(400, "?path= names the file to write")353 s = _session(_load_registry(), sid)354 p = _safe_path(Path(s["directory"]), path)355 if _kind(p) not in (None, "file"):356 raise HTTPException(400, f"{path} exists and is not a regular file")357 try:358 p.parent.mkdir(parents=True, exist_ok=True)359 except (FileExistsError, NotADirectoryError):360 raise HTTPException(400, f"a parent of {path} is not a directory")361 data = await request.body()362 with _open_regular(p, os.O_WRONLY | os.O_CREAT | os.O_TRUNC) as f:363 f.write(data)364 return {"written": path, "bytes": len(data)}365366367async def _json(request: Request) -> dict:368 raw = await request.body()369 if not raw:370 return {}371 try:372 data = json.loads(raw)373 except ValueError:374 raise HTTPException(400, "body must be JSON")375 if not isinstance(data, dict):376 raise HTTPException(400, "body must be a JSON object")377 return data378379380if __name__ == "__main__":381 uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")