Fix server.py, expand tasks to 40, add tests, clean observation schema

Dockerfile

@@ -18,7 +18,8 @@ COPY . .
 EXPOSE 7860
 
 # health check so HF Space knows when it's ready
-HEALTHCHECK --interval=10s --timeout=5s --start-period=10s --retries=3 \
+HEALTHCHECK --interval=10s --timeout=5s --start-period=15s --retries=3 \
   CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:7860/health', timeout=3)"
 
+# server:app refers to the root server.py which re-exports app from sql_env.server
 CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

README.md

@@ -12,7 +12,7 @@ tags:
 # SQL Correction RL Environment
 
 An **OpenEnv-compliant** reinforcement learning environment where an AI agent
-learns to fix broken SQL queries, a real task that developers face every day.
+learns to fix broken SQL queries — a real task that developers face every day.
 
 ---
 
@@ -21,25 +21,28 @@ learns to fix broken SQL queries, a real task that developers face every day.
 SQL errors are one of the most common and costly mistakes in software
 development. This environment trains agents to identify and correct SQL syntax
 and logical errors, ranging from simple typos to complex multi-join query
-reconstruction.
+reconstruction including column name mismatches.
 
 The environment provides **partial progress signals** at every step; the agent
 receives graded feedback even for near-correct answers, enabling meaningful
 learning across the full trajectory rather than sparse end-of-episode rewards.
+A **stagnation penalty** further discourages agents from repeating the same
+wrong answer across steps.
 
 ---
 
 ## Observation Space
 
-| Field              | Type            | Description                                              |
-|--------------------|-----------------|----------------------------------------------------------|
-| `task_id`          | string          | Unique identifier for the current task instance          |
-| `broken_query`     | string          | The malformed SQL query the agent must fix               |
-| `schema_context`   | string or null  | Table and column definitions when a task includes them   |
-| `error_hint`       | string or null  | Plain-language hint about the error (easy tasks only)    |
-| `step_number`      | integer         | Current step within the episode                          |
-| `previous_attempt` | string or null  | The agent's SQL output from the previous step            |
-| `feedback`         | string or null  | Grader feedback on the previous attempt                  |
+| Field              | Type            | Description                                                        |
+|--------------------|-----------------|--------------------------------------------------------------------|
+| `task_id`          | string          | Unique identifier for the current task instance                    |
+| `broken_query`     | string          | The malformed SQL query the agent must fix                         |
+| `schema_context`   | string or null  | Table and column definitions (hard tasks only)                     |
+| `error_hint`       | string or null  | Plain-language hint about the error (easy tasks only)              |
+| `step_number`      | integer         | Current step within the episode (0 = initial)                      |
+| `steps_remaining`  | integer         | Steps left before the episode ends                                 |
+| `previous_attempt` | string or null  | The agent's SQL output from the previous step                      |
+| `feedback`         | string or null  | Grader feedback on the previous attempt                            |
 
 ## Action Space
 
@@ -51,11 +54,11 @@ learning across the full trajectory rather than sparse end-of-episode rewards.
 
 ## Tasks
 
-| Name     | Difficulty | Max Steps | Description |
-|----------|------------|-----------|-------------|
-| `easy`   | Easy       | 5         | Fix a single syntax error (for example `FORM` -> `FROM`). Hint provided. |
-| `medium` | Medium     | 5         | Fix multiple errors including missing keywords and wrong clauses. No hint. |
-| `hard`   | Hard       | 4         | Fix complex multi-join queries with subtle errors and wrong clause ordering. Schema provided, no hint. |
+| Name     | Difficulty | Count | Max Steps | Description |
+|----------|------------|-------|-----------|-------------|
+| `easy`   | Easy       | 15    | 5         | Fix a single keyword typo (e.g. `FORM` → `FROM`). Hint provided. |
+| `medium` | Medium     | 15    | 5         | Fix multiple errors including missing keywords and wrong clauses. No hint. |
+| `hard`   | Hard       | 10    | 4         | Fix complex multi-join queries with subtle errors and wrong column names. Schema provided, no hint. |
 
 ---
 
@@ -65,9 +68,14 @@ learning across the full trajectory rather than sparse end-of-episode rewards.
 |-------|-----------|
 | `1.0` | Exact match after normalization (perfect fix) |
 | `0.7` | All correct tokens present, structure slightly off |
-| `0.4` | Most keywords correct and token overlap is high |
+| `0.4` | Most keywords correct and token overlap is high (≥85% keywords, ≥75% tokens) |
+| `0.3` | Partial keyword and structure match (≥65% keywords, ≥50% tokens) |
 | `0.2` | Basic `SELECT ... FROM ...` structure present |
-| `0.0` | Query still incorrect |
+| `0.0` | Response is not valid SQL |
+
+A **stagnation penalty** of `−0.1` is applied when the agent submits the same
+reward-equivalent answer for two or more consecutive steps, encouraging active
+correction rather than looping.
 
 Episodes terminate when reward = 1.0 (success) or max steps is reached.
 
@@ -88,24 +96,27 @@ uvicorn server:app --host 0.0.0.0 --port 7860
 
 # Test endpoints
 curl -X POST http://localhost:7860/reset \
-  -H "Content-Type: application/json" -d '{"task_name": "easy"}'
+  -H "Content-Type: application/json" -d '{"difficulty": "easy"}'
 
 curl -X POST http://localhost:7860/step \
   -H "Content-Type: application/json" \
-  -d '{"corrected_query": "SELECT * FROM users WHERE id = 1;"}'
+  -d '{"action": {"corrected_query": "SELECT * FROM users WHERE id = 1"}}'
+
+curl http://localhost:7860/tasks
+```
 
-curl -X POST http://localhost:7860/state \
-  -H "Content-Type: application/json" -d '{}'
+### Run Tests
+
+```bash
+pip install pytest
+pytest tests/ -v
 ```
 
 ### Docker
 
 ```bash
 docker build -t sql-correction-env .
-docker run -p 7860:7860 \
-  -e HF_TOKEN=your_token \
-  -e MODEL_NAME=Qwen/Qwen2.5-72B-Instruct \
-  sql-correction-env
+docker run -p 7860:7860 sql-correction-env
 ```
 
 ### Running Inference
@@ -116,10 +127,8 @@ export API_BASE_URL=https://router.huggingface.co/v1
 export MODEL_NAME=Qwen/Qwen2.5-72B-Instruct
 export ENV_URL=http://localhost:7860
 
-# Run each task
-SQL_ENV_TASK=easy   python inference.py
-SQL_ENV_TASK=medium python inference.py
-SQL_ENV_TASK=hard   python inference.py
+# Run all tasks
+python inference.py
 ```
 
 ---
@@ -138,10 +147,10 @@ SQL_ENV_TASK=hard   python inference.py
 
 ## API Endpoints
 
-| Method | Path      | Description                        |
-|--------|-----------|------------------------------------|
+| Method | Path      | Description                            |
+|--------|-----------|----------------------------------------|
 | POST   | `/reset`  | Start new episode, returns observation |
-| POST   | `/step`   | Submit action, returns result      |
-| POST   | `/state`  | Get current episode state          |
-| GET    | `/health` | Health check                       |
-# rebuilt
+| POST   | `/step`   | Submit action, returns result          |
+| POST   | `/state`  | Get current episode state              |
+| GET    | `/health` | Health check                           |
+| GET    | `/tasks`  | List available task difficulties       |

inference.py

@@ -26,12 +26,12 @@ except Exception:
     OpenAI = None
 
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "dummy")
-TASK_NAME = os.getenv("SQL_ENV_TASK", "easy")
-BENCHMARK = "sql-correction-env"
-ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
-MAX_STEPS = 8
+MODEL_NAME   = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "dummy")
+TASK_NAME    = os.getenv("SQL_ENV_TASK", "easy")
+BENCHMARK    = "sql-correction-env"
+ENV_URL      = os.getenv("ENV_URL", "http://localhost:7860")
+MAX_STEPS    = 8
 SUCCESS_SCORE_THRESHOLD = 0.5
 
 
@@ -52,7 +52,6 @@ def log_step(
 ) -> None:
     err = error if error else "null"
     done_val = str(done).lower()
