from __future__ import annotations import argparse from pathlib import Path from typing import Annotated, Any from fastapi import FastAPI, HTTPException, Query from fastapi.responses import HTMLResponse from pydantic import BaseModel, Field from server.environment import SecretsAuditEnvironment app = FastAPI(title="SecretsAuditEnv", version="1.0") environment = SecretsAuditEnvironment(repo_root=Path(__file__).resolve().parents[1]) class ResetRequest(BaseModel): task_id: int | str | None = None id: int | str | None = None difficulty: str | None = None class StepRequest(BaseModel): action: str = Field(..., min_length=1) def _parse_task_id(raw: Any) -> int: if raw is None: return 1 if isinstance(raw, int): task_id = raw else: text = str(raw).strip() if text.startswith("task_"): text = text.split("_", 1)[1] try: task_id = int(text) except ValueError as exc: raise HTTPException(status_code=422, detail="Invalid task_id value.") from exc if not 1 <= task_id <= 13: raise HTTPException(status_code=422, detail="task_id must be between 1 and 13.") return task_id _DIFFICULTY_TO_TASK_ID = {"easy": 1, "medium": 6, "hard": 11} @app.get("/") def root() -> dict: return { "name": "SecretsAuditEnv", "version": "1.0", "status": "healthy", "docs": "/docs", "web_ui": "/web", "tasks_endpoint": "/tasks", } @app.get("/health") def health() -> dict: """Health check endpoint — must return {"status": "healthy"} for OpenEnv.""" return {"status": "healthy"} @app.post("/mcp") def mcp(body: dict | None = None) -> dict: """MCP endpoint — required by OpenEnv runtime validator. Returns a valid JSON-RPC 2.0 response. Supports basic method routing for the OpenEnv MCP protocol. """ body = body or {} method = body.get("method", "") req_id = body.get("id", 1) if method == "initialize": return { "jsonrpc": "2.0", "id": req_id, "result": { "protocolVersion": "2024-11-05", "serverInfo": {"name": "SecretsAuditEnv", "version": "1.0"}, "capabilities": {}, }, } # Default: return a valid JSON-RPC 2.0 error for unknown methods return { "jsonrpc": "2.0", "id": req_id, "error": { "code": -32601, "message": f"Method not found: {method}", }, } @app.get("/metadata") def metadata() -> dict: """Metadata endpoint — required by OpenEnv runtime validator.""" return { "name": "SecretsAuditEnv", "description": "A 13-task RL benchmark for AI agents that find and fix secret leaks in git-backed codebases. Agents interact via bash commands and structured inspection actions, graded on security, health, and efficiency.", "version": "1.0", "author": "Team R2X", } @app.get("/schema") def schema() -> dict: """Schema endpoint — required by OpenEnv runtime validator.""" return { "action": { "type": "object", "properties": { "action": { "type": "string", "description": "Bash command or structured action (inspect_file, inspect_git_history, inspect_encoded)", } }, "required": ["action"], }, "observation": { "type": "object", "properties": { "visible_secrets": {"type": "array"}, "hidden_count_hint": {"type": "integer"}, "ranked_actions": {"type": "array"}, "top_blocker": {"type": "string"}, "step_budget": {"type": "integer"}, "steps_taken": {"type": "integer"}, "steps_remaining": {"type": "integer"}, "efficiency_bonus": {"type": "number"}, "conflict_map": {"type": "object"}, "security_score": {"type": "number"}, "health_score": {"type": "number"}, "reward": {"type": "number"}, }, }, "state": { "type": "object", "properties": { "task_id": {"type": "integer"}, "step_count": {"type": "integer"}, "reward": {"type": "number"}, "done": {"type": "boolean"}, }, }, } @app.get("/state") def get_state() -> dict: return environment.state() @app.get("/tasks", tags=["OpenEnv"]) def get_tasks() -> dict: return { "tasks": environment.list_tasks(), "action_schema": { "type": "object", "properties": { "action": { "type": "string", "description": "Bash command or structured inspection action", } }, "required": ["action"], }, } @app.post("/reset") def reset( request: ResetRequest | None = None, task_id: Annotated[str | int | None, Query()] = None, ) -> dict: try: raw_task_id: Any = task_id if raw_task_id is None and request is not None: if request.difficulty and request.task_id is None and request.id is None: raw_task_id = _DIFFICULTY_TO_TASK_ID.get(request.difficulty, 1) else: raw_task_id = request.task_id if request.task_id is not None else request.id resolved_task_id = _parse_task_id(raw_task_id) return environment.reset(resolved_task_id) except FileNotFoundError as exc: raise HTTPException(status_code=404, detail=str(exc)) from exc @app.post("/step") def step(request: StepRequest) -> dict: try: state = environment.step(request.action) session = (state.get("session") or {}) return { **state, "reward": session.get("reward", 0.0), "cumulative_reward": session.get("reward", 0.0), "done": float(session.get("reward", 0.0)) >= 0.99, "step": session.get("step_count", 0), "observation": session.get("observation", ""), } except RuntimeError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc @app.post("/grader", tags=["OpenEnv"]) def grader(body: dict) -> dict: """Grade a completed episode — required by OpenEnv validator. Accepts {task_id, step_rewards, step_infos} and returns {"score": float}. Score is always strictly in (0, 1). """ task_id = body.get("task_id") if not task_id: raise HTTPException(422, "'task_id' required") step_rewards = body.get("step_rewards", []) step_infos = body.get("step_infos", []) if step_rewards: avg_reward = sum(step_rewards) / len(step_rewards) else: # No rewards provided — use current session state try: task_num = _parse_task_id(task_id) state = environment.reset(task_num) session = (state.get("session") or {}) avg_reward = float(session.get("reward", 0.01)) except Exception: avg_reward = 0.01 # Clamp to strict (0, 1) — validator rejects 0.0 and 1.0 score = round(max(0.01, min(0.99, avg_reward)), 4) return { "task_id": task_id, "score": score, "breakdown": { "avg_reward": round(avg_reward, 4), "steps": len(step_rewards), }, } @app.post("/baseline", tags=["OpenEnv"]) def run_baseline(body: dict = {}) -> dict: """Run a baseline evaluation across tasks — required by OpenEnv.""" task_ids = ["task_1", "task_2", "task_3"] results = [] for tid in task_ids: try: task_num = _parse_task_id(tid) state = environment.reset(task_num) session = (state.get("session") or {}) reward = float(session.get("reward", 0.01)) score = round(max(0.01, min(0.99, reward)), 4) except Exception: score = 0.01 results.append({"task_id": tid, "score": score}) overall = round(sum(r["score"] for r in results) / max(len(results), 1), 4) return {"tasks": results, "summary": {"overall_score": overall}} @app.get("/web", response_class=HTMLResponse) async def web_ui() -> str: html_path = Path(__file__).parent / "web_ui.html" return html_path.read_text(encoding="utf-8") def main() -> None: parser = argparse.ArgumentParser(description="Run the SecretsAuditEnv server.") parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=7860) args = parser.parse_args() import uvicorn uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": main()