server/app.py

"""FastAPI application for ConfigDebugEnv.

Uses OpenEnv's create_fastapi_app() for standard framework compatibility
(WebSocket sessions, standard endpoints, grader discovery).

DEPLOYMENT: v2 - Session state management restored
"""
import json
import gradio as gr
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send
from typing import Callable

from openenv.core.env_server import create_fastapi_app
from server.models import ConfigDebugAction, ConfigDebugObservation, ConfigDebugState
from server.config_debug_environment import ConfigDebugEnvironment
from server.tasks.task_registry import get_task, TASK_ORDER
from server.step_middleware import StepPayloadWrapperMiddleware

# ---- Create the standard OpenEnv FastAPI app ----
# create_fastapi_app expects a callable (factory) that returns an Environment
app = create_fastapi_app(
    ConfigDebugEnvironment,       # factory / class — called per session
    ConfigDebugAction,            # action model (inherits Action)
    ConfigDebugObservation,       # observation model (inherits Observation)
)

# ---- FALLBACK: Global environment for session-less requests ----
# Validator may not send session IDs, so we maintain a global env instance
# that persists across requests for benchmark compatibility
GLOBAL_ENV = ConfigDebugEnvironment()
print("[APP_INIT] Global environment instance created for session-less requests")

# ---- Override OpenEnv's default /metadata route ----
# Remove the built-in metadata endpoint so we can replace it with task enumeration
remove_routes = ["/metadata", "/reset", "/step"]
for route_path in remove_routes:
    for i, route in enumerate(app.router.routes):
        if hasattr(route, "path") and route.path == route_path:
            app.router.routes.pop(i)
            print(f"[APP_INIT] Removed default {route_path} route for override")
            break

# ---- Middleware to fix /reset response schema ----
# OpenEnv returns {"observation": {...}, "reward": 0.0, "done": false}
# but validator expects {"observation": {...}, "info": {}}
class ResetSchemaFixMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        
        # Only fix /reset responses
        if request.url.path == "/reset" and request.method == "POST":
            if response.status_code == 200:
                try:
                    # Get response body
                    body = b""
                    async for chunk in response.body_iterator:
                        body += chunk
                    
                    data = json.loads(body)
                    
                    # Fix schema: Include base Observation fields at top level
                    # Reset response should have: observation, done, reward, metadata, info
                    if isinstance(data, dict) and "observation" in data:
                        fixed_data = {
                            "observation": data["observation"],
                            "done": data.get("done", False),
                            "reward": data.get("reward"),
                            "metadata": data.get("metadata", {}),
                            "info": {}
                        }
                        print("[MIDDLEWARE] Fixed /reset response schema - added base fields")
                        return JSONResponse(fixed_data, status_code=200)
                except Exception as e:
                    print(f"[MIDDLEWARE] Error fixing reset response: {e}")
        
        return response

app.add_middleware(ResetSchemaFixMiddleware)
app.add_middleware(StepPayloadWrapperMiddleware)

# ---- OVERRIDE: Custom /reset and /step for session-less validator ----
# When validator doesn't send session_id, use GLOBAL_ENV to persist state

@app.post("/reset")
async def reset_override(request: Request):
    """Custom reset handler - uses global env when no session_id."""
    session_id = request.headers.get("x-session-id")
    print(f"[RESET_OVERRIDE] session_id={session_id}")
    
    if not session_id or session_id == "":
        print(f"[RESET_OVERRIDE] No session_id provided, using GLOBAL_ENV")
        obs = GLOBAL_ENV.reset()
        return {
            "observation": obs.model_dump() if hasattr(obs, 'model_dump') else obs,
            "done": False,
            "reward": 0.0,
            "info": {}
        }
    
    # If session_id is provided, should let OpenEnv handle it
    # (but we removed those routes, so this is fallback only)
    print(f"[RESET_OVERRIDE] Session_id provided, using GLOBAL_ENV as fallback")
    obs = GLOBAL_ENV.reset()
    return {
        "observation": obs.model_dump() if hasattr(obs, 'model_dump') else obs,
        "done": False,
        "reward": 0.0,
        "info": {}
    }

