server.py

"""
server.py — FastAPI HTTP server for Incident Commander Arena.

Wraps IncidentEnvironment (from P1) and grade_episode (from grader.py)
behind a clean REST API so agents can interact over HTTP.

Endpoints:
    POST /reset          — start a new episode
    POST /step           — execute one agent action
    GET  /state          — inspect current episode state (debug)
    POST /grade          — score the finished episode
    GET  /scenarios      — list all loaded scenario IDs
    GET  /health         — liveness check

Run:
    uvicorn server:app --reload --port 8000

Install deps first:
    pip install fastapi uvicorn
"""

from __future__ import annotations

import sys
import os

# Make sure environment.py and grader.py are importable when running from any dir.
# Adjust this path if your folder layout is different.
sys.path.insert(0, os.path.dirname(__file__))

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import Optional

from environment import IncidentEnvironment
from grader import grade_episode

# ---------------------------------------------------------------------------
# App + environment setup
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Incident Commander Arena",
    description="SRE agent evaluation environment — REST API",
    version="1.0.0",
)

from fastapi.middleware.cors import CORSMiddleware

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # or ["http://localhost:5500"]
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Single shared environment instance — loaded once at startup.
# All episodes run through this one object; /reset wipes and resets state.
# Change the path if your scenarios folder is somewhere else.
SCENARIOS_DIR = os.environ.get("SCENARIOS_DIR", "./scenarios")

try:
    env = IncidentEnvironment(scenarios_dir=SCENARIOS_DIR)
except FileNotFoundError as e:
    print(f"\n[server] WARNING: {e}")
    print("[server] Server will start but /reset will fail until scenarios exist.\n")
    env = None  # Will be caught gracefully in /reset


# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------

class ResetRequest(BaseModel):
    task_id: Optional[str] = Field(
        default=None,
        description="Scenario ID to load. Omit to pick randomly.",
    )
    seed: Optional[int] = Field(
        default=None,
        description="Random seed for reproducibility.",
    )

class StepRequest(BaseModel):
    action: str = Field(
        description=(
            "One of: inspect_logs, query_metrics, check_dependencies, "
            "read_runbook, restart_service, rollback_deploy, "
            "page_engineer, update_status_page, mark_resolved"
        )
    )
    params: dict = Field(
        default={},
        description="Action-specific parameters. See environment.py for each action's required params.",
    )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _env_guard():
    """Raise 503 if environment failed to load (no scenarios folder yet)."""
    if env is None:
        raise HTTPException(
            status_code=503,
            detail=(
                f"Environment not initialised — no scenario JSON files found in '{SCENARIOS_DIR}'. "
                "Make sure P1 has shared at least one scenario file and restart the server."
            ),
        )

def _episode_guard():
    """Raise 400 if no episode is active yet."""
    if env.state is None:
        raise HTTPException(
            status_code=400,
            detail="No active episode. Call POST /reset first.",
        )

def _build_trajectory() -> list[dict]:
    """
    Convert env.timeline entries into the format grade_episode() expects.

    env.timeline stores lightweight summaries. The grader needs full
    action + params, so we reconstruct from the timeline entries.
    Note: env.timeline stores result_summary (a string), not the full result dict.
    The grader's root_cause dimension reads from mark_resolved params directly,
    so this is sufficient for all 5 scoring dimensions.
    """
    trajectory = []
    for entry in env.timeline:
        trajectory.append({
            "action": entry["action"],
            "params": entry["params"],
            "result": {"summary": entry.get("result_summary", "")},
        })
    return trajectory


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    """Liveness check. Returns server status and how many scenarios are loaded."""
    return {
        "status": "ok",
        "scenarios_loaded": len(env.scenarios) if env else 0,
        "scenarios_dir": SCENARIOS_DIR,
    }


@app.get("/scenarios")
def list_scenarios():
    """
    List all scenario IDs and their difficulty levels.
    Useful for agents or test harnesses that want to iterate over specific scenarios.
    """
    _env_guard()
    return {
        "scenarios": [
            {"id": sid, "difficulty": s["difficulty"]}
            for sid, s in env.scenarios.items()
        ]
    }