-    # Collapse newlines so the entire step fits on one line (spec requirement)
     action_clean = action.replace("\n", " ").replace("\r", "").strip()
     print(
         f"[STEP] step={step} action={action_clean} "
@@ -80,26 +79,28 @@ SYSTEM_PROMPT = textwrap.dedent(
     You will be shown a broken SQL query that contains typos or keyword errors.
     Fix ALL errors and return ONLY the corrected SQL query.
     No explanation, no markdown, no code blocks, no backticks.
-    Common errors: FORM->FROM, WEHRE->WHERE, GRUP->GROUP, HAVNG->HAVING,
-    ORDR->ORDER, INNE->INNER, LFT->LEFT, BETWEN->BETWEEN, DSC->DESC, SELCT->SELECT.
+    Common keyword typos to watch for:
+      FORM->FROM, WEHRE->WHERE, WHER->WHERE,
+      GRUP->GROUP, HAVNG->HAVING, ORDR->ORDER,
+      INNE->INNER, LFT->LEFT, BETWEN->BETWEEN,
+      DSC->DESC, SELCT->SELECT, LIMT->LIMIT.
+    Also watch for column name errors described in the schema context.
     """
 ).strip()
 
-
 SQL_REPLACEMENTS = {
-    "FORM": "FROM",
-    "WEHRE": "WHERE",
-    "WHER": "WHERE",
-    "GRUP": "GROUP",
-    "HAVNG": "HAVING",
-    "ORDR": "ORDER",
-    "INNE": "INNER",
-    "LFT": "LEFT",
+    "FORM":   "FROM",
+    "WEHRE":  "WHERE",
+    "WHER":   "WHERE",
+    "GRUP":   "GROUP",
+    "HAVNG":  "HAVING",
+    "ORDR":   "ORDER",
+    "INNE":   "INNER",
+    "LFT":    "LEFT",
     "BETWEN": "BETWEEN",
-    "DSC": "DESC",
-    "SELCT": "SELECT",
-    "LIMT": "LIMIT",
-    "DPT_ID": "DEPT_ID",
+    "DSC":    "DESC",
+    "SELCT":  "SELECT",
+    "LIMT":   "LIMIT",
 }
 
 
@@ -114,7 +115,9 @@ def heuristic_correct_sql(query: str) -> str:
 
 
 def get_model_action(
-    client: Optional["OpenAI"], obs: dict, history: List[str]
+    client: Optional["OpenAI"],
+    obs: dict,
+    history: List[str],
 ) -> str:
     """Return a corrected SQL string. Falls back to heuristic on any failure."""
     heuristic = heuristic_correct_sql(obs.get("broken_query", ""))
@@ -128,10 +131,11 @@ def get_model_action(
         Broken SQL query:
         {obs.get("broken_query", "")}
 
-        Schema context: {obs.get("schema_context") or "Not provided"}
-        Error hint: {obs.get("error_hint") or "None"}
-        Your previous attempt: {obs.get("previous_attempt") or "None"}
-        Feedback: {obs.get("feedback") or "None"}
+        Schema context:    {obs.get("schema_context") or "Not provided"}
+        Error hint:        {obs.get("error_hint") or "None"}
+        Steps remaining:   {obs.get("steps_remaining", "?")}
+        Previous attempt:  {obs.get("previous_attempt") or "None"}
+        Feedback:          {obs.get("feedback") or "None"}
 
         Recent history:
         {history_block}
@@ -145,7 +149,7 @@ def get_model_action(
             model=MODEL_NAME,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
+                {"role": "user",   "content": user_prompt},
             ],
             temperature=0.2,
             max_tokens=300,
@@ -167,17 +171,14 @@ async def run_task(task_name: str) -> None:
     Run one full episode for `task_name`.
 
     The [END] log line is ALWAYS emitted via the finally block, even if an
-    exception occurs mid-episode or the reset call fails. This is required
-    by the hackathon spec to avoid disqualification.
+    exception occurs mid-episode or the reset call fails.
     """
-    # Initialise all accumulators BEFORE the try so finally can always read them
     rewards: List[float] = []
-    history: List[str] = []
+    history: List[str]   = []
     steps_taken = 0
-    score = 0.0
-    success = False
+    score       = 0.0
+    success     = False
 
-    # Build LLM client (best-effort; None means heuristic-only mode)
     client = None
     if OpenAI is not None and API_KEY not in {"", "dummy"}:
         try:
@@ -191,47 +192,52 @@ async def run_task(task_name: str) -> None:
     try:
         http = httpx.AsyncClient(base_url=ENV_URL, timeout=60.0)
 
-        # --- reset ---------------------------------------------------------
+        # ── reset ───────────────────────────────────────────���──────────────
         reset_failed = False
         obs: dict = {}
         try:
-            reset_resp = await http.post("/reset", json={"difficulty": task_name})
+            reset_resp = await http.post(
+                "/reset", json={"difficulty": task_name}
+            )
             reset_resp.raise_for_status()
             reset_data = reset_resp.json()
+            # The openenv wrapper may nest the observation under "observation"
             obs = reset_data.get("observation", reset_data)
         except Exception as exc:
             print(f"[DEBUG] Reset failed: {exc}", flush=True)
-            # Do NOT return here — fall through to finally so [END] is always logged
             reset_failed = True
 
         if not reset_failed:
-            # --- step loop -------------------------------------------------
+            # ── step loop ──────────────────────────────────────────────────
             for step in range(1, MAX_STEPS + 1):
-                # Get action (never raises — heuristic is the ultimate fallback)
                 try:
                     action_str = get_model_action(client, obs, history)
                 except Exception as exc:
                     print(f"[DEBUG] Model action failed: {exc}", flush=True)
-                    action_str = heuristic_correct_sql(obs.get("broken_query", ""))
+                    action_str = heuristic_correct_sql(
+                        obs.get("broken_query", "")
+                    )
 
-                # Submit action to environment
                 try:
-                    step_resp = await http.post("/step", json={"action":{"corrected_query": action_str}})
+                    # Action must be wrapped under {"action": {...}}
+                    step_resp = await http.post(
+                        "/step",
+                        json={"action": {"corrected_query": action_str}},
+                    )
                     step_resp.raise_for_status()
                     result = step_resp.json()
                 except Exception as exc:
                     print(f"[DEBUG] Step {step} request failed: {exc}", flush=True)
-                    # Treat as a 0-reward terminal step so episode ends cleanly
                     rewards.append(0.0)
                     steps_taken = step
                     log_step(step, action_str, 0.0, True, str(exc))
                     break
 
-                obs = result.get("observation", obs)
+                obs    = result.get("observation", obs)
                 reward = float(result.get("reward", 0.0))
-                done = bool(result.get("done", False))
-                info = result.get("info")
-                error = info.get("error") if isinstance(info, dict) else None
+                done   = bool(result.get("done", False))
+                info   = result.get("info")
+                error  = info.get("error") if isinstance(info, dict) else None
 
                 rewards.append(reward)
                 steps_taken = step
@@ -244,23 +250,19 @@ async def run_task(task_name: str) -> None:
                 if done:
                     break
 
-            # Score = average reward across all steps, clamped to [0, 1]
             if rewards:
-                score = min(max(sum(rewards) / len(rewards), 0.0), 1.0)
+                score   = min(max(sum(rewards) / len(rewards), 0.0), 1.0)
             success = score >= SUCCESS_SCORE_THRESHOLD
 
     except Exception as exc:
-        # Catch-all for any unexpected error in the episode body
         print(f"[DEBUG] Unhandled episode error: {exc}", flush=True)
 
     finally:
-        # Always close the HTTP client
         if http is not None:
             try:
                 await http.aclose()
             except Exception as exc:
                 print(f"[DEBUG] HTTP close error: {exc}", flush=True)
-        # [END] MUST always be emitted — even after reset failure or exception
         log_end(success, steps_taken, score, rewards)
 
 
@@ -269,13 +271,8 @@ async def run_task(task_name: str) -> None:
 # ---------------------------------------------------------------------------
 
 async def main() -> None:
-    """
-    Run tasks according to SQL_ENV_TASK.
-    If SQL_ENV_TASK is a single valid difficulty, run only that task.
-    Otherwise run all three in sequence so all 3 tasks produce scores.
-    """
+    """Run all three difficulties in sequence so validator sees 3 [END] lines."""
     try:
-        # Always run all 3 tasks — validator counts 3 [END] lines
         for difficulty in ("easy", "medium", "hard"):
             await run_task(difficulty)
             print("", flush=True)

openenv.yaml

@@ -13,30 +13,36 @@ spec_version: 1
 name: sql-correction-env
 type: sequential
 runtime: docker
-app: server.app:app
+app: server:app
 port: 7860
 
 description: >
   An OpenEnv RL environment where an AI agent fixes broken SQL queries.
   Simulates a real developer task: identifying and correcting SQL syntax
   and logical errors across easy, medium, and hard difficulty levels.
+  Includes 40 tasks total with partial-credit grading and a stagnation penalty.
 
 author: SyncShift
 
 tasks:
   - id: easy
-    description: Fix a single syntax error. Error hint provided.
+    description: Fix a single keyword typo. Error hint provided.
     steps: 5
+    count: 15
     ideal_action: correct_sql
 
   - id: medium
     description: Fix multiple errors across keywords and clauses. No hint.
     steps: 5
+    count: 15
     ideal_action: correct_sql
 
   - id: hard
-    description: Fix many errors in complex multi-join queries. Schema provided.
+    description: >
+      Fix many errors in complex multi-join queries including column name
+      mismatches. Schema provided, no hint.
     steps: 4
+    count: 10
     ideal_action: correct_sql
 
 observation_space:
@@ -54,6 +60,8 @@ observation_space:
       nullable: true
     step_number:
       type: integer
+    steps_remaining:
+      type: integer
     previous_attempt:
       type: string
       nullable: true
@@ -68,17 +76,19 @@ action_space:
       type: string
 
 reward:
-  range: [0.001, 0.999]
+  range: [0.0, 1.0]
   description: >
-    0.999 = exact match, 0.7 = right tokens minor structure diff,
-    0.4 = most keywords correct, 0.2 = basic structure present, 0.001 = invalid SQL.
+    1.0 = exact match, 0.7 = right tokens minor structure diff,
+    0.4 = most keywords correct, 0.3 = partial match,
+    0.2 = basic structure present, 0.0 = invalid SQL.
+    Stagnation penalty of -0.1 applied after 2+ identical-reward steps.
 
 scoring:
-  reward_range: [0.001, 0.999]
+  reward_range: [0.0, 1.0]
   success_threshold: 0.5
   score_formula: mean(step_rewards)
 
 constraints:
   max_runtime_seconds: 1200
   max_memory_gb: 8
-  max_vcpu: 2
+  max_vcpu: 2

pyproject.toml

@@ -4,14 +4,17 @@ version = "1.0.0"
 description = "OpenEnv SQL Query Correction RL Environment"
 requires-python = ">=3.11"
 dependencies = [
-    "fastapi",
-    "uvicorn",
-    "pydantic",
-    "httpx",
-    "openai",
-    "openenv-core",
+    "fastapi>=0.135.0",
+    "uvicorn>=0.43.0",
+    "pydantic>=2.12.0",
+    "httpx>=0.28.1",
+    "openai>=2.7.2",
+    "openenv-core==0.2.3",
 ]
 
+[project.optional-dependencies]
+dev = ["pytest>=8.0.0"]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
@@ -21,3 +24,9 @@ packages = ["sql_env"]
 
 [project.scripts]
 server = "sql_env.server:main"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]

requirements.txt

@@ -1,6 +1,7 @@
-fastapi
-uvicorn
-pydantic
+fastapi>=0.135.0
+uvicorn>=0.43.0
+pydantic>=2.12.0
 httpx>=0.28.1
 openai>=2.7.2
-openenv-core==0.2.3
+openenv-core==0.2.3
+pytest>=8.0.0

server.log

@@ -1,2 +0,0 @@
-nohup: ignoring input
-ERROR:    Error loading ASGI app. Could not import module "app".

sql_env/__init__.py

@@ -11,4 +11,4 @@ __all__ = [
     "StepResult",
     "ALL_TASKS",
     "TASK_SETS",
-]
+]