@app.post("/step")
async def step_override(request: Request):
    """Custom step handler - uses global env when no session_id.
    
    Handles both direct and wrapped payload formats.
    """
    session_id = request.headers.get("x-session-id")
    print(f"[STEP_OVERRIDE] session_id={session_id}")
    
    # Parse request body to get fixed_config
    try:
        body = await request.json()
        print(f"[STEP_OVERRIDE] Raw body: {body}")
        
        # Extract fixed_config from wrapped or direct format
        fixed_config = None
        if "action" in body and isinstance(body["action"], dict):
            # Wrapped format: {"action": {"fixed_config": "..."}}
            fixed_config = body["action"].get("fixed_config", "{}")
        elif "fixed_config" in body:
            # Direct format: {"fixed_config": "..."}
            fixed_config = body["fixed_config"]
        else:
            print(f"[STEP_OVERRIDE] No fixed_config found in body")
            fixed_config = "{}"
        
        # Create action
        action = ConfigDebugAction(fixed_config=fixed_config)
        print(f"[STEP_OVERRIDE] Parsed action: fixed_config={action.fixed_config}")
        
        if not session_id or session_id == "":
            print(f"[STEP_OVERRIDE] No session_id provided, using GLOBAL_ENV")
            obs = GLOBAL_ENV.step(action)
            return {
                "observation": obs.model_dump() if hasattr(obs, 'model_dump') else obs,
                "done": obs.done,
                "reward": obs.reward,
                "info": {}
            }
        
        # If session_id is provided, use GLOBAL_ENV as fallback
        print(f"[STEP_OVERRIDE] Session_id provided, using GLOBAL_ENV as fallback")
        obs = GLOBAL_ENV.step(action)
        return {
            "observation": obs.model_dump() if hasattr(obs, 'model_dump') else obs,
            "done": obs.done,
            "reward": obs.reward,
            "info": {}
        }
        
    except Exception as e:
        print(f"[STEP_OVERRIDE] ERROR: {type(e).__name__}: {e}")
        raise

# ---- FALLBACK HANDLERS (deprecated - use /reset and /step above) ----
print("[APP_INIT] ConfigDebugEnvironment initialization started")
print(f"[APP_INIT] Loaded {len(TASK_ORDER)} tasks: {TASK_ORDER}")
for task_id in TASK_ORDER:
    try:
        task = get_task(task_id)
        print(f"[APP_INIT] Task '{task_id}' loaded: grader={task.grader.__name__}")
    except Exception as e:
        print(f"[APP_INIT] ERROR loading task '{task_id}': {str(e)}")



# ---- Custom endpoints ----

@app.get("/info")
def info():
    return {"name": "ConfigDebugEnv", "version": "1.0.0", "status": "running"}


@app.get("/health")
def health():
    return {"status": "healthy"}


@app.get("/debug")
def debug():
    """Debug endpoint to verify deployed code version and grader API."""
    from server.graders.grader_api import grade_task1, grade_task2, grade_task3, grade_task4, grade_task5, grade_task6, grade_task7
    from server.tasks.task1_json import BROKEN_CONFIG as b1
    from server.tasks.task2_yaml import BROKEN_CONFIG as b2
    
    return {
        "status": "ready",
        "version": "grader-fix-8e7ae76-bounds-0.01-0.99",
        "timestamp": "2026-04-11T",
        "grader_test": {
            "task1_broken_reward": float(grade_task1(b1)),
            "task2_broken_reward": float(grade_task2(b2)),
        },
        "message": "All graders return floats in (0, 1) strictly"
    }


@app.get("/diagnostics")
def diagnostics():
    """Diagnostics endpoint: validate all tasks and graders exactly as validator would.
    
    Tests both internal registry AND import paths from openenv.yaml manifest.
    """
    import importlib
    
    print("[VALIDATOR] GET /diagnostics called")
    tasks_info = []
    
    # Hardcoded mapping of task IDs to grader paths (from openenv.yaml)
    grader_paths = {
        "task1_json": "server.graders.grader_api:grade_task1",
        "task2_yaml": "server.graders.grader_api:grade_task2",
        "task3_dockerfile": "server.graders.grader_api:grade_task3",
        "task4_compose": "server.graders.grader_api:grade_task4",
        "task5_k8s": "server.graders.grader_api:grade_task5",
        "task6_github_actions": "server.graders.grader_api:grade_task6",
        "task7_nginx": "server.graders.grader_api:grade_task7",
    }
    
    for task_id in TASK_ORDER:
        try:
            # Test 1: Internal registry
            task = get_task(task_id)
            internal_callable = callable(task.grader)
            
            # Test 2: Import path validation (as validator would do it)
            grader_path = grader_paths.get(task_id)
            if ":" not in grader_path:
                raise ValueError(f"Invalid grader path format: {grader_path}")
            
            module_path, func_name = grader_path.split(":")
            module = importlib.import_module(module_path)
            grader_func = getattr(module, func_name)
            import_callable = callable(grader_func)
            
            tasks_info.append({
                "id": task_id,
                "grader_path": grader_path,
                "internal_registry": {
                    "grader_function": task.grader.__name__,
                    "callable": internal_callable
                },
                "manifest_import": {
                    "module": module_path,
                    "function": func_name,
                    "callable": import_callable,
                    "resolved": True
                },
                "status": "loaded" if (internal_callable and import_callable) else "partial"
            })
        except Exception as e:
            tasks_info.append({
                "id": task_id,
                "status": "error",
                "error": str(e),
                "error_type": type(e).__name__
            })
    
    total_loaded = sum(1 for t in tasks_info if t.get("status") == "loaded")
    
    return {
        "app_status": "running",
        "validation_method": "dual-check (internal registry + manifest import paths)",
        "total_tasks_expected": len(TASK_ORDER),
        "total_tasks_loaded": total_loaded,
        "all_valid": total_loaded == len(TASK_ORDER),
        "tasks": tasks_info
    }


