Upload 6 files

server/app.py

@@ -0,0 +1,17 @@
+"""
+Alias for server/main.py — required by openenv CLI structure.
+The actual application is defined in server/main.py.
+"""
+
+from server.main import app, run_server
+
+__all__ = ["app"]
+
+
+def main():
+    """Main function for openenv CLI compatibility."""
+    run_server()
+
+
+if __name__ == "__main__":
+    main()

server/environment.py

@@ -0,0 +1,509 @@
+"""SQL Quality Environment - Core environment manager."""
+
+from __future__ import annotations
+
+import random as _random
+import sqlite3
+import threading
+import uuid
+from typing import Any, Dict
+
+from fastapi import HTTPException
+
+from server.models import (
+    ActionType,
+    SQLAction,
+    SQLObservation,
+    SQLReward,
+    SQLState,
+    StepResponse,
+    TrajectoryStep,
+)
+from server.tasks import (
+    create_dedup_bug_task,
+    create_fk_integrity_task,
+    create_null_audit_task,
+    create_schema_drift_task,
+    create_type_repair_task,
+)
+
+TASK_CREATORS = {
+    "null_audit": create_null_audit_task,
+    "type_repair": create_type_repair_task,
+    "dedup_bug": create_dedup_bug_task,
+    "fk_integrity": create_fk_integrity_task,
+    "schema_drift": create_schema_drift_task,
+}
+
+
+class SQLQualityEnvironment:
+    """Main environment class for SQL Quality tasks."""
+
+    def __init__(self) -> None:
+        """Initialize environment."""
+        self._current_task_id: str | None = None
+        self._episode_id: str | None = None
+        self._attempt_number: int = 1
+        self._step_count: int = (
+            0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
+        )
+        self._max_steps_per_episode: int = 20  # prevent infinite loops
+        self._cumulative_reward: float = 0.0
+        self._done: bool = False
+        self._ground_truth: Dict[str, Any] = {}
+        self._cached_observation: SQLObservation | None = None
+        self._lock = threading.Lock()
+        # NEW: Multi-step trajectory support
+        self._trajectory: list = []
+        self._original_snapshot: dict = {}
+        self._original_conn: sqlite3.Connection | None = None
+        self._active: bool = False
+        self._last_feedback: str | None = None
+
+    def reset(self, task_id: str, seed: int | None = None) -> SQLObservation:
+        """Reset environment for a new episode.
+
+        Args:
+            task_id: One of the registered task IDs
+            seed: Optional deterministic episode seed
+
+        Returns:
+            SQLObservation for the new episode
+
+        Raises:
+            ValueError: If task_id is not valid
+        """
+        with self._lock:
+            # ── NEW: close any live connection from the previous episode ──────
+            if hasattr(self, "_original_conn") and self._original_conn is not None:
+                try:
+                    self._original_conn.close()
+                except Exception:
+                    pass  # Already closed or never opened — safe to ignore
+                self._original_conn = None
+
+            # Validate task_id
+            if task_id not in TASK_CREATORS:
+                raise ValueError(f"Invalid task_id: {task_id}")
+
+            # Generate new episode
+            self._episode_id = str(uuid.uuid4())
+            self._attempt_number = 1
+            self._step_count = (
+                0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
+            )
+            self._cumulative_reward = 0.0
+            self._done = False
+            self._current_task_id = task_id
+            self._trajectory = []
+            self._original_snapshot = {}
+            self._ground_truth = {}
+            self._cached_observation = None
+            self._last_feedback = None
+
+            # Create task-specific database to extract observation data
+            conn = None
+            try:
+                # Generate random episode seed for diversity
+                episode_seed = seed if seed is not None else _random.randint(0, 2**31 - 1)
+                creator = TASK_CREATORS[task_id]
+                conn, observation, self._ground_truth = creator(episode_seed)
+
+                # Store original connection for graders (read-only use)
+                self._original_conn = conn
+
+                # Snapshot the target table for data preservation grading
+                target_table = self._ground_truth.get("table", "")
+                if target_table:
+                    from server.graders import _snapshot_table
+
+                    self._original_snapshot = _snapshot_table(conn, target_table)
+
+                # Mark episode as active
+                self._active = True
+
+                # Cache observation for reuse in step()
+                self._cached_observation = observation.model_copy(
+                    update={"episode_id": self._episode_id}
+                )
+
+                # Keep connection open for multi-step operations
+                # Connection will be closed when episode ends or environment is reset
+
+            except Exception as e:
+                if conn is not None:
+                    conn.close()
+                raise e
+
+            return self._cached_observation
+
+    def step(self, action: SQLAction) -> StepResponse:
+        """
+        Execute one agent action and return the resulting observation and reward.
+
+        Routes by action_type:
+            QUERY   — execute SELECT on episode conn, return result in observation
+            PROFILE — compute column statistics, return in observation
+            FIX     — execute on fresh_test_db(), score partially, continue episode
+            SUBMIT  — call compute_episode_score(), end episode
+        """
+        with self._lock:
+            if not self._active:
+                raise HTTPException(
+                    status_code=422, detail="No active episode. Call reset() first."
+                )
+            if self._done:
+                raise HTTPException(
+                    status_code=422,
+                    detail="Episode is already done. Call reset() to start a new episode.",
+                )
+
+            # Check max steps guard
+            if self._step_count >= self._max_steps_per_episode:
+                raise HTTPException(
+                    status_code=422,
+                    detail="Max investigation steps reached. Call SUBMIT or reset().",
+                )
+
+            action = self._normalize_action(action)
+
+            action_type = action.action_type
+
+            # --- QUERY action ---
+            if action_type == ActionType.QUERY:
+                return self._handle_query(action)
+
+            # --- PROFILE action ---
+            elif action_type == ActionType.PROFILE:
+                return self._handle_profile(action)
+
+            # --- FIX action ---
+            elif action_type == ActionType.FIX:
+                return self._handle_fix(action)
+
+            # --- SUBMIT action (or legacy single-step) ---
+            else:
+                return self._handle_submit(action)
+
+    def _normalize_action(self, action: SQLAction) -> SQLAction:
+        """Infer v2 action_type for backward-compatible payloads without action_type."""
+        if "action_type" in action.model_fields_set:
+            return action
+
+        inferred = ActionType.SUBMIT
+        if action.column_name:
+            inferred = ActionType.PROFILE
+        elif action.fix_query:
+            inferred = ActionType.FIX
+        elif action.sql:
+            inferred = (
+                ActionType.QUERY
+                if action.sql.strip().upper().startswith("SELECT")
+                else ActionType.FIX
+            )
+        elif action.answer or action.root_cause_table or action.offending_key:
+            inferred = ActionType.SUBMIT
+
+        return action.model_copy(update={"action_type": inferred})
+
+    def _handle_query(self, action: SQLAction) -> StepResponse:
+        """
+        Handle a read-only QUERY action.
+
+        Executes the SQL on the episode connection (read-only SELECT enforced).
+        Returns a result preview in the observation. Does NOT end the episode.
+        Awards small step_reward (0.05) for queries that target the right area.
+        """
+        sql = (action.sql or "").strip()
+
+        # Enforce read-only: only SELECT allowed
+        if not sql.upper().startswith("SELECT"):
+            step_reward = 0.0
+            result_preview = "Error: QUERY action only accepts SELECT statements."
+        else:
+            try:
+                cursor = self._original_conn.execute(sql)
+                rows = cursor.fetchmany(5)
+                cols = [d[0] for d in cursor.description]
+                result_preview = str([dict(zip(cols, row)) for row in rows])[:300]
+                # Small reward for targeted queries
+                from server.graders import _query_reveals_issue
+
+                step_reward = (
+                    0.05 if _query_reveals_issue(sql, self._ground_truth) else 0.01
+                )
+            except sqlite3.Error as e:
+                result_preview = f"SQL Error: {e}"
+                step_reward = 0.0
+
+        # Add to trajectory
+        self._trajectory.append(
+            TrajectoryStep(
+                step_number=len(self._trajectory) + 1,
+                action_type=ActionType.QUERY,
+                sql_submitted=sql,
+                result_preview=result_preview[:200],
+                step_reward=step_reward,
+            )
+        )
+
+        self._cumulative_reward = min(
+            1.0, round(self._cumulative_reward + step_reward, 4)
+        )
+        self._last_feedback = result_preview[:200]
+        self._step_count += 1
+
+        # Build updated observation with trajectory history
+        obs = self._build_observation(result_preview=result_preview)
+
+        return StepResponse(
+            observation=obs,
+            reward=SQLReward(
+                score=self._cumulative_reward,
+                step_reward=step_reward,
+                feedback=result_preview[:200],
+            ),
+            done=False,
+            info={"action_type": "query", "step": len(self._trajectory)},
+        )
+
+    def _handle_profile(self, action: SQLAction) -> StepResponse:
+        """
+        Handle a PROFILE action — statistical summary of a named column.
+
+        Computes: null rate, distinct count, data type distribution, min/max.
+        Does NOT end the episode. Awards step_reward based on column relevance.
+        """
+        col = (action.column_name or "").strip()
+        result_preview = ""
+
+        if not col:
+            result_preview = "Error: column_name required for PROFILE action."
+        else:
+            # Try to find which table has this column
+            try:
+                tables = self._original_conn.execute(
+                    "SELECT name FROM sqlite_master WHERE type='table'"
+                ).fetchall()
+                profile_results = []
+                for (tbl,) in tables:
+                    try:
+                        cols = [
+                            row[1]
+                            for row in self._original_conn.execute(
+                                f"PRAGMA table_info({tbl})"
+                            ).fetchall()
+                        ]
+                        if col in cols:
+                            total = self._original_conn.execute(
+                                f"SELECT COUNT(*) FROM {tbl}"
+                            ).fetchone()[0]
+                            nulls = self._original_conn.execute(
+                                f"SELECT COUNT(*) FROM {tbl} WHERE {col} IS NULL"
+                            ).fetchone()[0]
+                            distinct = self._original_conn.execute(
+                                f"SELECT COUNT(DISTINCT {col}) FROM {tbl}"
+                            ).fetchone()[0]
+                            profile_results.append(
+                                f"{tbl}.{col}: total={total}, nulls={nulls} ({100*nulls//total if total else 0}%), distinct={distinct}"
+                            )
+                    except sqlite3.Error:
+                        pass
+                result_preview = (
+                    " | ".join(profile_results)
+                    if profile_results
+                    else f"Column '{col}' not found."
+                )
+            except sqlite3.Error as e:
+                result_preview = f"Profile error: {e}"
+
+        # Determine step_reward based on relevance to ground truth
+        ground_table = self._ground_truth.get("table", "").lower()
+        ground_col = self._ground_truth.get(
+            "column", self._ground_truth.get("null_column", "")
+        ).lower()
+        col_lower = col.lower()
+
+        if col_lower == ground_col:
+            step_reward = 0.08  # On-target: profiling exact relevant column
+        elif ground_table and ground_table in result_preview.lower():
+            step_reward = 0.05  # Near-target: profiled something from right table
+        elif col_lower in result_preview.lower() and len(col_lower) > 2:
+            step_reward = 0.01  # Off-target: profiled something real but irrelevant
+        else:
+            step_reward = 0.01  # Minimal: invalid column or junk probe
+
+        self._trajectory.append(
+            TrajectoryStep(
+                step_number=len(self._trajectory) + 1,
+                action_type=ActionType.PROFILE,
+                sql_submitted=f"PROFILE({col})",
+                result_preview=result_preview[:200],
+                step_reward=step_reward,
+            )
+        )
+
+        self._cumulative_reward = min(
+            1.0, round(self._cumulative_reward + step_reward, 4)
+        )
+        self._last_feedback = result_preview[:200]
+        self._step_count += 1
+        obs = self._build_observation(result_preview=result_preview)
+
+        return StepResponse(
+            observation=obs,
+            reward=SQLReward(
+                score=self._cumulative_reward,
+                step_reward=step_reward,
+                feedback=result_preview[:200],
+            ),
+            done=False,
+            info={"action_type": "profile", "column_profiled": col},
+        )
+
+    def _handle_fix(self, action: SQLAction) -> StepResponse:
+        """
+        Handle a FIX action — execute write query on isolated copy for preview scoring.
+
+        Executes on a _fresh_test_db() copy. Does NOT mutate episode state.
+        Returns intermediate feedback on whether fix appears to work.
+        Does NOT end episode — agent can follow up with verification query or SUBMIT.
+        Awards intermediate reward based on quick fix assessment.
+        """
+        fix_sql = action.effective_sql or ""
+        step_reward = 0.0
+        feedback = ""
+
+        if fix_sql:
+            from server.graders import grade_fix_correctness
+
+            quick_score, severity = grade_fix_correctness(
+                fix_sql,
+                self._original_conn,
+                self._original_snapshot,
+                self._ground_truth,
+            )
+            step_reward = quick_score * 0.3  # Partial credit for fix attempt
+            feedback = f"Fix preview: {severity} ({quick_score:.2f}). Submit to finalize score."
+        else:
+            feedback = "No SQL provided in fix action."
+
+        self._trajectory.append(
+            TrajectoryStep(
+                step_number=len(self._trajectory) + 1,
+                action_type=ActionType.FIX,
+                sql_submitted=fix_sql,
+                result_preview=feedback[:200],
+                step_reward=step_reward,
+            )
+        )
+
+        self._cumulative_reward = min(
+            1.0, round(self._cumulative_reward + step_reward, 4)
+        )
+        self._last_feedback = feedback
+        self._step_count += 1
+        obs = self._build_observation(result_preview=feedback)
+
+        return StepResponse(
+            observation=obs,
+            reward=SQLReward(
+                score=self._cumulative_reward,
+                step_reward=step_reward,
+                feedback=feedback,
+            ),
+            done=False,
+            info={
+                "action_type": "fix",
+                "severity_preview": severity if fix_sql else "N/A",
+            },
+        )
+
+    def _handle_submit(self, action: SQLAction) -> StepResponse:
+        """
+        Handle a SUBMIT action — compute final composite score and end episode.
+
+        Calls compute_episode_score() with the full trajectory.
+        Sets done=True. This is the only action that ends the episode.
+        """
+        from server.graders import compute_episode_score
+
+        reward = compute_episode_score(
+            trajectory=self._trajectory,
+            ground_truth=self._ground_truth,
+            original_conn=self._original_conn,
+            original_snapshot=self._original_snapshot,
+            submit_action=action,
+        )
+
+        self._done = True
+        self._active = False
+        self._cumulative_reward = min(
+            1.0, round(self._cumulative_reward + reward.score, 4)
+        )
+        self._last_feedback = reward.feedback
+        obs = self._build_observation()
+
+        # Close connection when episode ends
+        if self._original_conn:
+            self._original_conn.close()
+            self._original_conn = None
+
+        return StepResponse(
+            observation=obs,
+            reward=reward,
+            done=True,
+            info={
+                "action_type": "submit",
+                "final_score": reward.score,
+                "dimensions": reward.dimensions,
+                "trajectory_length": len(self._trajectory),
+            },
+        )
+
+    def _build_observation(self, result_preview: str = "") -> SQLObservation:
+        """Build current observation including trajectory history."""
+        history = [
+            {
+                "step": s.step_number,
+                "type": s.action_type.value,
+                "sql": (s.sql_submitted or "")[:80],
+                "result": s.result_preview[:80],
+                "reward": s.step_reward,
+            }
+            for s in self._trajectory[-5:]  # Last 5 steps only (context window)
+        ]
+        return SQLObservation(
+            task_id=self._current_task_id,
+            task_description=self._cached_observation.task_description,
+            schema_info=self._cached_observation.schema_info,
+            sample_rows=self._cached_observation.sample_rows,
+            attempt_number=self._attempt_number,
+            episode_id=self._episode_id,
+            trajectory_history=history,
+            investigation_hint=result_preview[:200] if result_preview else None,
+            previous_feedback=(
+                self._last_feedback if hasattr(self, "_last_feedback") else None
+            ),
+            max_attempts=self._cached_observation.max_attempts,
+        )
+
+    def state(self) -> SQLState:
+        """Get current environment state.
+
+        Returns:
+            SQLState with current episode information
+
+        Raises:
+            ValueError: If no episode has been started
+        """
+        if self._episode_id is None:
+            raise ValueError("No episode started. Call reset() first.")
+
+        return SQLState(
+            task_id=self._current_task_id,
+            episode_id=self._episode_id,
+            attempt_number=self._attempt_number,
+            cumulative_reward=self._cumulative_reward,
+            done=self._done,
+        )

