Spaces:

aadi-joshi
/

openenvhack

Sleeping

Fix inference Task4 hint, stale docstring, Dockerfile port comment

c979cbb 7 months ago

16.1 kB

"""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": {}}

	"""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": {}}