sql_env/env.py

@@ -13,7 +13,13 @@ class SQLCorrectionEnv:
 
     The agent receives a broken SQL query and must return the corrected version.
     Reward is shaped across the full trajectory — partial credit is given for
-    incremental improvements, penalizing stagnation and infinite loops.
+    incremental improvements, penalizing stagnation.
+
+    Usage::
+
+        env = SQLCorrectionEnv(difficulty="easy")
+        obs = await env.reset()
+        result = await env.step(SQLAction(corrected_query="SELECT * FROM users"))
     """
 
     def __init__(self, difficulty: str = "easy", task_index: Optional[int] = None):
@@ -30,7 +36,7 @@ class SQLCorrectionEnv:
         self._last_reward: float = 0.0
         self._stagnation_count: int = 0
 
-    # ── OpenEnv Interface ─────────────────────────────────────
+    # ── OpenEnv Interface ─────────────────────────────────────────────────────
 
     async def reset(self) -> SQLObservation:
         """Reset the environment and return the initial observation."""
@@ -61,27 +67,24 @@ class SQLCorrectionEnv:
 
         self._step_count += 1
 
-        # grade the action
         reward_model = grade(action, self._task)
         reward = reward_model.value
 
-        # detect stagnation (same reward twice in a row) — penalize
+        # Stagnation penalty
         if abs(reward - self._last_reward) < 0.01 and self._step_count > 1:
             self._stagnation_count += 1
             if self._stagnation_count >= 2:
-                reward = max(0.001, reward - 0.1)  # stagnation penalty
+                reward = max(0.0, reward - 0.1)
         else:
             self._stagnation_count = 0
 
         self._last_reward = reward
 
-        # generate feedback for the next observation
         feedback = generate_feedback(action, self._task, reward_model)
         self._last_feedback = feedback
         self._previous_attempt = action.corrected_query
 
-        # episode ends on perfect score or max steps reached
-        done = reward_model.value >= 0.99 or self._step_count >= self._task.max_steps
+        done = reward_model.value >= 0.95 or self._step_count >= self._task.max_steps
         self._done = done
 
         obs = self._make_observation()
@@ -95,7 +98,7 @@ class SQLCorrectionEnv:
                 "step": self._step_count,
                 "max_steps": self._task.max_steps,
                 "task_id": self._task.task_id,
-            }
+            },
         )
 
     async def state(self) -> dict:
@@ -117,16 +120,18 @@ class SQLCorrectionEnv:
         self._task = None
         self._done = True
 
-    # ── Internal ──────────────────────────────────────────────
+    # ── Internal ──────────────────────────────────────────────────────────────
 
     def _make_observation(self) -> SQLObservation:
         assert self._task is not None
+        steps_remaining = max(0, self._task.max_steps - self._step_count)
         return SQLObservation(
             task_id=self._task.task_id,
             broken_query=self._task.broken_query,
             schema_context=self._task.schema_context,
             error_hint=self._task.error_hint if self.difficulty == "easy" else None,
             step_number=self._step_count,
+            steps_remaining=steps_remaining,
             previous_attempt=self._previous_attempt,
             feedback=self._last_feedback,
         )

sql_env/grader.py

@@ -2,11 +2,6 @@ import re
 from sql_env.models import SQLAction, SQLTask, SQLReward
 
 
-def _clamp(value: float) -> float:
-    """Ensure score is strictly between 0 and 1 (gt=0.0, lt=1.0)."""
-    return max(0.001, min(0.999, value))
-
-
 def _normalize(query: str) -> str:
     """Uppercase, collapse whitespace, strip trailing semicolons."""
     q = query.strip().upper()
@@ -27,7 +22,7 @@ def _sql_keywords_present(query: str) -> set:
         'ORDER', 'JOIN', 'INNER', 'LEFT', 'RIGHT', 'OUTER',
         'BETWEEN', 'DESC', 'ASC', 'LIMIT', 'COUNT', 'SUM',
         'AVG', 'MAX', 'MIN', 'AS', 'ON', 'AND', 'OR', 'NOT',
-        'IN', 'LIKE', 'IS', 'NULL', 'DISTINCT'
+        'IN', 'LIKE', 'IS', 'NULL', 'DISTINCT',
     }
     normed = _normalize(query)
     found = set()
@@ -39,34 +34,32 @@ def _sql_keywords_present(query: str) -> set:
 
 def grade(action: SQLAction, task: SQLTask) -> SQLReward:
     """
-    4-level grader with partial progress signals.
-    0.99 — exact normalized match
-    0.7  — same tokens, minor whitespace/alias differences
-    0.4  — key SQL keywords all present and correct table/column names
+    5-level grader with partial progress signals.
+
+    1.0  — exact normalized match (perfect fix)
+    0.7  — same token set, minor structural/whitespace differences
+    0.4  — most SQL keywords correct AND high token overlap
+    0.3  — partial keyword and structure match
     0.2  — basic SELECT/FROM structure present