server/graders.py

@@ -0,0 +1,843 @@
+"""Grading functions for SQL Quality Environment."""
+
+from __future__ import annotations
+
+from difflib import SequenceMatcher
+import sqlite3
+from typing import Any, Dict
+
+from server.models import ActionType, SQLAction, SQLReward, TrajectoryStep
+
+
+def _fresh_test_db(original_conn: sqlite3.Connection) -> sqlite3.Connection:
+    """Create an isolated in-memory copy of a sqlite database."""
+    copy = sqlite3.connect(":memory:")
+    for line in original_conn.iterdump():
+        copy.execute(line)
+    copy.commit()
+    return copy
+
+
+# Backward-compatible alias
+_make_fresh_copy = _fresh_test_db
+
+
+def _query_reveals_issue(sql: str, ground_truth: dict) -> bool:
+    """Return True when query references both ground-truth table and column."""
+    sql_lower = (sql or "").lower()
+    table = ground_truth.get("table", ground_truth.get("target_table", "")).lower()
+    col = ground_truth.get(
+        "column", ground_truth.get("null_column", ground_truth.get("target_column", ""))
+    ).lower()
+    return bool(table and col and table in sql_lower and col in sql_lower)
+
+
+def _measure_data_preservation(
+    test_conn: sqlite3.Connection,
+    original_snapshot: dict,
+    table_name: str,
+    key_col: str = "id",
+) -> float:
+    """Measure row/value preservation ratio after a fix."""
+    if not original_snapshot:
+        return 1.0
+
+    try:
+        cursor = test_conn.execute(f"SELECT * FROM {table_name}")
+        cols = [d[0] for d in cursor.description]
+        if key_col not in cols:
+            return 0.0
+        key_idx = cols.index(key_col)
+        post_fix = {row[key_idx]: dict(zip(cols, row)) for row in cursor.fetchall()}
+    except sqlite3.Error:
+        return 0.0
+
+    if not post_fix:
+        return 0.0
+
+    original_ids = set(original_snapshot.keys())
+    post_ids = set(post_fix.keys())
+    row_score = len(original_ids & post_ids) / max(1, len(original_ids))
+
+    matched_rows = original_ids & post_ids
+    if not matched_rows:
+        return round(row_score * 0.5, 4)
+
+    value_matches = 0
+    value_total = 0
+    for rid in matched_rows:
+        orig_row = original_snapshot[rid]
+        post_row = post_fix[rid]
+        for col in orig_row:
+            value_total += 1
+            try:
+                if abs(float(str(orig_row[col])) - float(str(post_row.get(col, 0)))) < 0.001:
+                    value_matches += 1
+            except (ValueError, TypeError):
+                if str(orig_row[col]) == str(post_row.get(col, "")):
+                    value_matches += 1
+    value_score = value_matches / max(1, value_total)
+    return round(row_score * 0.5 + value_score * 0.5, 4)
+
+
+def _snapshot_table(
+    conn: sqlite3.Connection,
+    table_name: str,
+    key_col: str = "id",
+    limit: int = 200,
+) -> dict:
+    """Capture table contents before a fix."""
+    try:
+        cursor = conn.execute(f"SELECT * FROM {table_name} LIMIT {limit}")
+        cols = [d[0] for d in cursor.description]
+        if key_col not in cols:
+            return {}
+        key_idx = cols.index(key_col)
+        return {row[key_idx]: dict(zip(cols, row)) for row in cursor.fetchall()}
+    except sqlite3.Error:
+        return {}
+
+
+def _snapshot_for_preservation(original_snapshot: dict, ground_truth: dict) -> dict:
+    """Filter snapshots to preservation-critical columns when configured."""
+    preserve_columns = ground_truth.get("preserve_columns")
+    if not preserve_columns:
+        return original_snapshot
+    filtered: dict = {}
+    for row_id, row in original_snapshot.items():
+        filtered[row_id] = {col: row[col] for col in preserve_columns if col in row}
+    return filtered
+
+
+def _null_audit_score(answer: str, ground_truth: dict) -> tuple[float, str, dict]:
+    """Schema-aware null_audit scoring with lexical near-miss partial credit."""
+    clean = (answer or "").strip()
+    if not clean:
+        return 0.0, "CATASTROPHIC", {
+            "answer_correct": 0.0,
+            "answer_in_schema": 0.0,
+            "lexical_near_miss": 0.0,
+        }
+
+    clean_lower = clean.lower()
+    correct = ground_truth.get("null_column", "").strip().lower()
+    all_cols = [c.lower() for c in ground_truth.get("all_columns", [])]
+
+    if clean_lower == correct:
+        return 1.0, "PERFECT", {
+            "answer_correct": 1.0,
+            "answer_in_schema": 0.0,
+            "lexical_near_miss": 0.0,
+        }
+
+    if clean_lower in all_cols:
+        return 0.2, "ACCEPTABLE", {
+            "answer_correct": 0.0,
+            "answer_in_schema": 0.2,
+            "lexical_near_miss": 0.0,
+        }
+
+    # Bounded lexical partial for out-of-schema near misses only.
+    best_similarity = 0.0
+    for col in all_cols:
+        best_similarity = max(best_similarity, SequenceMatcher(None, clean_lower, col).ratio())
+    if best_similarity >= 0.85:
+        lexical = 0.15
+    elif best_similarity >= 0.65:
+        lexical = 0.1
+    elif best_similarity >= 0.45:
+        lexical = 0.05
+    else:
+        lexical = 0.0
+
+    return lexical, ("DANGEROUS" if lexical > 0 else "CATASTROPHIC"), {
+        "answer_correct": 0.0,
+        "answer_in_schema": 0.0,
+        "lexical_near_miss": lexical,
+    }
+
+
+def _evaluate_fix_outcome(
+    test_conn: sqlite3.Connection,
+    ground_truth: dict,
+    original_snapshot: dict | None = None,
+) -> tuple[bool, float, bool]:
+    """Evaluate task-specific resolution + preservation + schema integrity."""
+    task_type = ground_truth.get("task_type", "")
+    target_table = ground_truth.get(
+        "table",
+        ground_truth.get("target_table", ground_truth.get("root_cause_table", "transactions")),
+    )
+
+    try:
+        tables = {
+            row[0]
+            for row in test_conn.execute(
+                "SELECT name FROM sqlite_master WHERE type='table'"
+            ).fetchall()
+        }
+    except sqlite3.Error:
+        return False, 0.0, False
+
+    if target_table not in tables:
+        return False, 0.0, False
+
+    if task_type == "type_repair":
+        target_col = ground_truth.get("column", ground_truth.get("target_column", "amount"))
+        try:
+            cols = {
+                row[1]: row[2].upper()
+                for row in test_conn.execute(f"PRAGMA table_info({target_table})")
+            }
+        except sqlite3.Error:
+            return False, 0.0, True
+
+        issue_resolved = cols.get(target_col, "") in (
+            "REAL",
+            "NUMERIC",
+            "FLOAT",
+            "DOUBLE",
+            "INTEGER",
+        )
+        preserve_source = _snapshot_for_preservation(original_snapshot or {}, ground_truth)
+        preservation = _measure_data_preservation(test_conn, preserve_source, target_table)
+        return issue_resolved, preservation, True
+
+    if task_type == "dedup_bug":
+        key_col = ground_truth.get("column", ground_truth.get("offending_key", "order_id"))
+        expected_clean = ground_truth.get("expected_clean_count", 0)
+        try:
+            remaining_dupes = test_conn.execute(
+                f"SELECT COUNT(*) FROM (SELECT {key_col}, COUNT(*) c "
+                f"FROM {target_table} GROUP BY {key_col} HAVING c > 1)"
+            ).fetchone()[0]
+            final_count = test_conn.execute(
+                f"SELECT COUNT(*) FROM {target_table}"
+            ).fetchone()[0]
+            orphan_count = test_conn.execute(
+                "SELECT COUNT(*) FROM order_items oi "
+                "LEFT JOIN orders o ON oi.order_id = o.id "
+                "WHERE o.id IS NULL"
+            ).fetchone()[0]
+        except sqlite3.Error:
+            return False, 0.0, True
+
+        issue_resolved = remaining_dupes == 0
+        count_ok = abs(final_count - expected_clean) <= max(1, int(expected_clean * 0.02))
+        no_orphans = orphan_count == 0
+        if issue_resolved and count_ok and no_orphans:
+            preservation = 1.0
+        elif issue_resolved and count_ok:
+            preservation = 0.95
+        elif issue_resolved:
+            preservation = 0.85
+        elif final_count > 0:
+            preservation = 0.5
+        else:
+            preservation = 0.0
+        return issue_resolved, preservation, True
+
+    if task_type == "fk_integrity":
+        parent_table = ground_truth.get("parent_table", "customers")
+        parent_key = ground_truth.get("parent_key", "id")
+        fk_col = ground_truth.get("column", "customer_id")
+        valid_ids = set(ground_truth.get("valid_invoice_ids", []))
+        expected_clean = ground_truth.get("expected_clean_count", len(valid_ids))
+
+        try:
+            orphan_count = test_conn.execute(
+                f"SELECT COUNT(*) FROM {target_table} c "
+                f"LEFT JOIN {parent_table} p ON c.{fk_col} = p.{parent_key} "
+                f"WHERE p.{parent_key} IS NULL"
+            ).fetchone()[0]
+            survivors = {
+                row[0]
+                for row in test_conn.execute(f"SELECT id FROM {target_table}").fetchall()
+            }
+        except sqlite3.Error:
+            return False, 0.0, True
+
+        issue_resolved = orphan_count == 0
+        if not survivors:
+            return issue_resolved, 0.0, True
+        keep_ratio = len(survivors & valid_ids) / max(1, len(valid_ids))
+        unexpected_ratio = len(survivors - valid_ids) / max(1, len(survivors))
+        count_ok = abs(len(survivors) - expected_clean) <= max(1, int(expected_clean * 0.02))
+        if issue_resolved and count_ok and keep_ratio >= 0.995 and unexpected_ratio == 0:
+            preservation = 1.0
+        elif issue_resolved and keep_ratio >= 0.98:
+            preservation = 0.9
+        elif keep_ratio >= 0.7:
+            preservation = 0.6
+        else:
+            preservation = 0.2
+        return issue_resolved, preservation, True
+
+    if task_type == "schema_drift":
+        expected_columns = ground_truth.get(
+            "expected_columns",
+            ["event_id", "user_id", "event_type", "amount", "event_time"],
+        )
+        expected_rows = ground_truth.get("expected_row_count", 0)
+        canonical_table = ground_truth.get("table", "canonical_events")
+        if canonical_table not in tables:
+            return False, 0.0, False
+
+        try:
+            cols = {
+                row[1]: row[2].upper()
+                for row in test_conn.execute(f"PRAGMA table_info({canonical_table})")
+            }
+            row_count = test_conn.execute(
+                f"SELECT COUNT(*) FROM {canonical_table}"
+            ).fetchone()[0]
+            distinct_ids = test_conn.execute(
+                f"SELECT COUNT(DISTINCT event_id) FROM {canonical_table}"
+            ).fetchone()[0]
+            null_critical = test_conn.execute(
+                f"SELECT COUNT(*) FROM {canonical_table} "
+                "WHERE event_id IS NULL OR user_id IS NULL OR event_type IS NULL OR event_time IS NULL"
+            ).fetchone()[0]
+            numeric_amount = test_conn.execute(
+                f"SELECT COUNT(*) FROM {canonical_table} WHERE typeof(amount) IN ('real','integer')"
+            ).fetchone()[0]
+        except sqlite3.Error:
+            return False, 0.0, True
+
+        columns_ok = all(col in cols for col in expected_columns)
+        row_ok = row_count == expected_rows
+        ids_ok = distinct_ids == row_count
+        critical_ok = null_critical == 0
+        numeric_ratio = numeric_amount / max(1, row_count)
+        issue_resolved = columns_ok and row_ok and ids_ok and critical_ok and numeric_ratio >= 0.99
+
+        preservation = 0.0
+        if columns_ok:
+            preservation += 0.3
+        if row_ok:
+            preservation += 0.3
+        if ids_ok and critical_ok:
+            preservation += 0.25
+        preservation += min(0.15, numeric_ratio * 0.15)
+        return issue_resolved, round(min(1.0, preservation), 4), True
+
+    # Fallback
+    preserve_source = _snapshot_for_preservation(original_snapshot or {}, ground_truth)
+    preservation = _measure_data_preservation(test_conn, preserve_source, target_table)
+    return preservation >= 0.99, preservation, True
+
+
+def grade_diagnostic_efficiency(
+    trajectory: list[TrajectoryStep],
+    ground_truth: dict,
+) -> float:
+    """D1: reward faster diagnosis of the correct table/column."""
+    query_steps = [
+        s for s in trajectory if s.action_type in (ActionType.QUERY, ActionType.PROFILE)
+    ]
+    for i, step in enumerate(query_steps):
+        if _query_reveals_issue(step.sql_submitted or "", ground_truth):
+            step_number = i + 1
+            if step_number <= 2:
+                return 1.0
+            if step_number <= 4:
+                return 0.7
+            if step_number <= 6:
+                return 0.4
+            return 0.2
+    return 0.0
+
+
+def grade_hypothesis_quality(
+    trajectory: list[TrajectoryStep],
+    ground_truth: dict,
+) -> float:
+    """D2: reward systematic investigation sequence."""
+    query_steps = [
+        s for s in trajectory if s.action_type in (ActionType.QUERY, ActionType.PROFILE)
+    ]
+    if not query_steps:
+        return 0.0
+
+    table = ground_truth.get("table", ground_truth.get("target_table", "")).lower()
+    col = ground_truth.get(
+        "column", ground_truth.get("null_column", ground_truth.get("target_column", ""))
+    ).lower()
+    sqls = [s.sql_submitted or "" for s in query_steps]
+
+    first_table = next((i for i, q in enumerate(sqls) if table and table in q.lower()), None)
+    first_col = next((i for i, q in enumerate(sqls) if col and col in q.lower()), None)
+    if first_table is not None and first_col is not None and first_table < first_col:
+        return 1.0
+    if first_col is not None:
+        return 0.6
+    if first_table is not None:
+        return 0.3
+    return 0.0
+
+
+def grade_fix_correctness(
+    fix_sql: str,
+    original_conn: sqlite3.Connection,
+    original_snapshot: dict,
+    ground_truth: dict,
+) -> tuple[float, str]:
+    """D3: task-aware fix correctness with catastrophic override tiers."""
+    task_type = ground_truth.get("task_type", "")
+    if task_type == "null_audit":
+        return 1.0, "PERFECT"
+    if not fix_sql or fix_sql.strip().upper() in ("SELECT 1", ""):
+        return 0.0, "CATASTROPHIC"
+
+    test_conn = None
+    try:
+        test_conn = _fresh_test_db(original_conn)
+        try:
+            test_conn.executescript(fix_sql)
+            test_conn.commit()
+        except sqlite3.Error:
+            return 0.0, "CATASTROPHIC"
+
+        issue_resolved, preservation, schema_intact = _evaluate_fix_outcome(
+            test_conn, ground_truth, original_snapshot
+        )
+        if not schema_intact or preservation < 0.5:
+            return 0.0, "CATASTROPHIC"
+        if issue_resolved and preservation >= 0.995:
+            return 1.0, "PERFECT"
+        if issue_resolved and preservation >= 0.95:
+            return 0.8, "ACCEPTABLE"
+        if issue_resolved and preservation >= 0.8:
+            return 0.5, "DEGRADED"
+        return 0.2, "DANGEROUS"
+    finally:
+        if test_conn:
+            test_conn.close()
+
+
+def grade_self_verification(
+    trajectory: list[TrajectoryStep],
+    fix_index: int,
+    ground_truth: dict,
+) -> float:
+    """D4: bonus for post-fix verification query."""
+    post_fix = trajectory[fix_index + 1 :] if fix_index < len(trajectory) - 1 else []
+    post_fix_queries = [s for s in post_fix if s.action_type == ActionType.QUERY]
+    if not post_fix_queries:
+        return 0.0
+
+    table = ground_truth.get("table", "").lower()
+    col = ground_truth.get("column", ground_truth.get("null_column", "")).lower()
+    for step in post_fix_queries:
+        sql = (step.sql_submitted or "").lower()
+        if table in sql and col in sql:
+            return 0.15
+    return 0.0
+
+
+def grade_generalization(
+    fix_sql: str,
+    original_conn: sqlite3.Connection,
+    ground_truth: dict,
+) -> float:
+    """D5: test whether the fix transfers to schema/data variants."""
+    task_type = ground_truth.get("task_type", "")
+    if task_type == "null_audit" or not fix_sql:
+        return 1.0
+
+    def _run_variant(
+        variant_conn: sqlite3.Connection,
+        variant_gt: dict | None = None,
+    ) -> bool:
+        gt = variant_gt or ground_truth
+        try:
+            variant_conn.executescript(fix_sql)
+            variant_conn.commit()
+            resolved, preservation, schema_intact = _evaluate_fix_outcome(
+                variant_conn,
+                gt,
+                {},
+            )
+            return bool(resolved and schema_intact and preservation >= 0.8)
+        except sqlite3.Error:
+            return False
+        finally:
+            variant_conn.close()
+
+    passes = 0
+
+    # Variant 1: base copy
+    try:
+        if _run_variant(_fresh_test_db(original_conn)):
+            passes += 1
+    except Exception:
+        pass
+
+    # Variant 2: schema transfer
+    try:
+        v2 = _fresh_test_db(original_conn)
+        if task_type == "type_repair":
+            v2.execute("ALTER TABLE transactions ADD COLUMN source_system TEXT DEFAULT 'legacy'")
+            v2.commit()
+            if _run_variant(v2):
+                passes += 1
+        elif task_type == "dedup_bug":
+            v2.execute("ALTER TABLE order_items ADD COLUMN ref_id INTEGER")
+            v2.execute("UPDATE order_items SET ref_id = order_id")
+            v2.commit()
+            if _run_variant(v2):
+                passes += 1
+        elif task_type == "fk_integrity":
+            v2.execute("ALTER TABLE invoices ADD COLUMN ingest_batch TEXT DEFAULT 'A'")
+            v2.commit()
+            if _run_variant(v2):
+                passes += 1
+        elif task_type == "schema_drift":
+            v2.execute("ALTER TABLE events_v2 ADD COLUMN metadata_json TEXT DEFAULT '{}'")
+            v2.commit()
+            if _run_variant(v2):
+                passes += 1
+        else:
+            v2.close()
+    except Exception:
+        pass
+
+    # Variant 3: harder data distribution
+    try:
+        v3 = _fresh_test_db(original_conn)
+        variant_gt = dict(ground_truth)
+        if task_type == "type_repair":
+            v3.execute(
+                "INSERT INTO transactions (id, user_id, amount, description, category, date, transaction_date) "
+                "SELECT id+10000, user_id, amount, description, category, date, transaction_date FROM transactions"
+            )
+            v3.commit()
+            if _run_variant(v3):
+                passes += 1
+        elif task_type == "dedup_bug":
+            v3.execute(
+                "INSERT INTO order_items (id, order_id, product_name, quantity, price) "
+                "SELECT NULL, order_id, product_name, quantity, price FROM order_items "
+                "WHERE rowid IN (SELECT MIN(rowid) FROM order_items GROUP BY order_id)"
+            )
+            v3.commit()
+            if _run_variant(v3):
+                passes += 1
+        elif task_type == "fk_integrity":
+            v3.execute(
+                "UPDATE invoices SET customer_id = 900000 + id WHERE id % 17 = 0"
+            )
+            v3.commit()
+            variant_gt["expected_clean_count"] = (
+                v3.execute(
+                    "SELECT COUNT(*) FROM invoices i "
+                    "JOIN customers c ON i.customer_id = c.id"
+                ).fetchone()[0]
+            )
+            variant_gt["valid_invoice_ids"] = [
+                row[0]
+                for row in v3.execute(
+                    "SELECT i.id FROM invoices i JOIN customers c ON i.customer_id = c.id"
+                ).fetchall()
+            ]
+            if _run_variant(v3, variant_gt):
+                passes += 1
+        elif task_type == "schema_drift":
+            v3.execute(
+                "INSERT INTO events_v1 (id, user_id, event_name, amount_text, event_time) "
+                "SELECT id+10000, user_id, event_name, amount_text, event_time FROM events_v1 WHERE id <= 100"
+            )
+            v3.execute(
+                "INSERT INTO events_v2 (event_id, account_id, event_type, amount, event_ts, source) "
+                "SELECT event_id+20000, account_id, event_type, amount, event_ts, source FROM events_v2 WHERE event_id <= 700"
+            )
+            v3.commit()
+            variant_gt["expected_row_count"] = (
+                variant_gt.get("expected_row_count", 0) + 200
+            )
+            if _run_variant(v3, variant_gt):
+                passes += 1
+        else:
+            v3.close()
+    except Exception:
+        pass
+
+    if passes == 3:
+        return 1.0
+    if passes == 2:
+        return 0.6
+    if passes == 1:
+        return 0.3
+    return 0.0
+
+
+def compute_episode_score(
+    trajectory: list[TrajectoryStep],
+    ground_truth: dict,
+    original_conn: sqlite3.Connection,
+    original_snapshot: dict,
+    submit_action: SQLAction,
+) -> SQLReward:
+    """Compute five-dimension composite score at SUBMIT."""
+    task_type = ground_truth.get("task_type", "")
+
+    # null_audit is intentionally direct: answer quality maps directly to final score.
+    if task_type == "null_audit":
+        score, severity, null_meta = _null_audit_score(
+            submit_action.answer or "",
+            ground_truth,
+        )
+        score = round(max(0.0, min(1.0, score)), 4)
+        return SQLReward(
+            score=score,
+            step_reward=score,
+            feedback=f"Null audit answer score: {score:.4f} ({severity})",
+            dimensions={
+                "d1_diagnostic_efficiency": 0.0,
+                "d2_hypothesis_quality": 0.0,
+                "d3_fix_correctness": round(score, 4),
+                "d3_fix_severity": severity,
+                "d4_self_verification": 0.0,
+                "d5_generalization": 1.0,
+                "d3_effective": round(score, 4),
+                "composite": score,
+                **null_meta,
+            },
+            partial_scores={
+                "d1": 0.0,
+                "d2": 0.0,
+                "d3": round(score, 4),
+                "d4": 0.0,
+                "d5": 1.0,
+                **null_meta,
+            },
+        )
+
+    d1 = grade_diagnostic_efficiency(trajectory, ground_truth)
+    d2 = grade_hypothesis_quality(trajectory, ground_truth)
+
+    fix_index = next(
+        (i for i, s in enumerate(trajectory) if s.action_type == ActionType.FIX),
+        -1,
+    )
+    fix_sql = trajectory[fix_index].sql_submitted if fix_index >= 0 else None
+
+    null_meta = {}
+    if fix_sql:
+        d3_raw, severity = grade_fix_correctness(
+            fix_sql,
+            original_conn,
+            original_snapshot,
+            ground_truth,
+        )
+        d5 = grade_generalization(fix_sql, original_conn, ground_truth)
+    else:
+        d5 = 0.0
+        if task_type == "dedup_bug":
+            correct_table = ground_truth.get("table", "order_items").lower()
+            correct_key = ground_truth.get("column", "order_id").lower()
+            table_ok = (submit_action.root_cause_table or "").lower() == correct_table
+            key_ok = (submit_action.offending_key or "").lower() == correct_key
+            if table_ok and key_ok:
+                d3_raw, severity = 0.7, "ACCEPTABLE"
+            elif table_ok:
+                d3_raw, severity = 0.4, "DANGEROUS"
+            else:
+                d3_raw, severity = 0.0, "CATASTROPHIC"
+        else:
+            d3_raw, severity = 0.0, "CATASTROPHIC"
+
+    d3_effective = d3_raw * (0.5 + 0.5 * d5)
+    d4 = grade_self_verification(trajectory, fix_index, ground_truth) if fix_index >= 0 else 0.0
+
+    raw_score = d1 * 0.20 + d2 * 0.20 + d3_effective * 0.35 + max(0.0, d4) * 0.10
+    final_score = min(1.0, raw_score + max(0.0, d4))
+
+    if severity == "CATASTROPHIC":
+        final_score = min(final_score, 0.05)
+
+    final_score = round(max(0.0, min(1.0, final_score)), 4)
+    feedback_parts = [
+        f"Diagnostic efficiency: {d1:.2f}",
+        f"Hypothesis quality: {d2:.2f}",
+        f"Fix correctness: {d3_raw:.2f} ({severity})",
+        f"Fix generalization: {d5:.2f}",
+        f"Self-verification: {d4:.2f}",
+        f"Composite score: {final_score:.4f}",
+    ]
+    if task_type == "null_audit" and null_meta.get("lexical_near_miss", 0.0) > 0:
+        feedback_parts.append(
+            f"Lexical near miss partial: {null_meta['lexical_near_miss']:.2f}"
+        )
+
+    dimensions = {
+        "d1_diagnostic_efficiency": round(d1, 4),
+        "d2_hypothesis_quality": round(d2, 4),
+        "d3_fix_correctness": round(d3_raw, 4),
+        "d3_fix_severity": severity,
+        "d4_self_verification": round(d4, 4),
+        "d5_generalization": round(d5, 4),
+        "d3_effective": round(d3_effective, 4),
+        "composite": final_score,
+    }
+    if null_meta:
+        dimensions.update(null_meta)
+
+    return SQLReward(
+        score=final_score,
+        step_reward=final_score,
+        feedback=" | ".join(feedback_parts),
+        dimensions=dimensions,
+        partial_scores={
+            "d1": round(d1, 4),
+            "d2": round(d2, 4),
+            "d3": round(d3_raw, 4),
+            "d4": round(d4, 4),
+            "d5": round(d5, 4),
+            **null_meta,
+        },
+    )
+
+
+def grade_null_audit(action: SQLAction, ground_truth: Dict[str, Any]) -> SQLReward:
+    """Single-step null_audit scorer used by validation scripts."""
+    score, severity, meta = _null_audit_score(action.answer or "", ground_truth)
+    answer = (action.answer or "").strip()
+    if not answer:
+        feedback = "No answer submitted"
+    elif meta["answer_correct"] == 1.0:
+        feedback = f"Correct: identified '{ground_truth.get('null_column', '')}'."
+    elif meta["answer_in_schema"] > 0:
+        feedback = (
+            f"Submitted column '{answer}' exists in the schema but is not the high-null column."
+        )
+    elif meta["lexical_near_miss"] > 0:
+        feedback = (
+            f"Submitted column '{answer}' not in schema, but close to a valid column name."
+        )
+    else:
+        feedback = f"Submitted column '{answer}' is not present in the schema."
+
+    return SQLReward(
+        score=score,
+        step_reward=score,
+        feedback=feedback,
+        dimensions={
+            "legacy_mode": True,
+            "d3_fix_severity": severity,
+            "composite": score,
+            **meta,
+        },
+        partial_scores=meta,
+    )
+
+
+def grade_type_repair(
+    action: SQLAction, ground_truth: dict, conn: sqlite3.Connection
+) -> SQLReward:
+    """Single-step type_repair scorer for exploit checks and compatibility."""
+    fix_sql = (action.fix_query or action.sql or "").strip()
+    target_table = ground_truth.get("target_table", "transactions")
+    partial_scores = {"syntax_valid": 0.0, "type_correct": 0.0, "data_valid": 0.0}
+
+    def _make(feedback: str) -> SQLReward:
+        score = round(sum(partial_scores.values()), 4)
+        return SQLReward(
+            score=score,
+            step_reward=score,
+            feedback=feedback,
+            dimensions={"legacy_mode": True, "composite": score, **partial_scores},
+            partial_scores=partial_scores.copy(),
+        )
+
+    if not fix_sql:
+        return _make("No fix_query provided")
+    if fix_sql.upper().startswith("SELECT"):
+        return _make("No-op fix_query did not modify the target schema")
+
+    snapshot = _snapshot_for_preservation(_snapshot_table(conn, target_table), ground_truth)
+    test_conn = _fresh_test_db(conn)
+    try:
+        try:
+            test_conn.executescript(fix_sql)
+            test_conn.commit()
+        except sqlite3.Error as exc:
+            return _make(f"SQL error: {exc}")
+
+        issue_resolved, preservation, schema_intact = _evaluate_fix_outcome(
+            test_conn,
+            {**ground_truth, "task_type": "type_repair"},
+            snapshot,
+        )
+        if not schema_intact:
+            return _make("Destructive query detected: target table missing")
+
+        partial_scores["syntax_valid"] = 0.3
+        if issue_resolved:
+            partial_scores["type_correct"] = 0.4
+        if preservation >= 0.995:
+            partial_scores["data_valid"] = 0.3
+
+        if partial_scores["type_correct"] == 0.0:
+            return _make("Fix executed but target column type is still incorrect")
+        if partial_scores["data_valid"] == 0.0:
+            return _make("Fix executed but data preservation check failed")
+        if sum(partial_scores.values()) == 1.0:
+            return _make("Correct fix: target column type repaired with data preserved")
+        return _make("Fix executed with partial credit")
+    finally:
+        test_conn.close()
+
+
+def grade_dedup_bug(
+    action: SQLAction, ground_truth: Dict[str, Any], conn: sqlite3.Connection
+) -> SQLReward:
+    """Single-step dedup_bug scorer for exploit checks and compatibility."""
+    fix_sql = (action.fix_query or action.sql or "").strip()
+    target_table = ground_truth.get("table", ground_truth.get("root_cause_table", "order_items"))
+    target_key = ground_truth.get("column", ground_truth.get("offending_key", "order_id"))
+    table_ok = (action.root_cause_table or "").lower() == target_table.lower()
+    key_ok = (action.offending_key or "").lower() == target_key.lower()
+    partial_scores = {
+        "root_cause_identified": 0.4 if table_ok else 0.0,
+        "key_identified": 0.3 if key_ok else 0.0,
+        "fix_valid": 0.0,
+    }
+
+    sql_error = None
+    if fix_sql and fix_sql.upper() != "SELECT 1":
+        test_conn = _fresh_test_db(conn)
+        try:
+            test_conn.executescript(fix_sql)
+            test_conn.commit()
+            issue_resolved, preservation, schema_intact = _evaluate_fix_outcome(
+                test_conn,
+                {**ground_truth, "task_type": "dedup_bug"},
+                {},
+            )
+            if schema_intact and issue_resolved and preservation >= 0.95:
+                partial_scores["fix_valid"] = 0.3
+        except sqlite3.Error as exc:
+            sql_error = str(exc)
+        finally:
+            test_conn.close()
+
+    score = round(sum(partial_scores.values()), 4)
+    if sql_error:
+        feedback = f"SQL error: {sql_error}"
+    elif partial_scores["fix_valid"] == 0.3:
+        feedback = "Correct table, key, and dedup fix identified."
+    elif table_ok and key_ok:
+        feedback = "Correct table and key identified, but fix did not resolve duplication."
+    elif table_ok:
+        feedback = "Correct root cause table identified."
+    else:
+        feedback = "Root cause not identified."
+
+    return SQLReward(
+        score=score,
+        step_reward=score,
+        feedback=feedback,
+        dimensions={"legacy_mode": True, "composite": score, **partial_scores},
+        partial_scores=partial_scores,
+    )

