m64_9d6a8d0db091 / agent_gateway.py
382 lines, 15.0 KBraw
1"""Agent gateway for the opencode-inside example capsule.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, its only egress is the m64 gate (HTTP(S)_PROXY), which6lets out what the user granted at serve time (this capsule needs openrouter.ai:443, so7--allow-external-network), adds the OpenRouter key itself (the capsule holds only a8placeholder) and logs every connection. So this file keeps just the harness-neutral control API and9drives the local `opencode serve` on 127.0.0.1:4096.1011Routes (the agent.session kind; the m64 front proxies only these and /health):1213 GET /health 200 when opencode answers14 GET /v1/agent what this is15 POST /v1/sessions new session, fresh empty directory16 POST /v1/sessions/{id}/messages {"text": ...}; wait=true (default) blocks for the reply17 GET /v1/sessions/{id}/events SSE stream of this session's events18 GET /v1/sessions/{id}/files[?path=] list a dir / download a file19 PUT /v1/sessions/{id}/files?path= upload a file (raw body)20 POST /v1/sessions/{id}/abort stop the running turn21"""22import asyncio23import json24import logging25import mimetypes26import os27import stat28import time29import uuid30from contextlib import asynccontextmanager31from pathlib import Path32from urllib.parse import quote3334import httpx35import uvicorn36from fastapi import FastAPI, HTTPException, Request37from fastapi.responses import JSONResponse, StreamingResponse3839PORT = int(os.environ.get("M64_PORT", "9500"))40OPENCODE = os.environ.get("OPENCODE_URL", "http://127.0.0.1:4096")41PROVIDER = os.environ.get("AGENT_PROVIDER", "openrouter-compat")42MODEL = os.environ.get("AGENT_MODEL", "deepseek/deepseek-v4.1-flash")4344# What the user granted this run when serving it (m64 serve --allow-external-network /45# --allow-internal-network); the m64 gate enforces it.46_REACH = {47 "": "nothing: no egress at all",48 "external": "the public internet only (never the LAN or tailnet)",49 "internal": "the LAN and tailnet only",50 "external,internal": "the public internet, the LAN and the tailnet",51}52NETWORK = _REACH.get(os.environ.get("M64_NETWORK", ""), os.environ.get("M64_NETWORK", ""))53WORK = Path(os.environ.get("WORK_DIR", "/scratch/work"))54SESSIONS_DIR = WORK / "sessions"55REGISTRY = WORK / ".gateway" / "sessions.json"5657logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")58log = logging.getLogger("agent_gateway")596061# ── Session registry (lives in /scratch, so it goes with the run) ────────────────6263def _load_registry() -> dict:64 try:65 return json.loads(REGISTRY.read_text())66 except Exception:67 return {}686970def _save_registry(reg: dict) -> None:71 REGISTRY.parent.mkdir(parents=True, exist_ok=True)72 tmp = REGISTRY.with_suffix(".tmp")73 tmp.write_text(json.dumps(reg, indent=2))74 tmp.replace(REGISTRY)757677def _session(reg: dict, sid: str) -> dict:78 if sid not in reg:79 raise HTTPException(404, f"unknown session '{sid}'")80 return reg[sid]818283def _safe_path(root: Path, rel: str) -> Path:84 if "\x00" in rel:85 raise HTTPException(400, "invalid path")86 try:87 p = (root / rel.lstrip("/")).resolve()88 except (OSError, ValueError, RuntimeError): # e.g. a symlink loop89 raise HTTPException(400, "invalid path")90 if p != root.resolve() and root.resolve() not in p.parents:91 raise HTTPException(400, "path escapes the session directory")92 return p939495def _kind(p: Path) -> str | None:96 """"file", "dir", "other" (FIFO, socket, device...), or None if it does not exist."""97 try:98 mode = p.stat().st_mode99 except (FileNotFoundError, NotADirectoryError):100 return None101 except OSError:102 return "other"103 return "file" if stat.S_ISREG(mode) else "dir" if stat.S_ISDIR(mode) else "other"104105106# Files are written by the agent, so their bytes are untrusted. They are only ever served107# as downloads, sandboxed, never sniffed, and typed as one of these inert types or bytes.108_INERT_TYPES = {"text/plain", "image/png", "image/jpeg", "image/gif", "image/webp"}109_DOWNLOAD_HEADERS = {"X-Content-Type-Options": "nosniff", "Content-Security-Policy": "sandbox"}110111112def _open_regular(p: Path, flags: int):113 """Open p as a regular file, never blocking on a FIFO or following a final symlink."""114 try:115 fd = os.open(p, flags | os.O_NONBLOCK | os.O_NOFOLLOW, 0o644)116 except OSError:117 raise HTTPException(400, "not a regular file")118 if not stat.S_ISREG(os.fstat(fd).st_mode):119 os.close(fd)120 raise HTTPException(400, "not a regular file")121 os.set_blocking(fd, True)122 return os.fdopen(fd, "rb" if flags == os.O_RDONLY else "wb")123124125def _download_type(name: str) -> str:126 guessed = mimetypes.guess_type(name)[0]127 return guessed if guessed in _INERT_TYPES else "application/octet-stream"128129130# ── opencode translator ─────────────────────────────────────────────────────────131132class Opencode:133 """`opencode serve` HTTP API. Every call carries ?directory= so each session works in134 its own directory under $WORK/sessions."""135136 def __init__(self, base: str):137 self.base = base.rstrip("/")138 # trust_env=False: this is loopback; never route it through HTTPS_PROXY.139 self.http = httpx.AsyncClient(base_url=self.base, timeout=httpx.Timeout(30, read=None),140 trust_env=False)141142 async def healthy(self) -> bool:143 try:144 return (await self.http.get("/global/health", timeout=3)).status_code == 200145 except Exception:146 return False147148 async def create(self, directory: str, title: str | None) -> str:149 r = await self.http.post("/session", params={"directory": directory},150 json={"title": title} if title else {})151 r.raise_for_status()152 return r.json()["id"]153154 @staticmethod155 def _prompt_body(text: str) -> dict:156 return {"parts": [{"type": "text", "text": text}],157 "model": {"providerID": PROVIDER, "modelID": MODEL}}158159 async def send(self, sid: str, directory: str, text: str) -> dict:160 r = await self.http.post(f"/session/{sid}/message", params={"directory": directory},161 json=self._prompt_body(text))162 r.raise_for_status()163 reply = self._normalize(r.json())164 # One turn = several assistant messages (one per tool step); the response is only165 # the last. Gather the tool calls of the whole turn from the history.166 turn = []167 for m in reversed(await self.history(sid, directory)):168 if m["role"] == "user":169 break170 turn.append(m)171 reply["tools"] = [t for m in reversed(turn) for t in m["tools"]]172 reply["steps"] = len(turn)173 return reply174175 async def send_async(self, sid: str, directory: str, text: str) -> None:176 r = await self.http.post(f"/session/{sid}/prompt_async", params={"directory": directory},177 json=self._prompt_body(text))178 r.raise_for_status()179180 async def history(self, sid: str, directory: str) -> list:181 r = await self.http.get(f"/session/{sid}/message", params={"directory": directory})182 r.raise_for_status()183 return [self._normalize(m) for m in r.json()]184185 async def abort(self, sid: str, directory: str) -> None:186 await self.http.post(f"/session/{sid}/abort", params={"directory": directory})187188 async def events(self, sid: str, directory: str):189 """Yield SSE frames for events that belong to `sid`."""190 async with self.http.stream("GET", "/event", params={"directory": directory}) as r:191 async for line in r.aiter_lines():192 if not line.startswith("data:"):193 continue194 try:195 ev = json.loads(line[5:])196 except ValueError:197 continue198 props = ev.get("properties") or {}199 owner = (props.get("sessionID")200 or (props.get("info") or {}).get("sessionID")201 or (props.get("part") or {}).get("sessionID")202 or (props.get("info") or {}).get("id"))203 if owner == sid or ev.get("type") == "server.connected":204 yield f"data: {json.dumps(ev)}\n\n"205206 @staticmethod207 def _normalize(msg: dict) -> dict:208 info = msg.get("info") or {}209 parts = msg.get("parts") or []210 text = "".join(p.get("text", "") for p in parts if p.get("type") == "text")211 tools = [{"tool": p.get("tool"),212 "status": (p.get("state") or {}).get("status"),213 "input": (p.get("state") or {}).get("input"),214 "output": ((p.get("state") or {}).get("output") or "")[:4000]}215 for p in parts if p.get("type") == "tool"]216 out = {"id": info.get("id"), "role": info.get("role"), "text": text, "tools": tools}217 if info.get("role") == "assistant":218 out["tokens"] = info.get("tokens")219 out["finish"] = info.get("finish")220 if info.get("error"):221 out["error"] = info["error"]222 return out223224225harness = Opencode(OPENCODE)226227228def _report(phase: str, detail: str) -> None:229 """Append a phase line to m64's status file (if running under m64)."""230 if path := os.environ.get("M64_STATUS"):231 with open(path, "a") as f:232 f.write(json.dumps({"phase": phase, "detail": detail}) + "\n")233234235async def _report_ready_when_up() -> None:236 while not await harness.healthy():237 await asyncio.sleep(1)238 _report("ready", f"opencode up, model {MODEL} via OpenRouter")239240241@asynccontextmanager242async def _lifespan(app: FastAPI):243 SESSIONS_DIR.mkdir(parents=True, exist_ok=True)244 task = asyncio.create_task(_report_ready_when_up())245 yield246 task.cancel()247248249app = FastAPI(title="opencode-inside agent gateway", docs_url=None, redoc_url=None,250 openapi_url=None, lifespan=_lifespan)251252253@app.get("/health")254async def health():255 ok = await harness.healthy()256 return JSONResponse({"status": "ok" if ok else "starting", "harness": "opencode"},257 status_code=200 if ok else 503)258259260@app.get("/v1/agent")261def describe():262 return {263 "harness": "opencode", "llm": MODEL, "llm_provider": "OpenRouter",264 "sandbox": {265 "network": f"internal docker network; its only way out is the m64 gate, which lets out "266 f"{NETWORK}, adds the OpenRouter key itself (the agent never holds it) "267 "and logs every connection",268 "filesystem": "each session gets a fresh empty directory under the run's scratch",269 "privileges": "non-root (1000:1000), all capabilities dropped, no-new-privileges",270 },271 }272273274@app.post("/v1/sessions", status_code=201)275async def create_session(request: Request):276 body = await _json(request)277 reg = _load_registry()278 d = SESSIONS_DIR / uuid.uuid4().hex[:12]279 d.mkdir(parents=True)280 sid = await harness.create(str(d), body.get("title"))281 reg[sid] = {"id": sid, "title": body.get("title"), "created": time.time(),282 "directory": str(d)}283 _save_registry(reg)284 log.info("SESSION_CREATE %s dir=%s", sid, d)285 return {k: v for k, v in reg[sid].items() if k != "directory"}286287288@app.post("/v1/sessions/{sid}/messages")289async def send_message(sid: str, request: Request):290 body = await _json(request)291 text = body.get("text") or body.get("prompt")292 if not text or not isinstance(text, str):293 raise HTTPException(400, "body needs 'text'")294 s = _session(_load_registry(), sid)295 if body.get("wait", True):296 t0 = time.time()297 reply = await harness.send(sid, s["directory"], text)298 return {**reply, "elapsed_s": round(time.time() - t0, 1)}299 await harness.send_async(sid, s["directory"], text)300 return JSONResponse({"session_id": sid, "status": "busy",301 "events": f"GET /v1/sessions/{sid}/events"}, status_code=202)302303304@app.post("/v1/sessions/{sid}/abort")305async def abort(sid: str):306 s = _session(_load_registry(), sid)307 await harness.abort(sid, s["directory"])308 return {"aborted": sid}309310311@app.get("/v1/sessions/{sid}/events")312async def events(sid: str):313 s = _session(_load_registry(), sid)314 return StreamingResponse(harness.events(sid, s["directory"]), media_type="text/event-stream",315 headers={"Cache-Control": "no-cache"})316317318@app.get("/v1/sessions/{sid}/files")319async def get_files(sid: str, path: str = "."):320 s = _session(_load_registry(), sid)321 root = Path(s["directory"])322 p = _safe_path(root, path)323 kind = _kind(p)324 if kind is None:325 raise HTTPException(404, f"{path} not found")326 if kind == "other":327 raise HTTPException(400, f"{path} is not a regular file or directory")328 if kind == "dir":329 entries = []330 for c in sorted(p.iterdir()):331 k = _kind(c) or "other"332 entries.append({"name": c.name, "type": k,333 "size": c.stat().st_size if k == "file" else None})334 return {"path": str(p.relative_to(root)) if p != root else ".", "entries": entries}335 # The agent may swap the file for a FIFO after the check above: open without336 # blocking and check what was actually opened.337 f = _open_regular(p, os.O_RDONLY)338 size = os.fstat(f.fileno()).st_size339340 def chunks():341 with f:342 while data := f.read(65536):343 yield data344345 headers = {**_DOWNLOAD_HEADERS, "Content-Length": str(size),346 "Content-Disposition": "attachment; filename*=UTF-8''" + quote(p.name)}347 return StreamingResponse(chunks(), media_type=_download_type(p.name), headers=headers)348349350@app.put("/v1/sessions/{sid}/files", status_code=201)351async def put_file(sid: str, request: Request, path: str = ""):352 if not path.strip("/."):353 raise HTTPException(400, "?path= names the file to write")354 s = _session(_load_registry(), sid)355 p = _safe_path(Path(s["directory"]), path)356 if _kind(p) not in (None, "file"):357 raise HTTPException(400, f"{path} exists and is not a regular file")358 try:359 p.parent.mkdir(parents=True, exist_ok=True)360 except (FileExistsError, NotADirectoryError):361 raise HTTPException(400, f"a parent of {path} is not a directory")362 data = await request.body()363 with _open_regular(p, os.O_WRONLY | os.O_CREAT | os.O_TRUNC) as f:364 f.write(data)365 return {"written": path, "bytes": len(data)}366367368async def _json(request: Request) -> dict:369 raw = await request.body()370 if not raw:371 return {}372 try:373 data = json.loads(raw)374 except ValueError:375 raise HTTPException(400, "body must be JSON")376 if not isinstance(data, dict):377 raise HTTPException(400, "body must be a JSON object")378 return data379380381if __name__ == "__main__":382 uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")