@app.get("/metadata")
def metadata():
    """Override OpenEnv metadata endpoint with task enumeration for validator discovery.
    
    Task schema kept MINIMAL to avoid validator schema rejection:
    Only include fields the validator expects: id, has_grader
    """
    print("[VALIDATOR] GET /metadata called")
    return {
        "name": "ConfigDebugEnvironment",
        "description": "An environment for training AI agents to debug broken configuration files",
        "version": "1.0.0",
        "tasks": [
            {
                "id": tid,
                "has_grader": True,
            }
            for tid in TASK_ORDER
        ],
    }


@app.get("/tasks")
def tasks():
    return {
        "tasks": [
            {
                "id": tid,
                "name": get_task(tid).description,
                "difficulty": get_task(tid).difficulty,
                "file_type": get_task(tid).file_type,
                "num_bugs": get_task(tid).num_bugs,
                "has_grader": True,
            }
            for tid in TASK_ORDER
        ],
        "total_tasks": len(TASK_ORDER),
        "tasks_with_graders": len(TASK_ORDER),
    }


@app.post("/validate-graders")
def validate_graders():
    """RUNTIME AUDIT: Execute all graders with validator-style inputs and verify outputs.
    
    OpenEnv validator expects graders to return FLOAT ONLY (not tuples, not dicts).
    This tests the ACTUAL grader behavior that validator will encounter.
    
    For each task, we:
    1. Load broken config
    2. Call grader with it
    3. Verify output is float
    4. Verify bounds are strict (0, 1) - NOT 0.0 or 1.0
    5. Check for exceptions / NaN / inf
    
    Returns detailed audit of every grader's runtime behavior.
    """
    print("[VALIDATOR] POST /validate-graders called - RUNTIME AUDIT START")
    
    audit_results = []
    all_valid = True
    
    # Import all graders from wrapper layer (what validator uses)
    from server.graders.grader_api import (
        grade_task1, grade_task2, grade_task3, grade_task4,
        grade_task5, grade_task6, grade_task7,
    )
    
    graders = [
        (1, grade_task1),
        (2, grade_task2),
        (3, grade_task3),
        (4, grade_task4),
        (5, grade_task5),
        (6, grade_task6),
        (7, grade_task7),
    ]
    
    for task_num, grader_func in graders:
        task_id = TASK_ORDER[task_num - 1]
        
        try:
            task = get_task(task_id)
            broken_config = task.broken_config
            
            print(f"[VALIDATOR] Testing {task_id}: calling {grader_func.__name__}...")
            
            # Call the grader with broken config
            try:
                output = grader_func(broken_config)
                output_type = type(output).__name__
                
                # Validate output is FLOAT only
                is_float = isinstance(output, float)
                
                if is_float:
                    reward = output
                    in_bounds = 0 < reward < 1
                    is_exactly_zero = reward == 0.0
                    is_exactly_one = reward == 1.0
                    is_nan = reward != reward  # NaN check
                    is_inf = reward > 1e308 or reward < -1e308
                    
                    status = "valid" if in_bounds and not is_nan and not is_inf else "invalid"
                    if status != "valid":
                        all_valid = False
                else:
                    reward = None
                    status = "invalid"
                    all_valid = False
                    in_bounds = False
                    is_exactly_zero = False
                    is_exactly_one = False
                    is_nan = False
                    is_inf = False
                
                audit_results.append({
                    "task_id": task_id,
                    "grader": grader_func.__name__,
                    "call_status": "success",
                    "output": {
                        "type": output_type,
                        "value": reward if is_float else str(output),
                        "is_float": is_float,
                        "in_bounds": in_bounds,
                        "exactly_zero": is_exactly_zero,
                        "exactly_one": is_exactly_one,
                        "is_nan": is_nan,
                        "is_inf": is_inf,
                    },
                    "validation_status": status,
                })
                
                print(f"[VALIDATOR] {task_id}: {status} - output={reward if is_float else 'ERROR'}")
                
            except Exception as call_error:
                audit_results.append({
                    "task_id": task_id,
                    "grader": grader_func.__name__,
                    "call_status": "exception",
                    "exception": {
                        "type": type(call_error).__name__,
                        "message": str(call_error),
                    },
                    "validation_status": "invalid",
                })
                all_valid = False
                print(f"[VALIDATOR] {task_id}: EXCEPTION - {type(call_error).__name__}: {str(call_error)}")
                
        except Exception as task_load_error:
            audit_results.append({
                "task_id": task_id if 'task_id' in locals() else f"task{task_num}",
                "status": "task_load_error",
                "error": str(task_load_error),
            })
            all_valid = False
    
    print(f"[VALIDATOR] RUNTIME AUDIT COMPLETE - all_valid={all_valid}")
    
    return {
        "audit_type": "runtime_grader_execution",
        "timestamp": "2026-04-12",
        "validator_contract": "Each grader MUST return FLOAT ONLY in (0, 1) - NOT 0.0 or 1.0",
        "all_graders_valid": all_valid,
        "total_graders_tested": len(graders),
        "graders_passed": sum(1 for r in audit_results if r.get("validation_status") == "valid"),
        "audit_results": audit_results,
    }