@app.post("/reset")
def reset(body: ResetRequest = ResetRequest()):
    """
    Start a fresh episode.

    - Picks the scenario specified by task_id, or picks one randomly.
    - Wipes all previous episode state.
    - Returns the initial observation the agent sees (no ground_truth).

    Example:
        POST /reset
        { "task_id": "payment-service-bad-deploy" }
    """
    _env_guard()

    try:
        observation = env.reset(task_id=body.task_id, seed=body.seed)
    except ValueError as e:
        # Unknown task_id
        raise HTTPException(status_code=404, detail=str(e))

    return {
        "status": "episode_started",
        "task_id": env.current_task_id,
        "difficulty": env.state["difficulty"],
        "observation": observation,
    }


@app.post("/step")
def step(body: StepRequest):
    """
    Execute one agent action and get back the updated observation.

    Returns:
        observation  — what the agent can see after this action
        reward       — shaped float reward (useful for RL agents)
        done         — True when the episode is over (mark_resolved or 20 steps)
        info         — debug info (forbidden_penalty, action, params)

    Example:
        POST /step
        { "action": "inspect_logs", "params": { "service": "payment-svc" } }
    """
    _env_guard()
    _episode_guard()

    if env.episode_done:
        raise HTTPException(
            status_code=400,
            detail=(
                "Episode is already finished. "
                "Call POST /reset to start a new one, or POST /grade to score this one."
            ),
        )

    try:
        observation, reward, done, info = env.step({
            "action": body.action,
            "params": body.params,
        })
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))

    response = {
        "observation": observation,
        "reward": reward,
        "done": done,
        "info": info,
        "step": env.step_count,
        "max_steps": env.MAX_STEPS,
    }

    # Surface forbidden action warnings prominently so agents can't miss them
    if info.get("forbidden_penalty", 0) > 0:
        response["warning"] = info.get("forbidden_warning", "Forbidden action triggered.")

    return response


@app.post("/grade")
def grade():
    """
    Score the current (or just-finished) episode across all 5 dimensions.

    Can be called:
      - After the episode ends (done=True from /step)
      - Mid-episode for a preview score (not the final grade — state may still change)

    Returns weighted total score plus per-dimension breakdown.

    Scoring weights:
        root_cause         40%  — did the agent identify the right cause?
        forbidden_actions  25%  — did the agent avoid penalty actions?
        blast_radius       20%  — did the agent stay on the affected service?
        efficiency         10%  — how many steps vs the optimal count?
        escalation          5%  — did the agent escalate correctly (or not at all)?
    """
    _env_guard()
    _episode_guard()

    trajectory = _build_trajectory()
    ground_truth = env.state["ground_truth"]

    result = grade_episode(
        trajectory=trajectory,
        ground_truth=ground_truth,
        env_state=env.state,
    )

    return {
        "task_id": env.current_task_id,
        "difficulty": env.state["difficulty"],
        "episode_done": env.episode_done,
        "steps_taken": env.step_count,
        "scores": result,
        "note": (
            "Final score." if env.episode_done
            else "Preview score — episode not yet finished. Grading mid-episode."
        ),
    }


@app.get("/state")
def get_state():
    """
    Debug endpoint. Returns a full human-readable render of the current episode.

    This calls env.render() which shows active alerts, timeline, and service
    statuses — but still strips ground_truth. Safe to share with agents for
    debugging but not part of the standard agent protocol.
    """
    _env_guard()
    _episode_guard()

    return {
        "render": env.render(),
        "step_count": env.step_count,
        "episode_done": env.episode_done,
        "timeline": env.timeline,
    }


# ---------------------------------------------------------------------------
# Entry point (for running directly: python server.py)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)

def main(host: str = "0.0.0.0", port: int = 8000):
    import uvicorn
    uvicorn.run(app, host=host, port=port)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()
    main(host=args.host, port=args.port)
    
    def main(host: str = "0.0.0.0", port: int = 8000):
    import uvicorn
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()
    main(host=args.host, port=args.port)