"""RosDiff Stage 1 API (Orcanion internal, testers only). uvicorn api.main:app --host 0.0.0.0 --port 8000 Every endpoint except /health needs a tester key in the ``X-API-Key`` header (or ``Authorization: Bearer ``). """ from __future__ import annotations import datetime as dt import json import logging import shutil import time from collections.abc import Callable from dataclasses import asdict, dataclass from pathlib import Path from typing import Literal from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.responses import HTMLResponse, Response from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from rag.store import RagStore from . import budget from .auth import Allowlist, Tester from .codegen import DirectBackend, OpenCodeBackend, generate_code, mentioned_interfaces from .config import Settings, get_settings from .enhance import Enhanced from .enhance import enhance as enhance_prompt from .enhance import record as record_enhancement from .jobs import JobError, compute_usage, month_spend, refresh, start_isaac, start_review from .jobs import submit as submit_job from .llm import LLM, ChatClient, LLMError from .logstore import LogStore from .mjcf_check import validate_mjcf from .opencode_runner import OpenCodeWorkspace from .robots import ROBOT_NAMES, ROBOTS from .ros_check import load_interfaces from .routing import ModelSpec, all_models, is_free, plan_code, plan_scene from .runs import check_token, new_run_id, render_report, report_token from .sim_backend import SimBackend, make_backend from .simulate import generate_scene from .storage import R2Config, R2Storage, Storage from .zen import WithFreeModels, ZenLLM log = logging.getLogger("rosdiff") @dataclass class Services: settings: Settings store: RagStore logs: LogStore allowlist: Allowlist llm_factory: Callable[[], LLM] # scene generation backend: SimBackend | None = None # None: /simulate/run is disabled storage: Storage | None = None code_llm_factory: Callable[[], LLM] | None = None # direct ROS 2 code backend; None: same as llm_factory enhancer: Callable[..., Enhanced] = enhance_prompt # free prompt enhancement through OpenCode Zen free_llm_factory: Callable[[], LLM] | None = None # free debugging model (OpenCode Zen); None: paid only def default_services() -> Services: settings = get_settings() r2 = R2Config.from_env() if not settings.report_secret: log.warning("REPORT_SECRET is not set: report links will stop working when the server restarts") return Services( settings=settings, store=RagStore(settings.data_dir / "chroma"), logs=LogStore(settings.log_db), allowlist=Allowlist(settings.testers_file), llm_factory=lambda: ChatClient(settings), code_llm_factory=lambda: ChatClient(settings, model=settings.code_llm_model), free_llm_factory=lambda: ZenLLM(settings), backend=make_backend(settings), storage=R2Storage(r2) if r2 else None, ) # --------------------------------------------------------------------------- request/response models class SimulateRequest(BaseModel): description: str = Field(..., min_length=10, max_length=4000, description="The scenario, in plain language") robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Robot to include, or null for none") confirm_cost: bool = Field(False, description="Run even if the estimate is above CONFIRM_ABOVE_INR") enhance: bool = Field(True, description="Sharpen the request first with a free OpenCode Zen model") class SimulateResponse(BaseModel): id: str valid: bool mjcf: str | None robot: str | None validation: dict attempts: int sources: list[str] model: str how_to_run: str routing: dict | None = None # what RosDiff chose and why (tier, reason, simulator) enhanced_prompt: dict | None = None # {used, text, model, note}: the brief the generator actually got cost: dict | None = None # rupees this task used, the estimate, and the remaining budget class CodeRequest(BaseModel): request: str = Field(..., min_length=10, max_length=4000, description="What the ROS 2 package should do") distro: Literal["humble", "jazzy"] = "humble" language: Literal["python", "cpp"] = "python" package_name: str | None = Field(None, pattern=r"^[a-z][a-z0-9_]{1,62}$") check: bool = Field( False, description="Also build, lint, test and launch it in the ROS 2 sandbox (see /code/check)" ) confirm_cost: bool = Field(False, description="Run even if the estimate is above CONFIRM_ABOVE_INR") enhance: bool = Field(True, description="Sharpen the request first with a free OpenCode Zen model") class CodeResponse(BaseModel): id: str valid: bool files: dict[str, str] validation: dict attempts: int sources: list[str] backend: str model: str how_to_build: str check: dict | None = None # the sandbox job, when requested: {id, status, report_url} or {error} routing: dict | None = None cost: dict | None = None enhanced_prompt: dict | None = None class RunRequest(BaseModel): generation_id: str | None = Field(None, description="A valid scene from /simulate/generate (yours)") mjcf: str | None = Field(None, max_length=300_000, description="Or: an MJCF scene to validate and run") robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Only used with mjcf") duration_s: float = Field(10.0, gt=0, le=60) controller: Literal["hold", "passive"] = Field("hold", description="hold: keep the starting pose; passive: limp") class RunResponse(BaseModel): run_id: str status: str report_url: str status_url: str class RunStatus(BaseModel): run_id: str status: str error: str | None summary: dict | None report_url: str class SweepRequest(BaseModel): run_id: str | None = Field(None, description="Sweep the scene of one of your runs (linked on its report)") generation_id: str | None = Field(None, description="Or: a valid scene from /simulate/generate") mjcf: str | None = Field(None, max_length=300_000, description="Or: an MJCF scene to validate and sweep") robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Only used with mjcf") n_worlds: int = Field(64, ge=1, le=256, description="Randomized copies, simulated in parallel") duration_s: float = Field(5.0, gt=0, le=10) controller: Literal["hold", "passive"] = "hold" push_max_n: float = Field(80.0, ge=0, le=1000, description="Each world is pushed with 0..this many newtons") friction: tuple[float, float] = Field((0.5, 1.5), description="Friction scale range") object_mass: tuple[float, float] = Field((0.5, 2.0), description="Mass scale range for free objects") seed: int = 0 class ExportRequest(BaseModel): run_ids: list[str] = Field(..., min_length=1, max_length=100, description="Your finished runs, one robot") name: str = Field(..., pattern=r"^[a-z0-9][a-z0-9-]{1,60}$", description="Dataset name, e.g. g1-table-pushes") private: bool = True class FeedbackItem(BaseModel): target_type: Literal["generation", "job"] target_id: str rating: Literal["helpful", "not_helpful"] note: str = Field("", max_length=4000) class CheckRequest(BaseModel): generation_id: str = Field(..., description="A package from /code/generate (yours)") class JobResponse(BaseModel): id: str kind: str status: str report_url: str status_url: str class FeedbackRequest(BaseModel): rating: Literal["works", "partly", "broken"] note: str = Field("", max_length=4000) def _how_to_run(robot: str | None) -> str: if not robot: return "Save as scene.xml and run: python -m mujoco.viewer --mjcf scene.xml" spec = ROBOTS[robot] return ( f"Save as scene_rosdiff.xml inside mujoco_menagerie/{spec.directory}/ (next to {spec.model_file}) and run: " f"python -m mujoco.viewer --mjcf mujoco_menagerie/{spec.directory}/scene_rosdiff.xml" ) def _check_budget(s: Services, tester: str, plan, confirmed: bool) -> None: try: budget.check(s.logs, s.settings, tester, plan.estimate_inr, confirmed) except budget.BudgetError as e: raise HTTPException(status_code=e.status, detail={"message": str(e), **e.detail}) from e def _enhance(s: Services, wanted: bool, kind: str, text: str, **ctx) -> Enhanced: if not wanted: return Enhanced(text, text, False, note="not requested") try: return s.enhancer(s.settings, kind, text, **ctx) except Exception as e: # enhancement is a nicety: never fail a generation because of it log.warning("prompt enhancement crashed: %s", e) return Enhanced(text, text, False, note="enhancer error") def _enhanced_view(e: Enhanced) -> dict: return {"used": e.used, "text": e.text if e.used else None, "model": e.model, "note": e.note, "ms": e.ms} def _spec(model_id: str | None) -> ModelSpec: if is_free(model_id): return ModelSpec(model_id, "free", "OpenCode Zen", 0.0, 0.0, free=True) return all_models().get(model_id) or ModelSpec(model_id or "unknown", "?", "unknown", 0.0, 0.0) def _llm(s: Services, factory: Callable[[], LLM]) -> LLM: """The paid client, with `opencode/...` (free debugging) models routed to OpenCode Zen.""" paid = factory() return WithFreeModels(paid, s.free_llm_factory()) if s.free_llm_factory else paid def _book(s: Services, tester: str, gen_id: str, plan, attempts, brief: Enhanced | None = None) -> dict: """Record every model call of a generation in the rupee ledger; return what this task cost.""" before = budget.summary(s.logs, s.settings, tester)["used"] if brief is not None: record_enhancement(s.logs, tester, gen_id, plan.category, brief) for i, a in enumerate(attempts): spec = _spec(a.model) budget.record_call( s.logs, s.settings, tester=tester, generation_id=gen_id, category=plan.category if i == 0 else "debugging", model=spec, usage=a.usage, estimate_inr=plan.estimate_inr[0] if i == 0 else None, ) after = budget.summary(s.logs, s.settings, tester) return { "inr": round(after["used"] - before, 2), "estimate_inr": list(plan.estimate_inr), "budget": {"used": after["used"], "remaining": after["remaining"], "allocation": after["allocation"]}, } def _routing_view(plan, attempts) -> dict: tiers = [_spec(a.model).tier for a in attempts] view = { "complexity": plan.complexity, "tier": plan.tier, "reason": plan.reason, "tiers_used": tiers, "escalated": len({t for t in tiers if t != "free"}) > 1, "free_debug": "free" in tiers, "models": [a.model for a in attempts], } if plan.simulator: view.update(simulator=plan.simulator, simulator_reason=plan.simulator_reason) return view def _submit(s: Services, tester: str, kind: str, job_input: dict, **kw) -> dict: try: return submit_job(s, tester, kind, job_input, **kw) except JobError as e: raise HTTPException(status_code=e.status, detail=str(e)) from e # --------------------------------------------------------------------------- app def create_app(services: Services | None = None) -> FastAPI: app = FastAPI( title="RosDiff Stage 1", version="0.1.0", description="Internal tester build: RAG-grounded MJCF and ROS 2 code generation.", ) app.state.services = services app.mount("/static", StaticFiles(directory=Path(__file__).with_name("static")), name="static") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) def svc() -> Services: if app.state.services is None: app.state.services = default_services() return app.state.services def tester( x_api_key: str | None = Header(None), authorization: str | None = Header(None), s: Services = Depends(svc), ) -> Tester: key = x_api_key if not key and authorization and authorization.lower().startswith("bearer "): key = authorization[7:] if s.settings.owner_api_key and key == s.settings.owner_api_key: return Tester("owner", "", active=True) t = s.allowlist.identify(key) if t is None: raise HTTPException(status_code=401, detail="missing or unknown API key") return t @app.on_event("startup") def _startup_checks(): s = svc() if not s.settings.llm_api_key: log.warning("LLM_API_KEY is not set: generation endpoints will fail") if not s.settings.report_secret: log.warning("REPORT_SECRET is not set: report links will break on restart") @app.get("/health") def health(s: Services = Depends(svc)) -> dict: return { "status": "ok", "model": s.settings.llm_model, "code_model": s.settings.code_llm_model, "llm_configured": bool(s.settings.llm_api_key), "codegen_backend": s.settings.codegen_backend, "opencode_found": shutil.which(s.settings.opencode_bin) is not None, "prompt_enhance": s.settings.prompt_enhance and shutil.which(s.settings.opencode_bin) is not None, "enhance_model": s.settings.enhance_model, "gpu_share_usd": s.settings.compute_share_usd or None, "interface_catalogue": s.settings.interfaces_file.is_file(), "index": {name: s.store.count(name) for name in ("mujoco", "ros2")}, "robots": sorted(ROBOTS), "runs_enabled": s.backend is not None and s.storage is not None, "sim_backend": s.backend.name if s.backend else None, } @app.post("/simulate/generate", response_model=SimulateResponse) def simulate(req: SimulateRequest, t: Tester = Depends(tester), s: Services = Depends(svc)) -> SimulateResponse: start = time.monotonic() plan = plan_scene(req.description, req.robot, s.settings) # routed on the tester's own words _check_budget(s, t.name, plan, req.confirm_cost) brief = _enhance(s, req.enhance, "scene", req.description, robot=req.robot) logged_request = { **req.model_dump(exclude={"confirm_cost"}), "routing": plan.to_dict(), "enhanced": _enhanced_view(brief), } try: result = generate_scene(brief.text, req.robot, _llm(s, s.llm_factory), s.store, s.settings, plan=plan) except LLMError as e: gen_id = s.logs.record( tester=t.name, kind="simulate", request=logged_request, error=str(e), duration_ms=int((time.monotonic() - start) * 1000), ) record_enhancement(s.logs, t.name, gen_id, "simulation_ai", brief) raise HTTPException(status_code=502, detail=f"model call failed (log id {gen_id}): {e}") from e gen_id = s.logs.record( tester=t.name, kind="simulate", request=logged_request, model=result.model, sources=result.sources, attempts=[asdict(a) for a in result.attempts], output=result.mjcf, valid=result.valid, validation=result.validation.to_dict(), duration_ms=int((time.monotonic() - start) * 1000), ) return SimulateResponse( id=gen_id, valid=result.valid, mjcf=result.mjcf, robot=req.robot, validation=result.validation.to_dict(), attempts=len(result.attempts), sources=result.sources, model=result.model, how_to_run=_how_to_run(req.robot), routing=_routing_view(plan, result.attempts), cost=_book(s, t.name, gen_id, plan, result.attempts, brief), enhanced_prompt=_enhanced_view(brief), ) @app.post("/code/generate", response_model=CodeResponse) def code( req: CodeRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc) ) -> CodeResponse: start = time.monotonic() catalogue = load_interfaces(str(s.settings.interfaces_file)).get(req.distro, {}) plan = plan_code( req.request, req.distro, req.language, s.settings, len(mentioned_interfaces(req.request, catalogue)) ) _check_budget(s, t.name, plan, req.confirm_cost) brief = _enhance(s, req.enhance, "code", req.request, distro=req.distro, language=req.language) logged_request = { **req.model_dump(exclude={"confirm_cost"}), "routing": plan.to_dict(), "enhanced": _enhanced_view(brief), } try: if s.settings.codegen_backend == "opencode": with OpenCodeWorkspace(s.settings) as ws: result = generate_code( brief.text, req.distro, req.language, OpenCodeBackend(ws), s.store, s.settings, req.package_name, plan=plan, ) else: result = generate_code( brief.text, req.distro, req.language, DirectBackend(_llm(s, s.code_llm_factory or s.llm_factory)), s.store, s.settings, req.package_name, plan=plan, ) except (LLMError, Exception) as e: if not isinstance(e, LLMError): log.exception("code generation crashed") gen_id = s.logs.record( tester=t.name, kind="code", request=logged_request, error=str(e), duration_ms=int((time.monotonic() - start) * 1000), ) record_enhancement(s.logs, t.name, gen_id, "coding", brief) raise HTTPException(status_code=502, detail=f"generation failed (log id {gen_id}): {e}") from e gen_id = s.logs.record( tester=t.name, kind="code", request=logged_request, model=f"{result.model} via {result.backend}", sources=result.sources, attempts=[asdict(a) for a in result.attempts], output=json.dumps(result.files), valid=result.valid, validation=result.validation.to_dict(), duration_ms=int((time.monotonic() - start) * 1000), ) check = None if req.check and result.files: try: job = _start_check(s, t.name, gen_id, result.files, req.distro, req.language) check = {"id": job["id"], "status": job["status"], "report_url": _report_url(request, s, job["id"])} except HTTPException as e: check = {"error": e.detail} return CodeResponse( id=gen_id, check=check, routing=_routing_view(plan, result.attempts), cost=_book(s, t.name, gen_id, plan, result.attempts, brief), enhanced_prompt=_enhanced_view(brief), valid=result.valid, files=result.files, validation=result.validation.to_dict(), attempts=len(result.attempts), sources=result.sources, backend=result.backend, model=result.model, how_to_build=( f"Write the files into /src/, then: source /opt/ros/{req.distro}/setup.bash && cd && " "colcon build && source install/setup.bash && ros2 launch " ), ) def _report_url(request: Request, s: Services, run_id: str) -> str: base = s.settings.public_base_url or str(request.base_url).rstrip("/") return f"{base}/report/{run_id}?token={report_token(s.settings.report_secret, run_id)}" def _job_response(request: Request, s: Services, job: dict) -> JobResponse: base = s.settings.public_base_url or str(request.base_url).rstrip("/") return JobResponse( id=job["id"], kind=job["kind"], status=job["status"], report_url=_report_url(request, s, job["id"]), status_url=f"{base}/jobs/{job['id']}", ) def _start_check(s: Services, tester: str, gen_id: str, files: dict, distro: str, language: str) -> dict: return _submit( s, tester, "ros_check", {"files": files, "distro": distro}, run_id=new_run_id(), generation_id=gen_id, params={"distro": distro, "language": language}, ) def _scene_for(s: Services, t: Tester, run_id, generation_id, mjcf, robot) -> tuple[str, str | None, str | None]: """(mjcf, robot, generation_id) from exactly one of a run, a generation, or raw MJCF (validated here).""" if sum(x is not None for x in (run_id, generation_id, mjcf)) != 1: raise HTTPException(status_code=422, detail="give exactly one of run_id, generation_id or mjcf") if run_id: run = s.logs.get_run(run_id) if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": raise HTTPException(status_code=404, detail="no simulation run with that id for this key") if run.get("generation_id"): generation_id = run["generation_id"] else: key = ((run.get("output") or {}).get("files") or {}).get("scene") text = s.storage.read_text(key) if key and s.storage else None if not text: raise HTTPException(status_code=422, detail="that run's scene is not available (not finished?)") return text, run.get("robot"), None if generation_id: gen = s.logs.get(generation_id) if not gen or gen["tester"] != t.name or gen["kind"] != "simulate": raise HTTPException(status_code=404, detail="no scene generation with that id for this key") if not gen["valid"] or not gen["output"]: raise HTTPException(status_code=422, detail="that generation did not produce a valid scene") return gen["output"], gen["request"].get("robot"), generation_id check = validate_mjcf( mjcf, robot, s.settings.menagerie_dir, s.settings.mjcf_sim_seconds, s.settings.validate_timeout ) if not check.valid: raise HTTPException( status_code=422, detail={"message": "scene failed validation", "validation": check.to_dict()} ) return mjcf, robot, None @app.post("/simulate/sweep", response_model=JobResponse) def simulate_sweep(req: SweepRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): """Run a scene many times with randomized friction, masses and pushes (MuJoCo Warp on the GPU).""" mjcf, robot, gen_id = _scene_for(s, t, req.run_id, req.generation_id, req.mjcf, req.robot) params = { "n_worlds": req.n_worlds, "duration_s": req.duration_s, "controller": req.controller, "push_n": [0.0, req.push_max_n], "friction": list(req.friction), "object_mass": list(req.object_mass), "seed": req.seed, } job = _submit( s, t.name, "sweep", {"mjcf": mjcf, "robot": robot, **params}, run_id=new_run_id(), parent_id=req.run_id, generation_id=gen_id, robot=robot, duration_s=req.duration_s, controller=req.controller, params=params, ) return _job_response(request, s, job) @app.post("/simulate/run/{run_id}/review", response_model=JobResponse) def review_run(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): """Ask Cosmos Reason to watch the run's video (done automatically when AUTO_REVIEW is on).""" run = s.logs.get_run(run_id) if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": raise HTTPException(status_code=404, detail="no simulation run with that id for this key") run = refresh(s, run) try: job = start_review(s, run) except JobError as e: raise HTTPException(status_code=e.status, detail=str(e)) from e return _job_response(request, s, job) @app.post("/datasets/export", response_model=JobResponse) def export_dataset(req: ExportRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): """Turn finished runs into a LeRobot dataset (state, velocity, action, video, task) on the Hugging Face Hub.""" if s.storage is None: raise HTTPException(status_code=503, detail="storage is not configured (R2_* variables)") episodes, robots = [], set() for run_id in dict.fromkeys(req.run_ids): run = s.logs.get_run(run_id) if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": raise HTTPException(status_code=404, detail=f"no simulation run {run_id} for this key") run = refresh(s, run) files = (run.get("output") or {}).get("files") or {} if run["status"] != "COMPLETED" or "trajectory" not in files: raise HTTPException(status_code=422, detail=f"run {run_id} is not finished or has no trajectory") robots.add(run.get("robot")) task = "" if run.get("generation_id"): task = ((s.logs.get(run["generation_id"]) or {}).get("request") or {}).get("description", "") episodes.append( { "run_id": run_id, "task": task, "trajectory_url": s.storage.url(files["trajectory"]), "video_url": s.storage.url(files["video"]) if "video" in files else None, } ) if len(robots) > 1: raise HTTPException(status_code=422, detail="all runs in one dataset must use the same robot") robot = robots.pop() repo_id = f"{s.settings.hf_namespace}/rosdiff-{req.name}" job = _submit( s, t.name, "export", {"repo_id": repo_id, "robot": robot, "episodes": episodes, "private": req.private}, run_id=new_run_id(), robot=robot, params={"repo_id": repo_id, "runs": [e["run_id"] for e in episodes]}, ) return _job_response(request, s, job) @app.post("/simulate/run/{run_id}/isaac", response_model=JobResponse) def isaac_render(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): """Re-render a finished run photorealistically with Isaac Sim (limited to ISAAC_DAILY_LIMIT per day).""" run = s.logs.get_run(run_id) if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": raise HTTPException(status_code=404, detail="no simulation run with that id for this key") run = refresh(s, run) try: job = start_isaac(s, run) except JobError as e: raise HTTPException(status_code=e.status, detail=str(e)) from e return _job_response(request, s, job) @app.post("/code/check", response_model=JobResponse) def code_check(req: CheckRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): """Build, lint, test and launch a generated package in the ROS 2 sandbox. Returns at once; see report_url.""" gen = s.logs.get(req.generation_id) if not gen or gen["tester"] != t.name or gen["kind"] != "code": raise HTTPException(status_code=404, detail="no code generation with that id for this key") files = json.loads(gen["output"]) if gen.get("output") else {} if not files: raise HTTPException(status_code=422, detail="that generation produced no files") job = _start_check( s, t.name, req.generation_id, files, gen["request"]["distro"], gen["request"].get("language", "python") ) return _job_response(request, s, job) @app.post("/simulate/run", response_model=RunResponse) def simulate_run( req: RunRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc) ) -> RunResponse: """Run a validated scene on a cloud GPU. Returns immediately; follow report_url.""" if (req.generation_id is None) == (req.mjcf is None): raise HTTPException(status_code=422, detail="give exactly one of generation_id or mjcf") if req.duration_s > s.settings.run_max_duration_s: raise HTTPException(status_code=422, detail=f"duration_s must be <= {s.settings.run_max_duration_s}") if req.generation_id: gen = s.logs.get(req.generation_id) if not gen or gen["tester"] != t.name or gen["kind"] != "simulate": raise HTTPException(status_code=404, detail="no scene generation with that id for this key") if not gen["valid"] or not gen["output"]: raise HTTPException(status_code=422, detail="that generation did not produce a valid scene") mjcf, robot = gen["output"], gen["request"].get("robot") else: mjcf, robot = req.mjcf, req.robot check = validate_mjcf( mjcf, robot, s.settings.menagerie_dir, s.settings.mjcf_sim_seconds, s.settings.validate_timeout ) if not check.valid: raise HTTPException( status_code=422, detail={"message": "scene failed validation", "validation": check.to_dict()} ) # Usage visibility, not a cap: warn loudly when a key starts runs in a burst. since = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(minutes=10)).isoformat(timespec="seconds") recent = s.logs.runs_since(t.name, since) if recent >= s.settings.run_burst_warn: log.warning("tester %s has started %d runs in the last 10 minutes", t.name, recent + 1) run_id = new_run_id() job_input = {"mjcf": mjcf, "robot": robot, "duration_s": req.duration_s, "controller": req.controller} job = _submit( s, t.name, "sim", job_input, run_id=run_id, generation_id=req.generation_id, robot=robot, duration_s=req.duration_s, controller=req.controller, ) base = s.settings.public_base_url or str(request.base_url).rstrip("/") return RunResponse( run_id=run_id, status=job["status"], report_url=_report_url(request, s, run_id), status_url=f"{base}/simulate/run/{run_id}", ) @app.get("/simulate/run/{run_id}", response_model=RunStatus) def run_status(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): run = s.logs.get_run(run_id) if not run or run["tester"] != t.name: raise HTTPException(status_code=404, detail="no run with that id for this key") run = refresh(s, run) return RunStatus( run_id=run_id, status=run["status"], error=run.get("error"), summary=(run.get("output") or {}).get("summary"), report_url=_report_url(request, s, run_id), ) @app.get("/report/{run_id}", response_class=HTMLResponse) @app.get("/simulate/report/{run_id}", response_class=HTMLResponse) def run_report(run_id: str, request: Request, token: str | None = None, s: Services = Depends(svc)): """Browser page for any job; its link carries its own token, so no API key is needed.""" run = s.logs.get_run(run_id) if not run or not check_token(s.settings.report_secret, run_id, token): raise HTTPException(status_code=404, detail="report not found") run = refresh(s, run) children = [refresh(s, c) for c in s.logs.children(run_id)] return HTMLResponse( render_report(run, s.storage, s.settings.foxglove_org, children, link=lambda i: _report_url(request, s, i)), headers={"Cache-Control": "no-store", "Referrer-Policy": "no-referrer"}, ) @app.get("/jobs") def my_jobs(request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: """Your recent jobs, and how much of this month's compute budget is used (all testers together).""" jobs = s.logs.list_runs(t.name, 30) return { "budget": { "month_spend_usd_estimate": month_spend(s.settings, s.logs), "monthly_limit_usd": s.settings.monthly_credit_usd or None, "mine": compute_usage(s.settings, s.logs, t.name), }, "jobs": [ { "id": j["id"], "kind": j.get("kind") or "sim", "status": j["status"], "created_at": j["created_at"], "parent_id": j.get("parent_id"), "report_url": _report_url(request, s, j["id"]), } for j in jobs ], } @app.get("/jobs/{job_id}") def job_status(job_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: job = s.logs.get_run(job_id) if not job or job["tester"] != t.name: raise HTTPException(status_code=404, detail="no job with that id for this key") job = refresh(s, job) output = job.get("output") or {} return { "id": job_id, "kind": job.get("kind") or "sim", "status": job["status"], "error": job.get("error"), "parent_id": job.get("parent_id"), "summary": output.get("summary"), "files": sorted((output.get("files") or {}).keys()), "report_url": _report_url(request, s, job_id), "media": _media(s, job), } def _media(s: Services, job: dict) -> dict: """Presigned links for what the web app shows inline (videos), once a job has finished.""" files = (job.get("output") or {}).get("files") or {} if job["status"] != "COMPLETED" or s.storage is None: return {} names = {"video": "video", "video_webm": "video_webm", "isaac.mp4": "isaac_video"} return {label: s.storage.url(files[key]) for key, label in names.items() if key in files} @app.get("/me/history") def history(t: Tester = Depends(tester), s: Services = Depends(svc), limit: int = 60) -> dict: """Your recent scenes and packages, newest first (for the web app's sidebar).""" items = [] for g in s.logs.list(tester=t.name, limit=min(max(limit, 1), 200)): req = g.get("request") or {} text = req.get("description") or req.get("request") or "" items.append( { "id": g["id"], "kind": g["kind"], "created_at": g["created_at"], "valid": bool(g["valid"]) if g.get("valid") is not None else None, "title": text[:90], "robot": req.get("robot"), "distro": req.get("distro"), "language": req.get("language"), "rating": g.get("tester_rating"), } ) return {"tester": t.name, "items": items} @app.get("/generations/{gen_id}") def generation(gen_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: """One of your scenes or packages, with its runs and checks (to reopen it in the web app).""" g = s.logs.get(gen_id) if not g or g["tester"] != t.name: raise HTTPException(status_code=404, detail="no generation with that id for this key") req = g.get("request") or {} output = g.get("output") if g["kind"] == "code": output = json.loads(output) if output else {} jobs = [ { "id": j["id"], "kind": j.get("kind") or "sim", "status": j["status"], "created_at": j["created_at"], "report_url": _report_url(request, s, j["id"]), } for j in s.logs.list_runs(t.name, 200) if j.get("generation_id") == gen_id ] return { "id": gen_id, "kind": g["kind"], "created_at": g["created_at"], "request": req, "valid": bool(g["valid"]) if g.get("valid") is not None else None, "validation": g.get("validation") or {}, "output": output, "error": g.get("error"), "model": g.get("model"), "rating": g.get("tester_rating"), "jobs": list(reversed(jobs)), } @app.get("/code/{gen_id}/zip") def code_zip(gen_id: str, t: Tester = Depends(tester), s: Services = Depends(svc)) -> Response: """The generated package as a .zip, ready to unpack into a colcon workspace's src/.""" import io import zipfile g = s.logs.get(gen_id) if not g or g["tester"] != t.name or g["kind"] != "code" or not g.get("output"): raise HTTPException(status_code=404, detail="no code generation with that id for this key") buf = io.BytesIO() with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z: for path, content in json.loads(g["output"]).items(): z.writestr(path, content) return Response( buf.getvalue(), media_type="application/zip", headers={"Content-Disposition": f'attachment; filename="rosdiff-{gen_id}.zip"'}, ) @app.get("/", response_class=HTMLResponse, include_in_schema=False) def web_app() -> HTMLResponse: """The web app for testers (api/static/app).""" return HTMLResponse( (Path(__file__).with_name("static") / "app" / "index.html").read_text(), headers={"Cache-Control": "no-cache"}, ) @app.post("/feedback") def feedback_item(req: FeedbackItem, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: """👍/👎 (plus an optional note) on a generation or a job, stored with what RosDiff did for it.""" if req.target_type == "generation": g = s.logs.get(req.target_id) if not g or g["tester"] != t.name: raise HTTPException(status_code=404, detail="nothing with that id for this key") validation = g.get("validation") or {} context = { "kind": g["kind"], "valid": bool(g["valid"]) if g.get("valid") is not None else None, "model": g.get("model"), "routing": (g.get("request") or {}).get("routing"), "error": g.get("error"), "validation_error": validation.get("error") or validation.get("errors"), } s.logs.set_feedback(req.target_id, t.name, "works" if req.rating == "helpful" else "broken", req.note) else: j = s.logs.get_run(req.target_id) if not j or j["tester"] != t.name: raise HTTPException(status_code=404, detail="nothing with that id for this key") kind = j.get("kind") or "sim" context = { "kind": kind, "status": j["status"], "error": j.get("error"), "generation_id": j.get("generation_id"), "parent_id": j.get("parent_id"), "simulator": {"sim": "mujoco", "sweep": "mujoco_warp", "isaac": "isaac_sim"}.get(kind), "hardware": j.get("hardware"), "params": j.get("params"), "headline": ((j.get("output") or {}).get("summary") or {}).get("headline"), } s.logs.record_feedback(t.name, req.target_type, req.target_id, req.rating, req.note, context) return {"ok": True} @app.get("/me/usage") def my_usage(t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: """Your AI budget in rupees (and the team's, for admins). GPU jobs run on the separate compute credit.""" is_owner = t.name == "owner" out = { "tester": t.name, "is_owner": is_owner, "ai": budget.summary(s.logs, s.settings, t.name), "compute": { **compute_usage(s.settings, s.logs, t.name), "month_spend_usd_estimate": month_spend(s.settings, s.logs), "monthly_limit_usd": s.settings.monthly_credit_usd or None, "note": "Simulations and checks run on the cloud GPU credit, split equally between testers; " "it is separate from the rupee AI budget.", }, "admin": is_owner or t.name in s.settings.admin_testers, } if out["admin"]: names = [x.name for x in s.allowlist.testers()] out["team"] = budget.team(s.logs, s.settings, names) for row in out["team"]: row["gpu"] = compute_usage(s.settings, s.logs, row["tester"])["used"] return out @app.post("/feedback/{gen_id}") def feedback(gen_id: str, req: FeedbackRequest, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: """Tell us whether a generation actually worked when you ran it.""" if not s.logs.set_feedback(gen_id, t.name, req.rating, req.note): raise HTTPException(status_code=404, detail="no generation with that id for this key") return {"id": gen_id, "rating": req.rating} return app app = create_app()