-    0.01 — completely wrong
-    All scores are clamped to be strictly between 0.0 and 1.0.
+    0.0  — not recognizable SQL
     """
     agent = _normalize(action.corrected_query)
     correct = _normalize(task.canonical_answer)
 
-    # ── Level 1: Exact match ─────────────────────────────────
+    # ── Level 1: Exact match ─────────────────────────────────────────────────
     if agent == correct:
-        return SQLReward(
-            value=_clamp(0.999),
-            reason="Exact match — perfect correction."
-        )
+        return SQLReward(value=1.0, reason="Exact match — perfect correction.")
 
-    # ── Level 2: Same token set (right words, minor ordering) ─
+    # ── Level 2: Same token set (right words, minor ordering/alias diff) ─────
     agent_tokens = _tokenize(action.corrected_query)
     correct_tokens = _tokenize(task.canonical_answer)
     if agent_tokens == correct_tokens:
         return SQLReward(
-            value=_clamp(0.7),
-            reason="All correct tokens present but structure differs slightly."
+            value=0.7,
+            reason="All correct tokens present but structure differs slightly.",
         )
 
-    # ── Level 3: Most keywords correct + high token overlap ───
+    # ── Level 3: Most keywords correct + high token overlap ──────────────────
     correct_kws = _sql_keywords_present(task.canonical_answer)
     agent_kws = _sql_keywords_present(action.corrected_query)
     kw_overlap = len(correct_kws & agent_kws) / max(len(correct_kws), 1)
@@ -74,38 +67,49 @@ def grade(action: SQLAction, task: SQLTask) -> SQLReward:
 
     if kw_overlap >= 0.85 and token_overlap >= 0.75:
         return SQLReward(
-            value=_clamp(0.4),
-            reason=f"Most keywords correct ({kw_overlap:.0%} keyword match, {token_overlap:.0%} token match)."
+            value=0.4,
+            reason=(
+                f"Most keywords correct "
+                f"({kw_overlap:.0%} keyword match, {token_overlap:.0%} token match)."
+            ),
         )
 
-    # ── Level 3.5: Partial keyword and structure match ────────
-    if kw_overlap >= 0.7 and token_overlap >= 0.55:
+    # ── Level 3.5: Partial keyword and structure match ────────────────────────
+    if kw_overlap >= 0.65 and token_overlap >= 0.50:
         return SQLReward(
-            value=_clamp(0.3),
-            reason=f"Partial keyword and structure match ({kw_overlap:.0%} keyword match, {token_overlap:.0%} token match)."
+            value=0.3,
+            reason=(
+                f"Partial keyword and structure match "
+                f"({kw_overlap:.0%} keyword match, {token_overlap:.0%} token match)."
+            ),
         )
 
-    # ── Level 4: Basic structure present ─────────────────────
+    # ── Level 4: Basic structure present ─────────────────────────────────────
     if 'SELECT' in agent and 'FROM' in agent:
         return SQLReward(
-            value=_clamp(0.2),
-            reason="Basic SELECT/FROM structure present but significant errors remain."
+            value=0.2,
+            reason="Basic SELECT/FROM structure present but significant errors remain.",
         )
 
-    # ── Level 0: No recognizable SQL ─────────────────────────
-    return SQLReward(value=_clamp(0.001), reason="Response is not valid SQL.")
+    # ── Level 0: No recognizable SQL ─────────────────────────────────────────
+    return SQLReward(value=0.0, reason="Response is not valid SQL.")
 
 
 def generate_feedback(action: SQLAction, task: SQLTask, reward: SQLReward) -> str:
-    """Human-readable feedback shown in next observation."""
-    if reward.value >= 0.99:
+    """Human-readable feedback shown in the next observation."""
+    if reward.value >= 1.0:
         return "Correct! Query matches perfectly."
     if reward.value >= 0.7:
         return "Very close — check spacing or minor clause differences."
     if reward.value >= 0.4:
-        return "Good progress — most keywords are right, but check for typos in keywords or column names."
+        return (
+            "Good progress — most keywords are right, "
+            "but check for typos in keywords or column names."
+        )
     if reward.value >= 0.3:
         return "Partial match — right direction but several keywords or columns are off."
     if reward.value >= 0.2:
-        return "Basic structure is there — look carefully at every SQL keyword for typos."
-    return "The response doesn't look like valid SQL. Start with SELECT ... FROM ..."
+        return (
+            "Basic structure is there — look carefully at every SQL keyword for typos."
+        )
+    return "The response doesn't look like valid SQL. Start with SELECT ... FROM ..."

sql_env/models.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Any, Callable
+from typing import List, Optional, Any
 from openenv.core.env_server.types import Action, Observation, State
 from pydantic import Field, BaseModel
 
@@ -8,15 +8,27 @@ class SQLAction(Action):
 
 
 class SQLObservation(Observation):
+    """
+    Observation returned to the agent each step.
+
+    Fields:
+        task_id:          Unique identifier for the current task instance.
+        broken_query:     The malformed SQL query the agent must fix.
+        schema_context:   Table/column definitions (hard tasks only).
+        error_hint:       Plain-language hint about the error (easy tasks only).
+        step_number:      Current step within the episode (0 = initial observation).
+        steps_remaining:  How many steps are left before the episode ends.
+        previous_attempt: The agent's SQL output from the previous step.
+        feedback:         Grader feedback on the previous attempt.
+    """
     task_id: str
     broken_query: str
     schema_context: Optional[str] = None
     error_hint: Optional[str] = None
     step_number: int
+    steps_remaining: Optional[int] = None
     previous_attempt: Optional[str] = None
     feedback: Optional[str] = None
-    reward: float = 0.001
-    done: bool = False
 
 
 class SQLState(State):
@@ -30,7 +42,8 @@ class SQLState(State):
 
 
 class SQLReward(BaseModel):
-    value: float = Field(gt=0.0, lt=1.0)
+    # Allow full [0.0, 1.0] range so perfect matches can return exactly 1.0
+    value: float = Field(ge=0.0, le=1.0)
     reason: str
 
 
@@ -51,4 +64,4 @@ class StepResult(BaseModel):
     observation: SQLObservation
     reward: float
     done: bool
-    info: dict = Field(default_factory=dict)
+    info: dict = Field(default_factory=dict)

sql_env/server.py

@@ -2,10 +2,8 @@
 FastAPI server using openenv.core base classes — required for validator.
 """
 import random
-from typing import Optional
 from openenv.core.env_server.http_server import create_app
 from openenv.core.env_server.interfaces import Environment
-from openenv.core.env_server.types import State
 
 try:
     from sql_env.models import SQLAction, SQLObservation, SQLState
@@ -27,12 +25,13 @@ class SQLCorrectionEnvironment(Environment):
         self._done = False
         self._last_reward = 0.0
         self._rewards_history = []
+        self._stagnation_count = 0
 
     def reset(self, seed=None, episode_id=None, **kwargs) -> SQLObservation:
         actual_difficulty = (
-            kwargs.get("task_id") or
-            kwargs.get("difficulty") or
-            "easy"
+            kwargs.get("task_id")
+            or kwargs.get("difficulty")
+            or "easy"
         )
         self._difficulty = actual_difficulty
         tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])
@@ -41,39 +40,37 @@ class SQLCorrectionEnvironment(Environment):
         self._done = False
         self._last_reward = 0.0
         self._rewards_history = []
-        return SQLObservation(
-            task_id=self._current_task.task_id,
-            broken_query=self._current_task.broken_query,
-            schema_context=self._current_task.schema_context,
-            error_hint=self._current_task.error_hint,
-            step_number=0,
-            previous_attempt=None,
-            feedback=None,
-            reward=0.001,
-            done=False,
-        )
+        self._stagnation_count = 0
+        return self._make_observation(previous_attempt=None, feedback=None)
 
     def step(self, action: SQLAction) -> SQLObservation:
         if self._current_task is None:
             self.reset()
+
         self._step_count += 1
         reward_obj = grade(action, self._current_task)
         reward = reward_obj.value
+
+        # Stagnation penalty: penalize repeating the same score
+        if abs(reward - self._last_reward) < 0.01 and self._step_count > 1:
+            self._stagnation_count += 1
+            if self._stagnation_count >= 2:
+                reward = max(0.0, reward - 0.1)
+        else:
+            self._stagnation_count = 0
+
         self._last_reward = reward
         self._rewards_history.append(reward)
-        done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
+
+        done = (reward_obj.value >= 0.95) or (
+            self._step_count >= self._current_task.max_steps
+        )
         self._done = done
         feedback = generate_feedback(action, self._current_task, reward_obj)
-        return SQLObservation(
-            task_id=self._current_task.task_id,
-            broken_query=self._current_task.broken_query,
-            schema_context=self._current_task.schema_context,
-            error_hint=self._current_task.error_hint,
-            step_number=self._step_count,
+
+        return self._make_observation(
             previous_attempt=action.corrected_query,
             feedback=feedback,
-            reward=reward,
-            done=done,
         )
 
     @property
@@ -97,96 +94,33 @@ class SQLCorrectionEnvironment(Environment):
             last_reward=self._last_reward,
             rewards_history=self._rewards_history,
         )
-    def step(self, action: SQLAction) -> SQLObservation:
-        if self._current_task is None:
-            self.reset()
-        self._step_count += 1
-        reward_obj = grade(action, self._current_task)
-        reward = reward_obj.value
-        self._last_reward = reward
-        self._rewards_history.append(reward)
-        done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
-        self._done = done
-        feedback = generate_feedback(action, self._current_task, reward_obj)
+
+    # ── Internal helpers ──────────────────────────────────────────────────────
+
+    def _make_observation(
+        self,
+        previous_attempt: str | None,
+        feedback: str | None,
+    ) -> SQLObservation:
+        assert self._current_task is not None
+        max_steps = self._current_task.max_steps
+        steps_remaining = max(0, max_steps - self._step_count)
         return SQLObservation(
             task_id=self._current_task.task_id,
             broken_query=self._current_task.broken_query,
             schema_context=self._current_task.schema_context,
-            error_hint=self._current_task.error_hint,
+            # Only surface the hint on easy tasks
+            error_hint=(
+                self._current_task.error_hint
+                if self._difficulty == "easy"
+                else None
+            ),
             step_number=self._step_count,
-            previous_attempt=action.corrected_query,
+            steps_remaining=steps_remaining,
+            previous_attempt=previous_attempt,
             feedback=feedback,
-            reward=reward,
-            done=done,
         )
 
