Blame for server/environment.py

"""DBEREnvironment -- the core OpenEnv Environment implementation.

This class wires together:
  - fixtures  (corrupted DB per task, golden DB for grading)
  - safety    (pre/post-execution guards)
  - grader    (terminal score 0.0-1.0)
  - reward shaping (intermediate signals)

Episode lifecycle:
  reset(seed, task_id)  -> fresh DB, clean state, first observation
  step(action)          -> execute SQL or submit; returns updated observation
  state                 -> server-side state snapshot

See PROMPT.md sec2 for the full specification.
"""

from __future__ import annotations

import re
import sqlite3
from typing import Any, Dict, List, Optional

from openenv.core.env_server.interfaces import Environment, EnvironmentMetadata

from db_er.models import DBERAction, DBERObservation, DBERState
from server.fixtures import (
    CORE_TABLES,
    GOLDEN_ROW_COUNTS,
    TASK_TICKETS,
    create_corrupted_db,
    get_initial_row_counts,
)
from server.grader import compute_score, compute_violations
from server.safety import (
    SafetyViolation,
    budget_cost,
    check_post_execution,
    check_pre_execution,
    classify_sql,
)


#  Module-level shared episode state 
# openenv-core creates a NEW DBEREnvironment() instance per HTTP request
# (reset_handler and step_handler each call self._env_factory()).  All state
# must therefore live outside any instance - in this module-level object whose
# attributes are mutated in-place so every DBEREnvironment instance sees them.

class _EpisodeState:
    def __init__(self) -> None:
        self.conn: Optional[sqlite3.Connection] = None
        self.task_id: int = 1
        self.budget: int = 100
        self.violation_count: int = 0
        self.initial_row_counts: Dict[str, int] = {}
        self.safety_row_counts: Dict[str, int] = {}
        self.query_history: Dict[str, int] = {}
        self.step_count: int = 0
        self.episode_id: Optional[str] = None
        self.done: bool = False
        self.last_obs: Optional[DBERObservation] = None


_S = _EpisodeState()  # singleton - shared by all DBEREnvironment instances