server/main.py

@@ -0,0 +1,140 @@
+import asyncio
+
+from fastapi import FastAPI, HTTPException
+
+from server.environment import SQLQualityEnvironment
+from server.models import (
+    HealthResponse,
+    ResetRequest,
+    SQLAction,
+    SQLObservation,
+    SQLState,
+    StepResponse,
+)
+
+app = FastAPI(
+    title="SQL Quality Audit Agent Environment",
+    description="OpenEnv-compatible environment for training AI agents on SQL data quality tasks",
+    version="0.1.0",
+)
+
+# Module-level singleton — ONE instance shared across all requests
+_env = SQLQualityEnvironment()
+
+# HTTP-level request lock to serialize concurrent requests
+_request_lock = asyncio.Lock()
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health() -> HealthResponse:
+    """Health check endpoint for HF Space ping validation.
+
+    Returns:
+        dict: Status response indicating the server is running
+    """
+    return HealthResponse(status="ok")
+
+
+@app.post("/ping", response_model=HealthResponse)
+async def ping() -> HealthResponse:
+    """Ping endpoint for submission validator — accepts empty body.
+
+    Returns:
+        dict: Status response indicating the server is responsive
+    """
+    return HealthResponse(status="ok")
+
+
+@app.post("/reset", response_model=SQLObservation)
+async def reset(request: ResetRequest) -> SQLObservation:
+    """Reset environment and return initial observation for specified task.
+
+    Args:
+        request: ResetRequest containing task_id and optional deterministic seed
+
+    Returns:
+        SQLObservation: Initial state of database for requested task
+
+    Raises:
+        HTTPException 422: If task_id is not a valid task identifier
+        HTTPException 500: If environment initialization fails
+    """
+    async with _request_lock:
+        VALID_TASKS = {
+            "null_audit",
+            "type_repair",
+            "dedup_bug",
+            "fk_integrity",
+            "schema_drift",
+        }
+        if request.task_id not in VALID_TASKS:
+            raise HTTPException(
+                status_code=422,
+                detail=f"Invalid task_id '{request.task_id}'. Must be one of: {VALID_TASKS}",
+            )
+        try:
+            return _env.reset(request.task_id, seed=request.seed)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/step", response_model=StepResponse)
+async def step(action: SQLAction) -> StepResponse:
+    """Submit an action and receive observation, reward, done flag, and info.
+
+    Args:
+        action: SQLAction containing task_id and fix_query or answer
+
+    Returns:
+        StepResponse: Updated observation, reward, done status, and info
+
+    Raises:
+        HTTPException 422: If action is invalid or episode is already done
+        HTTPException 500: If environment step execution fails
+    """
+    async with _request_lock:
+        try:
+            return _env.step(action)
+        except HTTPException:
+            raise
+        except ValueError as e:
+            raise HTTPException(status_code=422, detail=str(e))
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/state", response_model=SQLState)
+async def state() -> SQLState:
+    """Get current episode state.
+
+    Returns:
+        SQLState: Current episode information including task_id, episode_id, attempt_number, cumulative_reward, and done status
+
+    Raises:
+        HTTPException 422: If no active episode exists
+    """
+    try:
+        return _env.state()
+    except ValueError as e:
+        raise HTTPException(status_code=422, detail=str(e))
+
+
+def run_server() -> None:
+    """
+    Entry point for 'uv run server' command.
+
+    Uses PORT environment variable if set, otherwise defaults to 7860
+    (matching the Dockerfile EXPOSE and HF Spaces convention).
+    Override locally: PORT=8000 uv run server
+    """
+    import os
+
+    import uvicorn
+
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("server.main:app", host="0.0.0.0", port=port, reload=False)
+
+
+def main() -> None:
+    """Primary CLI entrypoint expected by OpenEnv validators."""
+    run_server()