-    def state(self) -> SQLState:
-        if self._current_task is None:
-            return SQLState(
-                task_id="none",
-                difficulty="none",
-                step_count=0,
-                max_steps=0,
-                done=False,
-                last_reward=0.0,
-                rewards_history=[],
-            )
-        return SQLState(
-            task_id=self._current_task.task_id,
-            difficulty=self._difficulty,
-            step_count=self._step_count,
-            max_steps=self._current_task.max_steps,
-            done=self._done,
-            last_reward=self._last_reward,
-            rewards_history=self._rewards_history,
-        )
-    
-def step(self, action: SQLAction) -> SQLObservation:
-    # Auto-reset if no task loaded (create_app may use fresh instances)
-    if self._current_task is None:
-        self.reset()
-    
-    self._step_count += 1
-    reward_obj = grade(action, self._current_task)
-    reward = reward_obj.value
-    self._last_reward = reward
-    self._rewards_history.append(reward)
-    done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
-    self._done = done
-    feedback = generate_feedback(action, self._current_task, reward_obj)
-    return SQLObservation(
-        task_id=self._current_task.task_id,
-        broken_query=self._current_task.broken_query,
-        schema_context=self._current_task.schema_context,
-        error_hint=self._current_task.error_hint,
-        step_number=self._step_count,
-        previous_attempt=action.corrected_query,
-        feedback=feedback,
-        reward=reward,
-        done=done,
-    )
-
-    @property
-    def state(self) -> SQLState:
-        if self._current_task is None:
-            return SQLState(
-                task_id="none",
-                difficulty="none",
-                step_count=0,
-                max_steps=0,
-                done=False,
-                last_reward=0.001,
-                rewards_history=[],
-            )
-        return SQLState(
-            task_id=self._current_task.task_id,
-            difficulty=self._difficulty,
-            step_count=self._step_count,
-            max_steps=self._current_task.max_steps,
-            done=self._done,
-            last_reward=self._last_reward,
-            rewards_history=self._rewards_history,
-        )
 
 app = create_app(
     SQLCorrectionEnvironment,
@@ -196,6 +130,41 @@ app = create_app(
 )
 
 
+@app.get("/tasks")
+async def list_tasks():
+    """List available task difficulties with metadata."""
+    return {
+        "tasks": [
+            {
+                "id": "easy",
+                "difficulty": "easy",
+                "description": "Fix a single SQL keyword typo. Error hint provided.",
+                "max_steps": 5,
+                "count": 15,
+            },
+            {
+                "id": "medium",
+                "difficulty": "medium",
+                "description": (
+                    "Fix multiple errors across keywords and clauses. No hint."
+                ),
+                "max_steps": 5,
+                "count": 15,
+            },
+            {
+                "id": "hard",
+                "difficulty": "hard",
+                "description": (
+                    "Fix complex multi-join queries including column name errors. "
+                    "Schema provided, no hint."
+                ),
+                "max_steps": 4,
+                "count": 10,
+            },
+        ]
+    }
+
+
 def main():
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)
@@ -203,14 +172,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-from fastapi import Request
-
-@app.get("/tasks")
-async def list_tasks():
-    return {
-        "tasks": [
-            {"id": "easy", "difficulty": "easy", "description": "Fix a single syntax error.", "steps": 5, "ideal_action": "correct_sql", "has_grader": True, "grader": "sql_env.grader.grade"},
-            {"id": "medium", "difficulty": "medium", "description": "Fix multiple errors.", "steps": 5, "ideal_action": "correct_sql", "has_grader": True, "grader": "sql_env.grader.grade"},
-            {"id": "hard", "difficulty": "hard", "description": "Fix complex multi-join queries.", "steps": 4, "ideal_action": "correct_sql", "has_grader": True, "grader": "sql_env.grader.grade"},
-        ]
-    }

sql_env/tasks/__init__.py

@@ -8,4 +8,7 @@ ALL_TASKS = {
     "hard": HARD_TASKS,
 }
 
-TASK_SETS = ALL_TASKS
+# Alias used by server.py
+TASK_SETS = ALL_TASKS
+
+__all__ = ["ALL_TASKS", "TASK_SETS", "EASY_TASKS", "MEDIUM_TASKS", "HARD_TASKS"]

sql_env/tasks/easy.py

@@ -42,4 +42,84 @@ EASY_TASKS = [
         error_hint="There is a typo in a SQL keyword near the table name.",
         grader=grade,
     ),
-]
+    SQLTask(
+        task_id="easy_006",
+        difficulty="easy",
+        broken_query="SELECT id, email FORM users WHERE active = 1",
+        canonical_answer="SELECT id, email FROM users WHERE active = 1",
+        error_hint="There is a typo in a SQL keyword near the table name.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_007",
+        difficulty="easy",
+        broken_query="SELECT * FROM invoices WEHRE amount > 500",
+        canonical_answer="SELECT * FROM invoices WHERE amount > 500",
+        error_hint="There is a typo in the filtering keyword.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_008",
+        difficulty="easy",
+        broken_query="SELCT title, author FROM books",
+        canonical_answer="SELECT title, author FROM books",
+        error_hint="There is a typo in the first keyword of the query.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_009",
+        difficulty="easy",
+        broken_query="SELECT name FORM departments WHERE location = 'NYC'",
+        canonical_answer="SELECT name FROM departments WHERE location = 'NYC'",
+        error_hint="There is a typo in a SQL keyword near the table name.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_010",
+        difficulty="easy",
+        broken_query="SELECT product_name, price FROM catalog WEHRE category = 'electronics'",
+        canonical_answer="SELECT product_name, price FROM catalog WHERE category = 'electronics'",
+        error_hint="There is a typo in the filtering keyword.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_011",
+        difficulty="easy",
+        broken_query="SELECT * FORM sessions WHERE user_id = 42",
+        canonical_answer="SELECT * FROM sessions WHERE user_id = 42",
+        error_hint="There is a typo in a SQL keyword near the table name.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_012",
+        difficulty="easy",
+        broken_query="SELCT COUNT(*) FROM transactions WHERE type = 'credit'",
+        canonical_answer="SELECT COUNT(*) FROM transactions WHERE type = 'credit'",
+        error_hint="There is a typo in the first keyword of the query.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_013",
+        difficulty="easy",
+        broken_query="SELECT username, created_at FORM accounts WHERE verified = 1",
+        canonical_answer="SELECT username, created_at FROM accounts WHERE verified = 1",
+        error_hint="There is a typo in a SQL keyword near the table name.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_014",
+        difficulty="easy",
+        broken_query="SELECT * FROM reports WEHRE year = 2024",
+        canonical_answer="SELECT * FROM reports WHERE year = 2024",
+        error_hint="There is a typo in the filtering keyword.",
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="easy_015",
+        difficulty="easy",
+        broken_query="SELCT AVG(score) FROM results WHERE subject = 'math'",
+        canonical_answer="SELECT AVG(score) FROM results WHERE subject = 'math'",
+        error_hint="There is a typo in the first keyword of the query.",
+        grader=grade,
+    ),
+]

sql_env/tasks/hard.py

