server/app.py

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()