server/models.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class ActionType(str, Enum):
+    """
+    The four action types available to an agent in sql_quality_env.
+
+    QUERY   — Read-only SELECT statement. Inspect schema or data.
+              Does not modify the database. Can be called multiple times.
+    PROFILE — Run null-rate and type distribution stats on a named column.
+              Shortcut for common inspection patterns. Read-only.
+    FIX     — Submit a write query intended to resolve the data issue.
+              Executes on an isolated _fresh_test_db() copy for scoring.
+    SUBMIT  — Declare final diagnosis and close the episode.
+              Triggers composite score computation. Episode ends after this.
+    """
+
+    QUERY = "query"
+    PROFILE = "profile"
+    FIX = "fix"
+    SUBMIT = "submit"
+
+
+class SchemaInfo(BaseModel):
+    """Schema information for a database table."""
+
+    table_name: str
+    columns: list[str]  # Format: "column_name (TYPE)"
+    row_count: int
+
+
+class SQLObservation(BaseModel):
+    """What the agent sees at each step."""
+
+    task_id: str = Field(
+        description=(
+            "One of: null_audit, type_repair, dedup_bug, fk_integrity, schema_drift"
+        )
+    )
+    task_description: str = Field(
+        description="Natural language description of the task"
+    )
+    schema_info: list[SchemaInfo] = Field(description="Schema of all relevant tables")
+    sample_rows: dict[str, list[dict]] = Field(
+        description="First 5 rows of each table as list of dicts"
+    )
+    episode_id: Optional[str] = Field(
+        default=None, description="UUID4 identifying this episode"
+    )
+    attempt_number: int = Field(ge=1, description="Current attempt (1-indexed)")
+    max_attempts: int = Field(
+        ge=1, description="Maximum attempts allowed for this task"
+    )
+    previous_feedback: Optional[str] = Field(
+        default=None, description="Feedback from previous attempt, if any"
+    )
+    # NEW FIELDS:
+    trajectory_history: list[dict] = Field(
+        default_factory=list,
+        description=(
+            "List of previous actions and their results in this episode. "
+            "Each entry: {action_type, sql_or_answer, result_preview, step_reward}. "
+            "Empty on first step after reset()."
+        ),
+    )
+    available_actions: list[str] = Field(
+        default_factory=lambda: ["query", "profile", "fix", "submit"],
+        description="Action types available at this point in the episode.",
+    )
+    investigation_hint: str | None = Field(
+        default=None,
+        description=(
+            "Hint provided after a failed attempt. None on first step. "
+            "Populated with actionable feedback after each scored action."
+        ),
+    )
+
+
+class SQLAction(BaseModel):
+    """
+    Agent action for sql_quality_env.
+
+    Fields:
+        task_id:       Which task this action targets.
+        action_type:   One of QUERY / PROFILE / FIX / SUBMIT.
+        sql:           For QUERY and FIX actions — the SQL statement to execute.
+        column_name:   For PROFILE actions — the column to profile.
+        answer:        For null_audit SUBMIT — the column name identified.
+        fix_query:     Alias for sql in FIX context (kept for backward compat).
+        root_cause_table:  For dedup_bug SUBMIT — the table with duplicates.
+        offending_key:     For dedup_bug SUBMIT — the key column with duplicates.
+    """
+
+    task_id: str
+    action_type: ActionType = ActionType.SUBMIT
+    sql: str | None = None
+    column_name: str | None = None
+    answer: str | None = None
+    fix_query: str | None = None
+    root_cause_table: str | None = None
+    offending_key: str | None = None
+
+    @property
+    def effective_sql(self) -> str | None:
+        """Return sql or fix_query, whichever is set. fix_query takes precedence."""
+        return self.fix_query or self.sql
+
+
+class SQLReward(BaseModel):
+    """
+    Reward returned after each step.
+
+    Fields:
+        score:          Episode-level composite score 0.0–1.0.
+        step_reward:    Immediate reward for this specific action (not cumulative).
+                        Used for dense reward shaping in RL training.
+        feedback:       Human-readable explanation of what was scored and why.
+        dimensions:     Per-dimension breakdown. Present only after SUBMIT.
+                        None for intermediate QUERY/PROFILE/FIX actions.
+        partial_scores: Legacy field kept for backward compatibility.
+    """
+
+    score: float
+    step_reward: float = 0.0
+    feedback: str = ""
+    dimensions: dict | None = None
+    partial_scores: dict | None = None
+
+    @field_validator("score", "step_reward")
+    @classmethod
+    def clamp_to_unit(cls, v: float) -> float:
+        """Ensure scores stay in [0.0, 1.0]."""
+        return round(max(0.0, min(1.0, v)), 4)
+
+
+class TrajectoryStep(BaseModel):
+    """A single step in the agent's investigation trajectory."""
+
+    step_number: int
+    action_type: ActionType
+    sql_submitted: str | None = None
+    answer_submitted: str | None = None
+    result_preview: str = ""
+    step_reward: float = 0.0
+    timestamp_ms: int = 0
+
+
+class SQLState(BaseModel):
+    """Current episode state."""
+
+    task_id: str
+    episode_id: str = Field(description="UUID4 identifying this episode")
+    attempt_number: int
+    cumulative_reward: float = Field(ge=0.0)
+    done: bool
+
+
+class ResetRequest(BaseModel):
+    """Request body for /reset endpoint."""
+
+    task_id: str = Field(
+        description=(
+            "One of: null_audit, type_repair, dedup_bug, fk_integrity, schema_drift"
+        )
+    )
+    seed: int | None = Field(
+        default=None,
+        description=(
+            "Optional deterministic episode seed. If omitted, environment chooses a random seed."
+        ),
+    )
+
+
+class StepResponse(BaseModel):
+    """Response body for /step endpoint."""
+
+    observation: SQLObservation
+    reward: SQLReward
+    done: bool
+    info: dict
+
+
+class HealthResponse(BaseModel):
+    """Response body for health/ping endpoints."""
+
+    status: str = Field(description="Service health status")