class DBEREnvironment(Environment[DBERAction, DBERObservation, DBERState]):
    """Database Emergency Response OpenEnv environment.

    Supports 5 tasks of increasing difficulty:
      1 -- Phantom Duplicates  (Easy)
      2 -- Cascading Failure   (Medium)
      3 -- Payroll Black Hole  (Medium-Hard)
      4 -- Schema Drift        (Hard)
      5 -- Referential Maze    (Very Hard)
    """

    SUPPORTS_CONCURRENT_SESSIONS = False

    def get_metadata(self) -> EnvironmentMetadata:
        return EnvironmentMetadata(
            name="DB-ER: Database Emergency Response",
            description=(
                "An OpenEnv benchmark where an AI agent is paged into a production "
                "database incident and must diagnose and repair corruption using raw "
                "SQL queries under a strict action budget. Features 5 tasks spanning "
                "duplicate cleanup, FK restoration, payroll migration, schema repair, "
                "and multi-table referential integrity recovery. Graded deterministically "
                "via F1 score on exact row sets against a hidden golden database."
            ),
            version="1.0.0",
            author="Team maxout (Aadi Joshi, Kavya Bhand)",
            documentation_url="https://huggingface.co/spaces/aadi-joshi/openenvhack",
        )

    #  OpenEnv API 

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task_id: Optional[int] = None,
        **kwargs,
    ) -> DBERObservation:
        """Start a new episode.

        Args:
            seed:       Controls task selection when *task_id* is not given.
                        seed % 5 -> 0=task 1, 1=task 2, 2=task 3, 3=task 4, 4=task 5.
            episode_id: Opaque identifier stored in state.
            task_id:    Override task selection (1-5).
        """
        #  1. Pick task 
        if task_id is not None:
            if task_id not in (1, 2, 3, 4, 5):
                raise ValueError(f"task_id must be 1, 2, 3, 4, or 5 -- got {task_id!r}")
            _S.task_id = task_id
        elif seed is not None:
            _S.task_id = (seed % 5) + 1
        else:
            _S.task_id = 1

        #  2. Spin up a fresh corrupted DB 
        if _S.conn is not None:
            try:
                _S.conn.close()
            except Exception:
                pass

        _S.conn = create_corrupted_db(_S.task_id)
        _S.initial_row_counts = get_initial_row_counts(_S.conn, _S.task_id)
        # Safety floor = min(corrupted_initial, golden) per table.
        # - Task 1 (users 11->8 fix): min(11,8)=8  -> can delete 3 dupes, can't go below 8
        # - Task 2 (products 9->12 fix): min(9,12)=9 -> can INSERT more, can't delete any
        # - Task 3 (employees_old 18->15 fix): min(18,15)=15 -> can delete 3 dupes
        golden = GOLDEN_ROW_COUNTS[_S.task_id]
        _S.safety_row_counts = {
            table: min(_S.initial_row_counts.get(table, 0), golden.get(table, 0))
            for table in CORE_TABLES[_S.task_id]
        }

        #  3. Reset episode bookkeeping 
        _S.budget = 100
        _S.violation_count = compute_violations(_S.conn, _S.task_id)
        _S.query_history = {}
        _S.step_count = 0
        _S.episode_id = episode_id
        _S.done = False

        #  4. Build and return initial observation 
        obs = self._build_observation(
            last_action_type="NONE",
            last_query_result=[],
            error_logs="",
            reward=None,
            done=False,
        )
        _S.last_obs = obs
        return obs

    def step(
        self,
        action: DBERAction,
        timeout_s: Optional[float] = None,
        **kwargs,
    ) -> DBERObservation:
        """Execute one action and return the resulting observation."""
        if _S.conn is None or _S.done:
            raise RuntimeError(
                "Environment not initialised or episode already done. Call reset() first."
            )

        _S.step_count += 1

        if action.action_type == "execute_sql":
            return self._handle_execute_sql(action)
        elif action.action_type == "submit_resolution":
            return self._handle_submit_resolution(action)
        else:
            # Unknown action type - count as an error
            obs = self._build_observation(
                last_action_type="ERROR",
                last_query_result=[],
                error_logs=f"Unknown action_type: {action.action_type!r}",
                reward=-0.05,
                done=False,
            )
            _S.last_obs = obs
            return obs

    @property
    def state(self) -> DBERState:
        """Server-side state snapshot (for GET /state)."""
        return DBERState(
            task_id=_S.task_id,
            budget_remaining=_S.budget,
            violation_count=_S.violation_count,
            initial_row_counts=_S.initial_row_counts,
            query_history=dict(_S.query_history),
            episode_id=_S.episode_id,
            step_count=_S.step_count,
        )

    def close(self) -> None:
        # No-op: openenv-core calls close() after each HTTP request, but the
        # SQLite connection must persist across requests in _S (the singleton).
        pass

    #  Private helpers 

    def _handle_execute_sql(self, action: DBERAction) -> DBERObservation:
        query = (action.query or "").strip()
        if not query:
            return self._build_observation(
                last_action_type="ERROR",
                last_query_result=[],
                error_logs="execute_sql action requires a non-empty 'query' field.",
                reward=-0.05,
                done=False,
            )

        sql_type = classify_sql(query)
        core_tables = CORE_TABLES[_S.task_id]

        #  Pre-execution safety check 
        violation = check_pre_execution(query, core_tables)
        if violation:
            _S.done = True
            obs = self._build_observation(
                last_action_type="ERROR",
                last_query_result=[],
                error_logs=violation.message,
                reward=-1.0,
                done=True,
            )
            _S.last_obs = obs
            return obs

        #  Budget check 
        cost = budget_cost(sql_type)
        if _S.budget < cost:
            _S.done = True
            obs = self._build_observation(
                last_action_type="ERROR",
                last_query_result=[],
                error_logs=(
                    f"Budget exhausted. Remaining={_S.budget}, "
                    f"required={cost} for {sql_type}. Episode ended."
                ),
                reward=-0.5,
                done=True,
            )
            _S.last_obs = obs
            return obs

        #  Spam penalty 
        spam_penalty = 0.0
        if sql_type in ("SELECT", "META"):
            norm = _normalise_query(query)
            count = _S.query_history.get(norm, 0) + 1
            _S.query_history[norm] = count
            if count >= 3:
                spam_penalty = -0.05

        #  Execute the SQL 
        _S.budget -= cost

        prev_violation_count = _S.violation_count
        result_rows: List[Dict[str, Any]] = []
        error_msg = ""
        action_type_label = sql_type

        try:
            if sql_type == "MUTATION":
                # Use a savepoint so we can rollback on post-execution safety failure.
                # IMPORTANT: do NOT call conn.commit() before releasing the savepoint --
                # commit() would persist the changes before we can check row counts.
                _S.conn.execute("SAVEPOINT before_mutation")
                try:
                    cursor = _S.conn.execute(query)
                    affected = cursor.rowcount

                    #  Post-execution safety check (pre-commit) 
                    # Use golden row counts as floor so agents can delete corrupted rows
                    post_violation = check_post_execution(
                        _S.conn, _S.safety_row_counts
                    )
                    if post_violation:
                        # Roll back everything to before the mutation
                        _S.conn.execute("ROLLBACK TO before_mutation")
                        _S.conn.execute("RELEASE before_mutation")
                        _S.done = True
                        obs = self._build_observation(
                            last_action_type="ERROR",
                            last_query_result=[],
                            error_logs=post_violation.message,
                            reward=-1.0,
                            done=True,
                        )
                        _S.last_obs = obs
                        return obs

                    # Safe: commit by releasing the savepoint
                    _S.conn.execute("RELEASE before_mutation")

                    # Intermediate reward for 0-row mutations
                    if affected == 0:
                        spam_penalty += -0.1  # -0.1 for 0-row UPDATE/DELETE

                    result_rows = [{"rows_affected": affected}]
                    action_type_label = "MUTATION"

                except sqlite3.Error as exc:
                    try:
                        _S.conn.execute("ROLLBACK TO before_mutation")
                        _S.conn.execute("RELEASE before_mutation")
                    except Exception:
                        pass
                    error_msg = f"SQLite error: {exc}"
                    action_type_label = "ERROR"

            else:
                # SELECT / PRAGMA / META
                cursor = _S.conn.execute(query)
                rows = cursor.fetchall()
                if rows:
                    cols = [d[0] for d in cursor.description or []]
                    result_rows = [dict(zip(cols, row)) for row in rows]
                action_type_label = sql_type

        except sqlite3.Error as exc:
            error_msg = f"SQLite error: {exc}"
            action_type_label = "ERROR"
        except Exception as exc:
            error_msg = f"Unexpected error: {exc}"
            action_type_label = "ERROR"

        #  Recompute violations 
        _S.violation_count = compute_violations(_S.conn, _S.task_id)

        # Intermediate reward for reducing violations
        violation_reward = 0.0
        if _S.violation_count < prev_violation_count:
            reduction = prev_violation_count - _S.violation_count
            violation_reward = 0.2 * reduction

        total_reward = violation_reward + spam_penalty

        #  Budget exhausted mid-episode? 
        if _S.budget <= 0:
            _S.done = True

        obs = self._build_observation(
            last_action_type=action_type_label,
            last_query_result=result_rows if not error_msg else error_msg,
            error_logs=error_msg,
            reward=total_reward if total_reward != 0 else None,
            done=_S.done,
        )
        _S.last_obs = obs
        return obs

    def _handle_submit_resolution(self, action: DBERAction) -> DBERObservation:
        """Trigger the terminal grader and end the episode."""
        _S.done = True

        # Run the deterministic grader
        terminal_score = compute_score(_S.conn, _S.task_id)

        # False submission penalty - if the DB still has unresolved violations,
        # the agent submitted too early.  Apply -0.5 regardless of grader score.
        current_violations = compute_violations(_S.conn, _S.task_id)
        if current_violations > 0:
            final_reward = terminal_score - 0.5
        else:
            final_reward = terminal_score

        obs = self._build_observation(
            last_action_type="META",
            last_query_result=[
                {
                    "grader_score": terminal_score,
                    "message": (
                        f"Episode complete. Grader score: {terminal_score:.4f}. "
                        f"Notes: {action.notes or '(none)'}"
                    ),
                }
            ],
            error_logs="",
            reward=final_reward,
            done=True,
        )
        _S.last_obs = obs
        return obs

    #  Observation builder 

    def _build_observation(
        self,
        last_action_type: str,
        last_query_result,
        error_logs: str,
        reward: Optional[float],
        done: bool,
    ) -> DBERObservation:
        db_summary = _get_db_summary(_S.conn) if _S.conn else {}
        return DBERObservation(
            incident_ticket=TASK_TICKETS[_S.task_id],
            task_id=_S.task_id,
            last_action_type=last_action_type,
            last_query_result=last_query_result,
            error_logs=error_logs,
            db_summary=db_summary,
            violation_count=_S.violation_count,
            budget_remaining=_S.budget,
            step_count=_S.step_count,
            done=done,
            reward=reward,
        )


#  Utilities 

def _normalise_query(query: str) -> str:
    """Normalise a SQL query for spam detection (case-fold + collapse whitespace)."""
    q = query.upper()
    q = re.sub(r"\s+", " ", q).strip()
    return q


def _get_db_summary(conn: Optional[sqlite3.Connection]) -> Dict[str, Any]:
    """Return {"tables": [...], "row_counts": {...}} from the live connection."""
    if conn is None:
        return {"tables": [], "row_counts": {}}
    try:
        tables = [
            row[0]
            for row in conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
            ).fetchall()
        ]
        row_counts = {}
        for table in tables:
            row = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()
            row_counts[table] = row[0]
        return {"tables": tables, "row_counts": row_counts}
    except Exception:
        return {"tables": [], "row_counts": {}}