# ---- Gradio Web UI ----

_ui_env = ConfigDebugEnvironment()


def format_state(env):
    """Format environment state with progress bar and RL signals."""
    state = env.state
    progress_bar = "█" * int(state.progress_ratio * 10) + "░" * (10 - int(state.progress_ratio * 10))
    
    return f"""
Task Progress: {len(state.tasks_completed)+1}/7
Progress: {progress_bar} ({int(state.progress_ratio*100)}%)
Total Reward: {state.total_reward:.2f}

Current Task: {state.current_task_id}
Difficulty: {state.current_difficulty}

Bugs Found: {state.bugs_found_so_far}
Error: {state.current_error_message or 'None'}

Completed: {', '.join(state.tasks_completed) if state.tasks_completed else 'None'}
Remaining: {', '.join(state.tasks_remaining[:3]) if state.tasks_remaining else 'None'}
"""


def ui_get_state():
    """Get current environment state (inspectable state)."""
    return format_state(_ui_env)


def ui_reset():
    _ui_env.reset()
    obs = _ui_env._build_observation()
    return (
        f"Task: {obs.task_id} | Difficulty: {obs.difficulty} | Bugs: {obs.num_bugs}",
        obs.task_description,
        obs.broken_config,
        obs.error_message,
        format_state(_ui_env),
        "Environment reset. Submit a fixed config to begin.",
    )


def ui_step(fixed_config):
    if _ui_env._done:
        return (
            "All tasks completed!",
            "",
            "",
            "Episode done. Click Reset to start again.",
            format_state(_ui_env),
            f"Final score: {_ui_env.total_reward:.1f} / {len(TASK_ORDER)}.0",
        )

    action = ConfigDebugAction(fixed_config=fixed_config)
    obs = _ui_env.step(action)

    history = f"Reward: {obs.reward:.2f} | Bugs found: {obs.bugs_found_so_far}/{obs.num_bugs}\nFeedback: {obs.error_message}"
    return (
        f"Task: {obs.task_id} | Difficulty: {obs.difficulty} | Bugs: {obs.num_bugs}",
        obs.task_description,
        obs.broken_config,
        obs.error_message,
        format_state(_ui_env),
        history,
    )


with gr.Blocks(title="ConfigDebugEnv") as demo:
    gr.Markdown("# ConfigDebugEnv")
    gr.Markdown("An RL environment for debugging broken config files across 7 real-world formats.")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Agent Interface")
            task_info = gr.Textbox(label="Current Task", interactive=False)
            task_desc = gr.Textbox(label="Task Description", interactive=False, lines=2)
            broken_config = gr.Textbox(label="Broken Config", interactive=False, lines=10)
            error_msg = gr.Textbox(label="Error Message", interactive=False, lines=2)

            gr.Markdown("### Take Action")
            fixed_config_input = gr.Textbox(label="Your Fixed Config", placeholder="Paste your fixed configuration here...", lines=10)
            with gr.Row():
                reset_btn = gr.Button("Reset Environment", variant="secondary")
                step_btn = gr.Button("Step", variant="primary")
                state_btn = gr.Button("Get State", variant="secondary")

        with gr.Column(scale=1):
            gr.Markdown("### State Observer")
            state_display = gr.Textbox(label="Current State (with RL Signals)", interactive=False, lines=14)
            history_display = gr.Textbox(label="Action History / Reward", interactive=False, lines=4)

    reset_btn.click(
        fn=ui_reset,
        outputs=[task_info, task_desc, broken_config, error_msg, state_display, history_display],
    )
    step_btn.click(
        fn=ui_step,
        inputs=[fixed_config_input],
        outputs=[task_info, task_desc, broken_config, error_msg, state_display, history_display],
    )
    state_btn.click(
        fn=ui_get_state,
        outputs=[state_display],
    )


app = gr.mount_gradio_app(app, demo, path="/")


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


if __name__ == "__main__":
    main()