@@ -5,8 +5,24 @@ HARD_TASKS = [
     SQLTask(
         task_id="hard_001",
         difficulty="hard",
-        broken_query="SELCT e.name, d.dept_name, SUM(s.amount) AS total_sales FORM employees e LFT JOIN departments d ON e.dpt_id = d.id INNE JOIN sales s ON e.id = s.emp_id WHER s.sale_date BETWEN '2024-01-01' AND '2024-12-31' GRUP BY e.name, d.dept_name HAVNG SUM(s.amount) > 10000 ORDR BY total_sales DSC",
-        canonical_answer="SELECT e.name, d.dept_name, SUM(s.amount) AS total_sales FROM employees e LEFT JOIN departments d ON e.dept_id = d.id INNER JOIN sales s ON e.id = s.emp_id WHERE s.sale_date BETWEEN '2024-01-01' AND '2024-12-31' GROUP BY e.name, d.dept_name HAVING SUM(s.amount) > 10000 ORDER BY total_sales DESC",
+        broken_query=(
+            "SELCT e.name, d.dept_name, SUM(s.amount) AS total_sales "
+            "FORM employees e LFT JOIN departments d ON e.dpt_id = d.id "
+            "INNE JOIN sales s ON e.id = s.emp_id "
+            "WHER s.sale_date BETWEN '2024-01-01' AND '2024-12-31' "
+            "GRUP BY e.name, d.dept_name "
+            "HAVNG SUM(s.amount) > 10000 "
+            "ORDR BY total_sales DSC"
+        ),
+        canonical_answer=(
+            "SELECT e.name, d.dept_name, SUM(s.amount) AS total_sales "
+            "FROM employees e LEFT JOIN departments d ON e.dept_id = d.id "
+            "INNER JOIN sales s ON e.id = s.emp_id "
+            "WHERE s.sale_date BETWEEN '2024-01-01' AND '2024-12-31' "
+            "GROUP BY e.name, d.dept_name "
+            "HAVING SUM(s.amount) > 10000 "
+            "ORDER BY total_sales DESC"
+        ),
         schema_context=(
             "employees(id INT, name VARCHAR, dept_id INT, salary DECIMAL)\n"
             "departments(id INT, dept_name VARCHAR)\n"
@@ -19,8 +35,26 @@ HARD_TASKS = [
     SQLTask(
         task_id="hard_002",
         difficulty="hard",
-        broken_query="SELECT c.name, COUNT(o.id) AS order_count, SUM(oi.qty * p.price) AS revenue FORM customers c LFT JOIN orders o ON c.id = o.customer_id LFT JOIN order_items oi ON o.id = oi.order_id INNE JOIN products p ON oi.product_id = p.id WHER o.created_at >= '2024-01-01' GRUP BY c.name HAVNG revenue > 5000 ORDR BY revenue DSC LIMT 10",
-        canonical_answer="SELECT c.name, COUNT(o.id) AS order_count, SUM(oi.qty * p.price) AS revenue FROM customers c LEFT JOIN orders o ON c.id = o.customer_id LEFT JOIN order_items oi ON o.id = oi.order_id INNER JOIN products p ON oi.product_id = p.id WHERE o.created_at >= '2024-01-01' GROUP BY c.name HAVING revenue > 5000 ORDER BY revenue DESC LIMIT 10",
+        broken_query=(
+            "SELECT c.name, COUNT(o.id) AS order_count, SUM(oi.qty * p.price) AS revenue "
+            "FORM customers c LFT JOIN orders o ON c.id = o.customer_id "
+            "LFT JOIN order_items oi ON o.id = oi.order_id "
+            "INNE JOIN products p ON oi.product_id = p.id "
+            "WHER o.created_at >= '2024-01-01' "
+            "GRUP BY c.name "
+            "HAVNG revenue > 5000 "
+            "ORDR BY revenue DSC LIMT 10"
+        ),
+        canonical_answer=(
+            "SELECT c.name, COUNT(o.id) AS order_count, SUM(oi.qty * p.price) AS revenue "
+            "FROM customers c LEFT JOIN orders o ON c.id = o.customer_id "
+            "LEFT JOIN order_items oi ON o.id = oi.order_id "
+            "INNER JOIN products p ON oi.product_id = p.id "
+            "WHERE o.created_at >= '2024-01-01' "
+            "GROUP BY c.name "
+            "HAVING revenue > 5000 "
+            "ORDER BY revenue DESC LIMIT 10"
+        ),
         schema_context=(
             "customers(id INT, name VARCHAR, email VARCHAR)\n"
             "orders(id INT, customer_id INT, created_at DATETIME)\n"
@@ -34,13 +68,284 @@ HARD_TASKS = [
     SQLTask(
         task_id="hard_003",
         difficulty="hard",
-        broken_query="SELECT dept, AVG(salary) AS avg_sal, MAX(salary) AS max_sal FORM employees WHER hire_date BETWEN '2020-01-01' AND '2023-12-31' AND status = 'active' GRUP BY dept HAVNG AVG(salary) > 60000 ORDR BY avg_sal DSC",
-        canonical_answer="SELECT dept, AVG(salary) AS avg_sal, MAX(salary) AS max_sal FROM employees WHERE hire_date BETWEEN '2020-01-01' AND '2023-12-31' AND status = 'active' GROUP BY dept HAVING AVG(salary) > 60000 ORDER BY avg_sal DESC",
+        broken_query=(
+            "SELECT dept, AVG(salary) AS avg_sal, MAX(salary) AS max_sal "
+            "FORM employees "
+            "WHER hire_date BETWEN '2020-01-01' AND '2023-12-31' AND status = 'active' "
+            "GRUP BY dept "
+            "HAVNG AVG(salary) > 60000 "
+            "ORDR BY avg_sal DSC"
+        ),
+        canonical_answer=(
+            "SELECT dept, AVG(salary) AS avg_sal, MAX(salary) AS max_sal "
+            "FROM employees "
+            "WHERE hire_date BETWEEN '2020-01-01' AND '2023-12-31' AND status = 'active' "
+            "GROUP BY dept "
+            "HAVING AVG(salary) > 60000 "
+            "ORDER BY avg_sal DESC"
+        ),
+        schema_context=(
+            "employees(id INT, name VARCHAR, dept VARCHAR, salary DECIMAL, "
+            "hire_date DATE, status VARCHAR)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_004",
+        difficulty="hard",
+        broken_query=(
+            "SELCT p.name, cat.category_name, SUM(oi.quantity) AS total_sold "
+            "FORM products p "
+            "INNE JOIN categories cat ON p.cat_id = cat.id "
+            "INNE JOIN order_items oi ON p.id = oi.prod_id "
+            "INNE JOIN orders o ON oi.order_id = o.id "
+            "WHER o.status = 'completed' AND o.order_date >= '2024-01-01' "
+            "GRUP BY p.name, cat.category_name "
+            "ORDR BY total_sold DSC LIMT 20"
+        ),
+        canonical_answer=(
+            "SELECT p.name, cat.category_name, SUM(oi.quantity) AS total_sold "
+            "FROM products p "
+            "INNER JOIN categories cat ON p.category_id = cat.id "
+            "INNER JOIN order_items oi ON p.id = oi.product_id "
+            "INNER JOIN orders o ON oi.order_id = o.id "
+            "WHERE o.status = 'completed' AND o.order_date >= '2024-01-01' "
+            "GROUP BY p.name, cat.category_name "
+            "ORDER BY total_sold DESC LIMIT 20"
+        ),
+        schema_context=(
+            "products(id INT, name VARCHAR, category_id INT, price DECIMAL)\n"
+            "categories(id INT, category_name VARCHAR)\n"
+            "order_items(id INT, order_id INT, product_id INT, quantity INT)\n"
+            "orders(id INT, status VARCHAR, order_date DATE)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_005",
+        difficulty="hard",
+        broken_query=(
+            "SELECT a.title, u.username, COUNT(c.id) AS comment_count, "
+            "AVG(r.score) AS avg_score "
+            "FORM articles a "
+            "INNE JOIN users u ON a.author_id = u.id "
+            "LFT JOIN comments c ON a.id = c.article_id "
+            "LFT JOIN ratings r ON a.id = r.article_id "
+            "WHER a.published_at BETWEN '2024-01-01' AND '2024-06-30' "
+            "GRUP BY a.title, u.username "
+            "HAVNG COUNT(c.id) > 5 "
+            "ORDR BY avg_score DSC"
+        ),
+        canonical_answer=(
+            "SELECT a.title, u.username, COUNT(c.id) AS comment_count, "
+            "AVG(r.score) AS avg_score "
+            "FROM articles a "
+            "INNER JOIN users u ON a.author_id = u.id "
+            "LEFT JOIN comments c ON a.id = c.article_id "
+            "LEFT JOIN ratings r ON a.id = r.article_id "
+            "WHERE a.published_at BETWEEN '2024-01-01' AND '2024-06-30' "
+            "GROUP BY a.title, u.username "
+            "HAVING COUNT(c.id) > 5 "
+            "ORDER BY avg_score DESC"
+        ),
+        schema_context=(
+            "articles(id INT, title VARCHAR, author_id INT, published_at DATE)\n"
+            "users(id INT, username VARCHAR, email VARCHAR)\n"
+            "comments(id INT, article_id INT, user_id INT, body TEXT)\n"
+            "ratings(id INT, article_id INT, user_id INT, score FLOAT)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_006",
+        difficulty="hard",
+        broken_query=(
+            "SELCT w.warehouse_name, p.name, SUM(inv.qty) AS stock_total "
+            "FORM warehouses w "
+            "INNE JOIN inventory inv ON w.id = inv.wrhs_id "
+            "INNE JOIN products p ON inv.product_id = p.id "
+            "WHER inv.last_updated >= '2024-01-01' "
+            "GRUP BY w.warehouse_name, p.name "
+            "HAVNG SUM(inv.qty) < 50 "
+            "ORDR BY stock_total ASC"
+        ),
+        canonical_answer=(
+            "SELECT w.warehouse_name, p.name, SUM(inv.qty) AS stock_total "
+            "FROM warehouses w "
+            "INNER JOIN inventory inv ON w.id = inv.warehouse_id "
+            "INNER JOIN products p ON inv.product_id = p.id "
+            "WHERE inv.last_updated >= '2024-01-01' "
+            "GROUP BY w.warehouse_name, p.name "
+            "HAVING SUM(inv.qty) < 50 "
+            "ORDER BY stock_total ASC"
+        ),
+        schema_context=(
+            "warehouses(id INT, warehouse_name VARCHAR, location VARCHAR)\n"
+            "inventory(id INT, warehouse_id INT, product_id INT, qty INT, "
+            "last_updated DATE)\n"
+            "products(id INT, name VARCHAR, sku VARCHAR, price DECIMAL)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_007",
+        difficulty="hard",
+        broken_query=(
+            "SELECT s.student_name, co.course_name, "
+            "AVG(g.grade) AS avg_grade, COUNT(g.id) AS assignments_done "
+            "FORM students s "
+            "INNE JOIN enrollments en ON s.id = en.student_id "
+            "INNE JOIN courses co ON en.course_id = co.id "
+            "INNE JOIN grades g ON s.id = g.std_id AND co.id = g.course_id "
+            "WHER en.semester = 'Fall2024' "
+            "GRUP BY s.student_name, co.course_name "
+            "HAVNG AVG(g.grade) >= 70 "
+            "ORDR BY avg_grade DSC"
+        ),
+        canonical_answer=(
+            "SELECT s.student_name, co.course_name, "
+            "AVG(g.grade) AS avg_grade, COUNT(g.id) AS assignments_done "
+            "FROM students s "
+            "INNER JOIN enrollments en ON s.id = en.student_id "
+            "INNER JOIN courses co ON en.course_id = co.id "
+            "INNER JOIN grades g ON s.id = g.student_id AND co.id = g.course_id "
+            "WHERE en.semester = 'Fall2024' "
+            "GROUP BY s.student_name, co.course_name "
+            "HAVING AVG(g.grade) >= 70 "
+            "ORDER BY avg_grade DESC"
+        ),
+        schema_context=(
+            "students(id INT, student_name VARCHAR, email VARCHAR)\n"
+            "enrollments(id INT, student_id INT, course_id INT, semester VARCHAR)\n"
+            "courses(id INT, course_name VARCHAR, credits INT)\n"
+            "grades(id INT, student_id INT, course_id INT, grade FLOAT)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_008",
+        difficulty="hard",
+        broken_query=(
+            "SELCT e.name, m.name AS manager_name, d.dept_name, "
+            "e.salary, AVG(e2.salary) AS dept_avg "
+            "FORM employees e "
+            "LFT JOIN employees m ON e.manager_id = m.id "
+            "INNE JOIN departments d ON e.dept_id = d.id "
+            "INNE JOIN employees e2 ON e2.dept_id = e.dept_id "
+            "WHER e.salary > 50000 "
+            "GRUP BY e.name, m.name, d.dept_name, e.salary "
+            "HAVNG e.salary > AVG(e2.salary) "
+            "ORDR BY e.salary DSC"
+        ),
+        canonical_answer=(
+            "SELECT e.name, m.name AS manager_name, d.dept_name, "
+            "e.salary, AVG(e2.salary) AS dept_avg "
+            "FROM employees e "
+            "LEFT JOIN employees m ON e.manager_id = m.id "
+            "INNER JOIN departments d ON e.dept_id = d.id "
+            "INNER JOIN employees e2 ON e2.dept_id = e.dept_id "
+            "WHERE e.salary > 50000 "
+            "GROUP BY e.name, m.name, d.dept_name, e.salary "
+            "HAVING e.salary > AVG(e2.salary) "
+            "ORDER BY e.salary DESC"
+        ),
+        schema_context=(
+            "employees(id INT, name VARCHAR, dept_id INT, manager_id INT, "
+            "salary DECIMAL)\n"
+            "departments(id INT, dept_name VARCHAR, budget DECIMAL)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_009",
+        difficulty="hard",
+        broken_query=(
+            "SELECT t.tag_name, COUNT(DISTINCT pt.post_id) AS post_count, "
+            "AVG(p.views) AS avg_views "
+            "FORM tags t "
+            "INNE JOIN post_tags pt ON t.id = pt.tag_id "
+            "INNE JOIN posts p ON pt.post_id = p.id "
+            "INNE JOIN users u ON p.user_id = u.id "
+            "WHER p.created_at >= '2024-01-01' AND u.role = 'author' "
+            "GRUP BY t.tag_name "
+            "HAVNG COUNT(DISTINCT pt.post_id) > 10 "
+            "ORDR BY avg_views DSC LIMT 15"
+        ),
+        canonical_answer=(
+            "SELECT t.tag_name, COUNT(DISTINCT pt.post_id) AS post_count, "
+            "AVG(p.views) AS avg_views "
+            "FROM tags t "
+            "INNER JOIN post_tags pt ON t.id = pt.tag_id "
+            "INNER JOIN posts p ON pt.post_id = p.id "
+            "INNER JOIN users u ON p.user_id = u.id "
+            "WHERE p.created_at >= '2024-01-01' AND u.role = 'author' "
+            "GROUP BY t.tag_name "
+            "HAVING COUNT(DISTINCT pt.post_id) > 10 "
+            "ORDER BY avg_views DESC LIMIT 15"
+        ),
+        schema_context=(
+            "tags(id INT, tag_name VARCHAR)\n"
+            "post_tags(post_id INT, tag_id INT)\n"
+            "posts(id INT, user_id INT, title VARCHAR, views INT, "
+            "created_at DATE)\n"
+            "users(id INT, username VARCHAR, role VARCHAR)"
+        ),
+        error_hint=None,
+        max_steps=4,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="hard_010",
+        difficulty="hard",
+        broken_query=(
+            "SELCT proj.name AS project_name, emp.name AS employee_name, "
+            "SUM(ts.hours) AS total_hours, ts.week_start "
+            "FORM projects proj "
+            "INNE JOIN project_members pm ON proj.id = pm.proj_id "
+            "INNE JOIN employees emp ON pm.employee_id = emp.id "
+            "INNE JOIN timesheets ts ON emp.id = ts.emp_id "
+            "AND proj.id = ts.project_id "
+            "WHER ts.week_start BETWEN '2024-01-01' AND '2024-03-31' "
+            "AND proj.status = 'active' "
+            "GRUP BY proj.name, emp.name, ts.week_start "
+            "HAVNG SUM(ts.hours) > 40 "
+            "ORDR BY total_hours DSC"
+        ),
+        canonical_answer=(
+            "SELECT proj.name AS project_name, emp.name AS employee_name, "
+            "SUM(ts.hours) AS total_hours, ts.week_start "
+            "FROM projects proj "
+            "INNER JOIN project_members pm ON proj.id = pm.project_id "
+            "INNER JOIN employees emp ON pm.employee_id = emp.id "
+            "INNER JOIN timesheets ts ON emp.id = ts.employee_id "
+            "AND proj.id = ts.project_id "
+            "WHERE ts.week_start BETWEEN '2024-01-01' AND '2024-03-31' "
+            "AND proj.status = 'active' "
+            "GROUP BY proj.name, emp.name, ts.week_start "
+            "HAVING SUM(ts.hours) > 40 "
+            "ORDER BY total_hours DESC"
+        ),
         schema_context=(
-            "employees(id INT, name VARCHAR, dept VARCHAR, salary DECIMAL, hire_date DATE, status VARCHAR)"
+            "projects(id INT, name VARCHAR, status VARCHAR, budget DECIMAL)\n"
+            "project_members(id INT, project_id INT, employee_id INT, role VARCHAR)\n"
+            "employees(id INT, name VARCHAR, dept_id INT, hourly_rate DECIMAL)\n"
+            "timesheets(id INT, employee_id INT, project_id INT, "
+            "hours DECIMAL, week_start DATE)"
         ),
         error_hint=None,
         max_steps=4,
         grader=grade,
     ),
-]
+]

sql_env/tasks/medium.py

@@ -42,4 +42,84 @@ MEDIUM_TASKS = [
         error_hint=None,
         grader=grade,
     ),
-]
+    SQLTask(
+        task_id="medium_006",
+        difficulty="medium",
+        broken_query="SELECT category, AVG(price) AS avg_price FORM products GRUP BY category ORDR BY avg_price DESC",
+        canonical_answer="SELECT category, AVG(price) AS avg_price FROM products GROUP BY category ORDER BY avg_price DESC",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_007",
+        difficulty="medium",
+        broken_query="SELECT u.email, COUNT(o.id) AS order_count FORM users u LFT JOIN orders o ON u.id = o.user_id GRUP BY u.email",
+        canonical_answer="SELECT u.email, COUNT(o.id) AS order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.email",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_008",
+        difficulty="medium",
+        broken_query="SELECT id, title FORM articles WEHRE published = 1 ORDR BY created_at DESC LIMT 10",
+        canonical_answer="SELECT id, title FROM articles WHERE published = 1 ORDER BY created_at DESC LIMIT 10",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_009",
+        difficulty="medium",
+        broken_query="SELECT region, SUM(revenue) AS total FORM sales GRUP BY region HAVNG SUM(revenue) > 100000",
+        canonical_answer="SELECT region, SUM(revenue) AS total FROM sales GROUP BY region HAVING SUM(revenue) > 100000",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_010",
+        difficulty="medium",
+        broken_query="SELCT p.name, s.quantity FORM products p INNE JOIN stock s ON p.id = s.product_id WEHRE s.quantity < 10",
+        canonical_answer="SELECT p.name, s.quantity FROM products p INNER JOIN stock s ON p.id = s.product_id WHERE s.quantity < 10",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_011",
+        difficulty="medium",
+        broken_query="SELECT customer_id, MAX(amount) AS max_order FORM orders GRUP BY customer_id ORDR BY max_order DESC",
+        canonical_answer="SELECT customer_id, MAX(amount) AS max_order FROM orders GROUP BY customer_id ORDER BY max_order DESC",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_012",
+        difficulty="medium",
+        broken_query="SELECT t.name, COUNT(e.id) AS member_count FORM teams t LFT JOIN employees e ON t.id = e.team_id GRUP BY t.name",
+        canonical_answer="SELECT t.name, COUNT(e.id) AS member_count FROM teams t LEFT JOIN employees e ON t.id = e.team_id GROUP BY t.name",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_013",
+        difficulty="medium",
+        broken_query="SELCT month, SUM(sales) FORM monthly_revenue WEHRE year = 2024 GRUP BY month ORDR BY month ASC",
+        canonical_answer="SELECT month, SUM(sales) FROM monthly_revenue WHERE year = 2024 GROUP BY month ORDER BY month ASC",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_014",
+        difficulty="medium",
+        broken_query="SELECT c.name, AVG(r.rating) AS avg_rating FORM courses c INNE JOIN reviews r ON c.id = r.course_id GRUP BY c.name HAVNG AVG(r.rating) >= 4.0",
+        canonical_answer="SELECT c.name, AVG(r.rating) AS avg_rating FROM courses c INNER JOIN reviews r ON c.id = r.course_id GROUP BY c.name HAVING AVG(r.rating) >= 4.0",
+        error_hint=None,
+        grader=grade,
+    ),
+    SQLTask(
+        task_id="medium_015",
+        difficulty="medium",
+        broken_query="SELECT DISTINCT country FORM customers WEHRE signup_date > '2023-01-01' ORDR BY country ASC",
+        canonical_answer="SELECT DISTINCT country FROM customers WHERE signup_date > '2023-01-01' ORDER BY country ASC",
+        error_hint=None,
+        grader=grade,
+    ),
+]

tests/test_env.py

@@ -0,0 +1,267 @@
+"""
+tests/test_env.py — Smoke tests and grader unit tests.
+
+Run with:  pytest tests/ -v
+"""
+import asyncio
+import pytest
+from sql_env.models import SQLAction, SQLTask
+from sql_env.grader import grade, generate_feedback
+from sql_env.env import SQLCorrectionEnv
+from sql_env.tasks import ALL_TASKS, EASY_TASKS, MEDIUM_TASKS, HARD_TASKS
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _make_task(broken: str, canonical: str, difficulty: str = "easy") -> SQLTask:
+    return SQLTask(
+        task_id="test_task",
+        difficulty=difficulty,
+        broken_query=broken,
+        canonical_answer=canonical,
+    )
+
+
+def _action(query: str) -> SQLAction:
+    return SQLAction(corrected_query=query)
+
+
+# ── Grader unit tests ─────────────────────────────────────────────────────────
+
+class TestGrader:
+    def test_exact_match_returns_one(self):
+        task = _make_task(
+            "SELECT * FORM users",
+            "SELECT * FROM users",
+        )
+        reward = grade(_action("SELECT * FROM users"), task)
+        assert reward.value == 1.0
+
+    def test_exact_match_case_insensitive(self):
+        task = _make_task(
+            "SELECT * FORM users",
+            "SELECT * FROM users",
+        )
+        reward = grade(_action("select * from users"), task)
+        assert reward.value == 1.0
+
+    def test_exact_match_trailing_semicolon(self):
+        task = _make_task(
+            "SELECT * FORM users",
+            "SELECT * FROM users",
+        )
+        reward = grade(_action("SELECT * FROM users;"), task)
+        assert reward.value == 1.0
+
+    def test_wrong_answer_not_one(self):
+        task = _make_task(
+            "SELECT * FORM users",
+            "SELECT * FROM users",
+        )
+        reward = grade(_action("SELECT * FORM users"), task)
+        assert reward.value < 1.0
+
+    def test_completely_wrong_returns_zero(self):
+        task = _make_task(
+            "SELECT * FORM users",
+            "SELECT * FROM users",
+        )
+        reward = grade(_action("hello world"), task)
+        assert reward.value == 0.0
+
+    def test_basic_structure_returns_02(self):
+        task = _make_task(
+            "SELECT * FORM users WHERE id = 1",
+            "SELECT * FROM users WHERE id = 1",
+        )
+        # Correct structure, still has FROM typo
+        reward = grade(_action("SELECT * FORM users WHERE id = 1"), task)
+        assert reward.value == pytest.approx(0.2, abs=0.05)
+
+    def test_reward_range(self):
+        task = _make_task(
+            "SELCT * FORM users WEHRE id = 1",
+            "SELECT * FROM users WHERE id = 1",
+        )
+        for query in [
+            "hello world",
+            "SELECT * FORM users",
+            "SELECT * FROM users WHERE id = 1",
+            "select * from users where id = 1",
+        ]:
+            reward = grade(_action(query), task)
+            assert 0.0 <= reward.value <= 1.0, (
+                f"Reward {reward.value} out of [0, 1] for query: {query}"
+            )
+
+    def test_feedback_not_empty(self):
+        task = _make_task("SELECT * FORM users", "SELECT * FROM users")
+        reward = grade(_action("SELECT * FROM users"), task)
+        fb = generate_feedback(_action("SELECT * FROM users"), task, reward)
+        assert isinstance(fb, str) and len(fb) > 0
+
+
+# ── Task catalogue tests ──────────────────────────────────────────────────────
+
+class TestTaskCatalogue:
+    def test_easy_task_count(self):
+        assert len(EASY_TASKS) >= 10, "Need at least 10 easy tasks"
+
+    def test_medium_task_count(self):
+        assert len(MEDIUM_TASKS) >= 10, "Need at least 10 medium tasks"
+
+    def test_hard_task_count(self):
+        assert len(HARD_TASKS) >= 5, "Need at least 5 hard tasks"
+
+    def test_all_task_ids_unique(self):
+        all_ids = [t.task_id for tasks in ALL_TASKS.values() for t in tasks]
+        assert len(all_ids) == len(set(all_ids)), "Duplicate task IDs found"
+
+    def test_easy_tasks_have_hints(self):
+        for task in EASY_TASKS:
+            assert task.error_hint is not None and len(task.error_hint) > 0, (
+                f"Easy task {task.task_id} missing error_hint"
+            )
+
+    def test_hard_tasks_have_schema(self):
+        for task in HARD_TASKS:
+            assert task.schema_context is not None and len(task.schema_context) > 0, (
+                f"Hard task {task.task_id} missing schema_context"
+            )
+
+    def test_canonical_answers_are_valid_sql(self):
+        """Canonical answers must at least contain SELECT and FROM."""
+        for difficulty, tasks in ALL_TASKS.items():
+            for task in tasks:
+                upper = task.canonical_answer.upper()
+                assert "SELECT" in upper, (
+                    f"{task.task_id}: canonical_answer missing SELECT"
+                )
+                assert "FROM" in upper, (
+                    f"{task.task_id}: canonical_answer missing FROM"
+                )
+
+    def test_grading_canonical_answer_returns_perfect(self):
+        """Every task must return 1.0 when given its own canonical answer."""
+        for difficulty, tasks in ALL_TASKS.items():
+            for task in tasks:
+                action = _action(task.canonical_answer)
+                reward = grade(action, task)
+                assert reward.value == 1.0, (
+                    f"{task.task_id}: canonical answer did not score 1.0 "
+                    f"(got {reward.value})"
+                )
+
+    def test_grading_broken_query_below_perfect(self):
+        """Broken queries must score below 1.0."""
+        for difficulty, tasks in ALL_TASKS.items():
+            for task in tasks:
+                action = _action(task.broken_query)
+                reward = grade(action, task)
+                assert reward.value < 1.0, (
+                    f"{task.task_id}: broken query unexpectedly scored 1.0"
+                )
+
+
+# ── Environment integration tests ─────────────────────────────────────────────
+
+class TestEnvironment:
+    def test_reset_returns_observation(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy")
+            obs = await env.reset()
+            assert obs.task_id is not None
+            assert obs.broken_query is not None
+            assert obs.step_number == 0
+            assert obs.steps_remaining == 5
+
+        asyncio.run(run())
+
+    def test_step_returns_result(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy")
+            await env.reset()
+            result = await env.step(_action("SELECT * FROM users WHERE id = 1"))
+            assert 0.0 <= result.reward <= 1.0
+            assert isinstance(result.done, bool)
+            assert result.observation.step_number == 1
+
+        asyncio.run(run())
+
+    def test_steps_remaining_decrements(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy")
+            await env.reset()
+            result = await env.step(_action("SELECT * FROM x"))
+            assert result.observation.steps_remaining == 4
+
+        asyncio.run(run())
+
+    def test_correct_answer_terminates(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy", task_index=0)
+            await env.reset()
+            canonical = EASY_TASKS[0].canonical_answer
+            result = await env.step(_action(canonical))
+            assert result.done is True
+            assert result.reward == pytest.approx(1.0)
+
+        asyncio.run(run())
+
+    def test_max_steps_terminates(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy", task_index=0)
+            await env.reset()
+            result = None
+            for _ in range(5):
+                result = await env.step(_action("SELECT * FORM users"))
+            assert result.done is True
+
+        asyncio.run(run())
+
+    def test_done_episode_raises(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy", task_index=0)
+            await env.reset()
+            canonical = EASY_TASKS[0].canonical_answer
+            await env.step(_action(canonical))  # this terminates
+            with pytest.raises(RuntimeError):
+                await env.step(_action("SELECT 1"))
+
+        asyncio.run(run())
+
+    def test_medium_hint_hidden(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="medium")
+            obs = await env.reset()
+            assert obs.error_hint is None
+
+        asyncio.run(run())
+
+    def test_hard_schema_present(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="hard")
+            obs = await env.reset()
+            assert obs.schema_context is not None
+
+        asyncio.run(run())
+
+    def test_state_reflects_progress(self):
+        async def run():
+            env = SQLCorrectionEnv(difficulty="easy", task_index=0)
+            await env.reset()
+            await env.step(_action("SELECT * FORM users"))
+            state = await env.state()
+            assert state["step_count"] == 1
+            assert state["done"] is False
+
+        asyncio.run(run())
+
+    def test_all_difficulties_reset(self):
+        async def run():
+            for diff in ["easy", "medium", "hard"]:
+                env = SQLCorrectionEnv(difficulty=diff)
+                obs = await env.reset()
+                assert obs.broken_query is not None
+
+        asyncio.run(run())