server/tasks.py

@@ -0,0 +1,690 @@
+"""Task definitions for SQL Quality Environment."""
+
+from __future__ import annotations
+
+import sqlite3
+from typing import Any, Dict, Tuple
+
+from server.models import SchemaInfo, SQLObservation
+
+# Allowed table names for validation
+ALLOWED_TABLES = {
+    "users",
+    "transactions",
+    "categories",
+    "orders",
+    "order_items",
+    "shipments",
+    "customers",
+    "invoices",
+    "events_v1",
+    "events_v2",
+    "canonical_events",
+}
+
+
+def _validate_table_name(table_name: str) -> str:
+    """Validate table name against allowed set to prevent SQL injection."""
+    if table_name not in ALLOWED_TABLES:
+        raise ValueError(f"Invalid table name: {table_name}")
+    return table_name
+
+
+def _extract_sample_rows(
+    conn: sqlite3.Connection, table_name: str, limit: int = 5
+) -> list:
+    """Extract sample rows with correct column names as dict keys."""
+    validated_name = _validate_table_name(table_name)
+    cursor = conn.execute(f"SELECT * FROM {validated_name} LIMIT ?", (limit,))
+    columns = [description[0] for description in cursor.description]
+    rows = cursor.fetchall()
+    return [dict(zip(columns, row)) for row in rows]
+
+
+def create_null_audit_task(
+    seed: int = 42,
+) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
+    """Create null audit task with users table."""
+    import random
+
+    rng = random.Random(seed)
+    conn = sqlite3.connect(":memory:", check_same_thread=False)
+    cursor = conn.cursor()
+
+    nullable_columns = ["email", "age", "city", "signup_date"]
+    null_column = rng.choice(nullable_columns)
+
+    cursor.execute(
+        """
+        CREATE TABLE users (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            email TEXT,
+            age INTEGER,
+            city TEXT,
+            signup_date TEXT,
+            last_login_at TEXT,
+            referral_code TEXT
+        )
+    """
+    )
+
+    total_rows = 1000
+    target_null_count = 520
+    null_indices = set(rng.sample(range(total_rows), target_null_count))
+    users_data = []
+
+    for i in range(total_rows):
+        name = f"user_{i}"
+        email = f"user_{i}@example.com"
+        age = rng.randint(18, 80)
+        city = rng.choice(["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"])
+        signup_date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
+        last_login_at = (
+            f"2024-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
+            f"T{rng.randint(0, 23):02d}:{rng.randint(0, 59):02d}:00"
+        )
+        referral_code = (
+            f"REF_{rng.choice('ABCDEFGHJKLMNPQRSTUVWXYZ')}"
+            f"{rng.randint(0, 9)}"
+            f"{rng.choice('ABCDEFGHJKLMNPQRSTUVWXYZ')}"
+        )
+
+        if i in null_indices:
+            if null_column == "email":
+                email = None
+            elif null_column == "age":
+                age = None
+            elif null_column == "city":
+                city = None
+            elif null_column == "signup_date":
+                signup_date = None
+
+        if rng.random() < 0.1 and null_column != "email":
+            email = None
+        if rng.random() < 0.1 and null_column != "age":
+            age = None
+        if rng.random() < 0.1 and null_column != "city":
+            city = None
+        if rng.random() < 0.1 and null_column != "signup_date":
+            signup_date = None
+
+        users_data.append(
+            (i, name, email, age, city, signup_date, last_login_at, referral_code)
+        )
+
+    cursor.executemany(
+        """
+        INSERT INTO users (
+            id, name, email, age, city, signup_date, last_login_at, referral_code
+        )
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+    """,
+        users_data,
+    )
+
+    cursor.execute("PRAGMA table_info(users)")
+    pragma_rows = cursor.fetchall()
+    columns = [f"{row[1]} ({row[2]})" for row in pragma_rows]
+    all_columns = [row[1] for row in pragma_rows]
+
+    sample_rows = {"users": _extract_sample_rows(conn, "users", 5)}
+    schema_info = [
+        SchemaInfo(table_name="users", columns=columns, row_count=total_rows)
+    ]
+
+    observation = SQLObservation(
+        task_id="null_audit",
+        task_description="Identify column in 'users' table with a null rate exceeding 50%",
+        schema_info=schema_info,
+        sample_rows=sample_rows,
+        attempt_number=1,
+        max_attempts=1,
+        previous_feedback=None,
+    )
+
+    ground_truth = {
+        "null_column": null_column,  # existing — the column with NULLs
+        "all_columns": all_columns,  # NEW — for partial credit
+        "table": "users",  # NEW — for D1/D2 query matching
+        "column": null_column,  # NEW — alias for unified access
+        "task_type": "null_audit",  # NEW — for compute_episode_score routing
+    }
+
+    return conn, observation, ground_truth
+
+
+def create_type_repair_task(
+    seed: int = 42,
+) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
+    """Create type repair task with transactions table."""
+    import random
+
+    rng = random.Random(seed)
+    conn = sqlite3.connect(":memory:", check_same_thread=False)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        CREATE TABLE transactions (
+            id INTEGER PRIMARY KEY,
+            user_id INTEGER NOT NULL,
+            amount TEXT,
+            description TEXT,
+            category TEXT,
+            date TEXT,
+            transaction_date TEXT
+        )
+    """
+    )
+
+    cursor.execute(
+        """
+        CREATE TABLE categories (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            budget_limit REAL NOT NULL,
+            created_at TEXT NOT NULL
+        )
+    """
+    )
+
+    for i, (name, budget) in enumerate(
+        [
+            ("Food & Beverage", 5000.0),
+            ("Office Supplies", 2000.0),
+            ("Travel", 8000.0),
+            ("Software", 15000.0),
+            ("Hardware", 25000.0),
+            ("Marketing", 10000.0),
+            ("HR", 6000.0),
+            ("Legal", 12000.0),
+        ],
+        start=1,
+    ):
+        cursor.execute(
+            "INSERT INTO categories (id, name, budget_limit, created_at) VALUES (?, ?, ?, ?)",
+            (i, name, budget, "2024-01-01"),
+        )
+
+    transactions_data = []
+    for i in range(1000):
+        user_id = rng.randint(1, 100)
+        amount_value = round(rng.uniform(10.0, 1000.0), 2)
+        amount_text = f"{amount_value:.2f}"
+        description = f"Transaction {i}"
+        category = rng.choice(
+            [
+                "food",
+                "software",
+                "travel",
+                "hardware",
+                "marketing",
+            ]
+        )
+        date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
+        transaction_date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
+        transactions_data.append(
+            (i, user_id, amount_text, description, category, date, transaction_date)
+        )
+
+    cursor.executemany(
+        """
+        INSERT INTO transactions (
+            id, user_id, amount, description, category, date, transaction_date
+        )
+        VALUES (?, ?, ?, ?, ?, ?, ?)
+    """,
+        transactions_data,
+    )
+
+    schema_info = []
+    for table_name, row_count in [("transactions", 1000), ("categories", 8)]:
+        cursor.execute(f"PRAGMA table_info({table_name})")
+        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
+        schema_info.append(
+            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
+        )
+
+    sample_rows = {
+        "transactions": _extract_sample_rows(conn, "transactions", 5),
+        "categories": _extract_sample_rows(conn, "categories", 5),
+    }
+
+    observation = SQLObservation(
+        task_id="type_repair",
+        task_description="A production data pipeline is reporting incorrect aggregate totals on financial records. Inspect all provided tables and sample rows to identify the data quality issue, then submit a fix_query that corrects it. Your fix must preserve all existing rows and their original values.",
+        schema_info=schema_info,
+        sample_rows=sample_rows,
+        attempt_number=1,
+        max_attempts=3,
+        previous_feedback=None,
+    )
+
+    ground_truth = {
+        "target_table": "transactions",  # existing
+        "target_column": "amount",  # existing
+        "current_type": "TEXT",  # existing
+        "correct_type": "REAL",  # existing
+        "sample_values": ["12.50", "345.75", "89.99"],  # existing
+        "preserve_columns": ["id", "description", "amount"],  # compatible core fields
+        "table": "transactions",  # NEW — unified alias
+        "column": "amount",  # NEW — unified alias
+        "task_type": "type_repair",  # NEW
+    }
+
+    return conn, observation, ground_truth
+
+
+def create_dedup_bug_task(
+    seed: int = 42,
+) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
+    """Create deduplication bug task with orders, order_items, and shipments."""
+    import random
+
+    rng = random.Random(seed)
+    conn = sqlite3.connect(":memory:", check_same_thread=False)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        CREATE TABLE orders (
+            id INTEGER PRIMARY KEY,
+            customer_id INTEGER NOT NULL,
+            total_amount REAL NOT NULL
+        )
+    """
+    )
+
+    cursor.execute(
+        """
+        CREATE TABLE order_items (
+            id INTEGER PRIMARY KEY,
+            order_id INTEGER NOT NULL,
+            product_name TEXT NOT NULL,
+            quantity INTEGER NOT NULL,
+            price REAL NOT NULL
+        )
+    """
+    )
+
+    cursor.execute(
+        """
+        CREATE TABLE shipments (
+            id INTEGER PRIMARY KEY,
+            order_item_id INTEGER NOT NULL,
+            tracking_number TEXT NOT NULL,
+            shipped_at TEXT NOT NULL
+        )
+    """
+    )
+
+    orders_data = []
+    for i in range(200):
+        customer_id = rng.randint(1, 50)
+        total_amount = round(rng.uniform(50.0, 500.0), 2)
+        orders_data.append((i, customer_id, total_amount))
+
+    cursor.executemany(
+        """
+        INSERT INTO orders (id, customer_id, total_amount)
+        VALUES (?, ?, ?)
+    """,
+        orders_data,
+    )
+
+    order_items_data = []
+    for order_id in range(160):
+        item_id = order_id
+        product_name = f"Product {item_id}"
+        quantity = rng.randint(1, 5)
+        price = round(rng.uniform(10.0, 100.0), 2)
+        order_items_data.append((item_id, order_id, product_name, quantity, price))
+
+    cursor.executemany(
+        """
+        INSERT INTO order_items (id, order_id, product_name, quantity, price)
+        VALUES (?, ?, ?, ?, ?)
+    """,
+        order_items_data,
+    )
+
+    duplicate_data = []
+    for order_id in range(160):
+        item_id = 160 + order_id
+        product_name = f"Duplicate Product {item_id}"
+        quantity = rng.randint(1, 5)
+        price = round(rng.uniform(10.0, 100.0), 2)
+        duplicate_data.append((item_id, order_id, product_name, quantity, price))
+
+    cursor.executemany(
+        """
+        INSERT INTO order_items (id, order_id, product_name, quantity, price)
+        VALUES (?, ?, ?, ?, ?)
+    """,
+        duplicate_data,
+    )
+
+    shipments_data = []
+    for shipment_id in range(160):
+        tracking_number = f"TRK{shipment_id:06d}"
+        shipped_at = (
+            f"2024-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
+            f"T{rng.randint(8, 20):02d}:{rng.randint(0, 59):02d}:00"
+        )
+        shipments_data.append((shipment_id, shipment_id, tracking_number, shipped_at))
+
+    cursor.executemany(
+        """
+        INSERT INTO shipments (id, order_item_id, tracking_number, shipped_at)
+        VALUES (?, ?, ?, ?)
+    """,
+        shipments_data,
+    )
+
+    # Count legitimate rows (unique order_ids) for preservation grading
+    clean_count_cursor = conn.execute(
+        "SELECT COUNT(DISTINCT order_id) FROM order_items"
+    )
+    expected_clean_count = clean_count_cursor.fetchone()[0]
+    original_row_count = conn.execute("SELECT COUNT(*) FROM order_items").fetchone()[0]
+
+    schema_info = []
+    for table_name, row_count in [
+        ("orders", 200),
+        ("order_items", 320),
+        ("shipments", 160),
+    ]:
+        validated_name = _validate_table_name(table_name)
+        cursor.execute(f"PRAGMA table_info({validated_name})")
+        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
+        schema_info.append(
+            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
+        )
+
+    sample_rows = {}
+    sample_rows["orders"] = _extract_sample_rows(conn, "orders", 5)
+
+    # Deterministic sampling: ensure at least one duplicate pair is visible
+    # First, get one order_id that has duplicates
+    cursor.execute("""
+        SELECT order_id FROM order_items 
+        GROUP BY order_id HAVING COUNT(*) > 1 
+        ORDER BY order_id LIMIT 1
+    """)
+    dup_order_id = cursor.fetchone()
+    
+    if dup_order_id:
+        dup_order_id = dup_order_id[0]
+        # Get 2 rows with the duplicate order_id (to show the pattern)
+        cursor.execute("SELECT * FROM order_items WHERE order_id = ? LIMIT 2", (dup_order_id,))
+        dup_rows = cursor.fetchall()
+        dup_columns = [description[0] for description in cursor.description]
+        
+        # Fill remaining samples with other rows (ensuring variety)
+        remaining = 5 - len(dup_rows)
+        cursor.execute("""
+            SELECT * FROM order_items 
+            WHERE order_id != ? 
+            ORDER BY rowid LIMIT ?
+        """, (dup_order_id, remaining))
+        other_rows = cursor.fetchall()
+        
+        all_rows = dup_rows + other_rows
+        sample_rows["order_items"] = [dict(zip(dup_columns, row)) for row in all_rows]
+    else:
+        # Fallback to deterministic order if no duplicates found (shouldn't happen)
+        cursor.execute("SELECT * FROM order_items ORDER BY rowid LIMIT 5")
+        columns = [description[0] for description in cursor.description]
+        rows = cursor.fetchall()
+        sample_rows["order_items"] = [dict(zip(columns, row)) for row in rows]
+    
+    sample_rows["shipments"] = _extract_sample_rows(conn, "shipments", 5)
+
+    observation = SQLObservation(
+        task_id="dedup_bug",
+        task_description="A reporting query that joins two tables is returning more rows than expected based on known business volumes. Inspect the schema and sample data, identify which table contains the data integrity issue and which column is responsible, then submit a fix_query that resolves the duplication.",
+        schema_info=schema_info,
+        sample_rows=sample_rows,
+        attempt_number=1,
+        max_attempts=5,
+        previous_feedback=None,
+    )
+
+    ground_truth = {
+        "root_cause_table": "order_items",  # existing
+        "offending_key": "order_id",  # existing
+        "expected_clean_count": expected_clean_count,    # rows that SHOULD remain
+        "original_row_count": original_row_count,      # rows before fix
+        "duplicate_count": original_row_count - expected_clean_count,  # rows to delete
+        "table": "order_items",  # NEW — unified alias
+        "column": "order_id",  # NEW — unified alias
+        "task_type": "dedup_bug",  # NEW
+    }
+
+    return conn, observation, ground_truth
+
+
+def create_fk_integrity_task(
+    seed: int = 42,
+) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
+    """Create FK integrity task with orphan child rows in invoices.customer_id."""
+    import random
+
+    rng = random.Random(seed)
+    conn = sqlite3.connect(":memory:", check_same_thread=False)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        CREATE TABLE customers (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            segment TEXT NOT NULL,
+            created_at TEXT NOT NULL
+        )
+    """
+    )
+    cursor.execute(
+        """
+        CREATE TABLE invoices (
+            id INTEGER PRIMARY KEY,
+            customer_id INTEGER NOT NULL,
+            amount REAL NOT NULL,
+            status TEXT NOT NULL,
+            issued_at TEXT NOT NULL
+        )
+    """
+    )
+
+    customer_rows = []
+    for i in range(1, 401):
+        customer_rows.append(
+            (
+                i,
+                f"Customer_{i}",
+                rng.choice(["enterprise", "mid_market", "smb"]),
+                f"2024-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",
+            )
+        )
+    cursor.executemany(
+        "INSERT INTO customers (id, name, segment, created_at) VALUES (?, ?, ?, ?)",
+        customer_rows,
+    )
+
+    invoice_rows = []
+    for i in range(1, 1201):
+        invoice_rows.append(
+            (
+                i,
+                rng.randint(1, 400),
+                round(rng.uniform(25.0, 2500.0), 2),
+                rng.choice(["issued", "paid", "overdue"]),
+                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",
+            )
+        )
+
+    orphan_invoice_ids = set(rng.sample(range(1, 1201), 180))
+    rewritten_rows = []
+    valid_invoice_ids = []
+    for row in invoice_rows:
+        invoice_id, customer_id, amount, status, issued_at = row
+        if invoice_id in orphan_invoice_ids:
+            rewritten_rows.append(
+                (invoice_id, 10000 + rng.randint(1, 1000), amount, status, issued_at)
+            )
+        else:
+            rewritten_rows.append(row)
+            valid_invoice_ids.append(invoice_id)
+
+    cursor.executemany(
+        "INSERT INTO invoices (id, customer_id, amount, status, issued_at) VALUES (?, ?, ?, ?, ?)",
+        rewritten_rows,
+    )
+
+    schema_info = []
+    for table_name, row_count in [("customers", 400), ("invoices", 1200)]:
+        cursor.execute(f"PRAGMA table_info({table_name})")
+        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
+        schema_info.append(
+            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
+        )
+
+    sample_rows = {
+        "customers": _extract_sample_rows(conn, "customers", 5),
+        "invoices": _extract_sample_rows(conn, "invoices", 5),
+    }
+
+    observation = SQLObservation(
+        task_id="fk_integrity",
+        task_description=(
+            "A billing pipeline shows referential integrity drift between parent and child records. "
+            "Inspect the schema and samples, then submit a fix_query that removes orphan relationships "
+            "while preserving valid invoice records."
+        ),
+        schema_info=schema_info,
+        sample_rows=sample_rows,
+        attempt_number=1,
+        max_attempts=4,
+        previous_feedback=None,
+    )
+
+    ground_truth = {
+        "task_type": "fk_integrity",
+        "table": "invoices",
+        "column": "customer_id",
+        "parent_table": "customers",
+        "parent_key": "id",
+        "expected_clean_count": len(valid_invoice_ids),
+        "valid_invoice_ids": valid_invoice_ids,
+        "orphan_count": len(orphan_invoice_ids),
+        "preserve_columns": ["id", "customer_id", "amount", "status", "issued_at"],
+    }
+    return conn, observation, ground_truth
+
+
+def create_schema_drift_task(
+    seed: int = 42,
+) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
+    """Create schema drift task with two upstream versions that need canonicalization."""
+    import random
+
+    rng = random.Random(seed)
+    conn = sqlite3.connect(":memory:", check_same_thread=False)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        CREATE TABLE events_v1 (
+            id INTEGER PRIMARY KEY,
+            user_id INTEGER NOT NULL,
+            event_name TEXT NOT NULL,
+            amount_text TEXT NOT NULL,
+            event_time TEXT NOT NULL
+        )
+    """
+    )
+    cursor.execute(
+        """
+        CREATE TABLE events_v2 (
+            event_id INTEGER PRIMARY KEY,
+            account_id INTEGER NOT NULL,
+            event_type TEXT NOT NULL,
+            amount REAL NOT NULL,
+            event_ts TEXT NOT NULL,
+            source TEXT NOT NULL
+        )
+    """
+    )
+
+    v1_rows = []
+    for i in range(1, 601):
+        amount = round(rng.uniform(5.0, 500.0), 2)
+        v1_rows.append(
+            (
+                i,
+                rng.randint(1, 300),
+                rng.choice(["purchase", "refund", "chargeback"]),
+                f"{amount:.2f}",
+                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}T{rng.randint(0,23):02d}:{rng.randint(0,59):02d}:00",
+            )
+        )
+    cursor.executemany(
+        "INSERT INTO events_v1 (id, user_id, event_name, amount_text, event_time) VALUES (?, ?, ?, ?, ?)",
+        v1_rows,
+    )
+
+    v2_rows = []
+    for i in range(601, 1201):
+        v2_rows.append(
+            (
+                i,
+                rng.randint(1, 300),
+                rng.choice(["purchase", "refund", "chargeback"]),
+                round(rng.uniform(5.0, 500.0), 2),
+                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}T{rng.randint(0,23):02d}:{rng.randint(0,59):02d}:00",
+                rng.choice(["batch_a", "batch_b"]),
+            )
+        )
+    cursor.executemany(
+        "INSERT INTO events_v2 (event_id, account_id, event_type, amount, event_ts, source) VALUES (?, ?, ?, ?, ?, ?)",
+        v2_rows,
+    )
+
+    schema_info = []
+    for table_name, row_count in [("events_v1", 600), ("events_v2", 600)]:
+        cursor.execute(f"PRAGMA table_info({table_name})")
+        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
+        schema_info.append(
+            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
+        )
+
+    sample_rows = {
+        "events_v1": _extract_sample_rows(conn, "events_v1", 5),
+        "events_v2": _extract_sample_rows(conn, "events_v2", 5),
+    }
+
+    observation = SQLObservation(
+        task_id="schema_drift",
+        task_description=(
+            "Two upstream event producers emit the same business concept with drifted schemas. "
+            "Create a canonical_events artifact that normalizes both streams into one aligned schema "
+            "with numeric amounts and consistent timestamps."
+        ),
+        schema_info=schema_info,
+        sample_rows=sample_rows,
+        attempt_number=1,
+        max_attempts=4,
+        previous_feedback=None,
+    )
+
+    ground_truth = {
+        "task_type": "schema_drift",
+        "table": "canonical_events",
+        "column": "amount",
+        "source_tables": ["events_v1", "events_v2"],
+        "expected_columns": ["event_id", "user_id", "event_type", "amount", "event_time"],
+        "expected_row_count": 1200,
+        "preserve_columns": ["event_id", "user_id", "event_type", "amount", "event_time"],
+    }
+    return conn, observation, ground_truth