fix: restore UI, add /tasks /grade, clamp scores to (0.05, 0.95)

inference.py

@@ -7,80 +7,114 @@ from openai import OpenAI
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-ENV_URL = os.getenv("ENV_URL", "http://localhost:8000") # defaults to common uvicorn port but user specified 7860
+ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
 MAX_STEPS = 6
 TEMPERATURE = 0.1
 MAX_TOKENS = 500
 
 client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
 
-SYSTEM_PROMPT = """You are an expert SQL developer.
-For write_query tasks: write a correct SQL SELECT query.
-For fix_query tasks: fix the broken SQL provided.
-For optimize_query tasks: rewrite slow SQL using CTEs or JOINs instead of subqueries.
+SYSTEM_PROMPT = """You are an expert SQL developer working with a SQLite database.
 
-ALWAYS respond with ONLY a JSON object like this:
-{"action_type": "write_query", "sql": "SELECT ...", "explanation": "reason"}
+Your job depends on the task type:
+- write_query: Write a correct SQL SELECT query from scratch
+- fix_query: You are shown broken SQL. Fix it so it runs correctly
+- optimize_query: You are shown slow SQL. Rewrite it using CTEs (WITH clause) 
+  or window functions (ROW_NUMBER, RANK) instead of correlated subqueries
 
-Rules: Only SELECT allowed. No DROP DELETE INSERT UPDATE CREATE ALTER."""
+RESPONSE FORMAT — always respond with ONLY this JSON:
+{"action_type": "write_query", "sql": "SELECT ...", "explanation": "brief reason"}
 
-def build_prompt(obs: dict) -> str:
-    obs_data = obs.get("observation", obs)
-    task_desc = obs_data.get("task_description", "")
-    schema_info = obs_data.get("schema_info", "")
-    sample_data = obs_data.get("sample_data", "")
-    hints = obs_data.get("metadata", [])
-    last_sql = obs_data.get("last_sql", "")
-    last_result = obs_data.get("last_result", "")
-    last_error = obs_data.get("last_error", "")
-    step_count = obs_data.get("step_count", 0)
-    feedback = obs_data.get("feedback", "")
+RULES:
+- Only SELECT statements are allowed
+- No DROP, DELETE, INSERT, UPDATE, CREATE, ALTER, TRUNCATE
+- For fix_query: keep the same intent, just fix the bugs
+- For optimize_query: use WITH clause or window functions
+- Respond ONLY with the JSON object, no other text"""
+
+def build_prompt(obs: dict, step: int, history: list) -> str:
+    parts = []
+    parts.append(f"=== TASK ===\n{obs.get('task_description', '')}")
     
-    prompt = f"Task Description:\n{task_desc}\n\nSchema Info:\n{schema_info}\n\nSample Data:\n{sample_data}\n"
+    schema = obs.get('schema_info', '')
+    if schema:
+        parts.append(f"=== DATABASE SCHEMA ===\n{chr(10).join(schema.split(chr(10))[:30])}")
+        
+    data = obs.get('sample_data', '')
+    if data:
+        parts.append(f"=== SAMPLE DATA ===\n{chr(10).join(data.split(chr(10))[:20])}")
+        
+    hints = obs.get('expected_description', '')
     if hints:
-        prompt += f"\nHints: {', '.join(hints)}\n"
+        parts.append(f"=== HINTS ===\n{hints}")
+        
+    last_sql = obs.get('last_sql')
     if last_sql:
-        prompt += f"\nLast SQL Submitted: {last_sql}\n"
+        parts.append(f"=== YOUR PREVIOUS SQL ===\n{last_sql}")
+        
+    last_result = obs.get('last_result')
     if last_result:
-        prompt += f"Result of Last SQL: {last_result}\n"
+        parts.append(f"=== RESULT OF PREVIOUS SQL ===\n{last_result}")
+        
+    last_error = obs.get('last_error')
     if last_error:
-        prompt += f"Error Message: {last_error}\n"
-    if step_count > 0 and feedback:
-        prompt += f"Feedback from Grader: {feedback}\n"
+        parts.append(f"=== ERROR ===\n{last_error}")
+        
+    feedback = obs.get('feedback', '')
+    if feedback and step > 1:
+        parts.append(f"=== GRADER FEEDBACK ===\n{feedback}")
+        
+    if history:
+        hist_str = "\n".join(history[-3:])
+        parts.append(f"=== RECENT HISTORY ===\n{hist_str}")
         
-    prompt += "\nRespond with JSON action only."
-    return prompt
+    parts.append("Respond with ONLY a JSON action object.")
+    return "\n\n".join(parts)
 
 def parse_action(response_text: str, task_type: str) -> dict:
+    text = response_text.strip()
     try:
-        return json.loads(response_text)
+        return json.loads(text)
     except json.JSONDecodeError:
-        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
-        if json_match:
-            try:
-                return json.loads(json_match.group(0))
-            except json.JSONDecodeError:
-                pass
-        
-        sql_match = re.search(r'(?i)SELECT\s+.*', response_text, re.DOTALL)
-        sql = sql_match.group(0).strip() if sql_match else "SELECT 1"
-        if sql.endswith('```'): sql = sql[:-3].strip()
-        
-        return {"action_type": task_type, "sql": sql, "explanation": "Fallback parsed"}
+        pass
+        
+    match = re.search(r'\{[^{}]+\}', text)
+    if match:
+        try:
+            return json.loads(match.group(0))
+        except json.JSONDecodeError:
+            pass
+            
+    sql_match = re.search(r'(?:SELECT|WITH).+', text, re.DOTALL | re.IGNORECASE)
+    if sql_match:
+        sql = sql_match.group(0).strip()
+        if sql.endswith("```"): sql = sql[:-3].strip()
+        return {"action_type": task_type, "sql": sql, "explanation": "Regex parsed"}
+        
+    return {"action_type": task_type, "sql": "SELECT 1", "explanation": "parse failed"}
 
 def run_episode(task_id: str) -> dict:
-    res = requests.post(f"{ENV_URL}/reset", json={"task_id": task_id})
-    if res.status_code != 200:
+    try:
+        res = requests.post(f"{ENV_URL}/reset", json={"task_id": task_id, "difficulty": None})
+        res.raise_for_status()
+        obs_obj = res.json()
+        obs = obs_obj["observation"]
+    except Exception as e:
+        print(f"  Error resetting environment: {e}")
         return {"task_id": task_id, "best_reward": 0.05}
-    obs_obj = res.json()
+        
+    print(f"  Task: {task_id} | {obs.get('task_type', '')}")
+    print(f"  Desc: {obs.get('task_description', '')[:100]}...")
+    
+    best_reward = 0.0
+    history = []
     
-    best_reward = 0.05
-    for step in range(MAX_STEPS):
-        obs = obs_obj.get("observation", obs_obj)
+    for step in range(1, MAX_STEPS + 1):
         if obs.get("done", False):
+            print(f"  Done at step {step}")
             break
             
-        prompt = build_prompt(obs_obj)
+        prompt = build_prompt(obs, step, history)
         
         try:
             completion = client.chat.completions.create(
@@ -92,73 +126,81 @@ def run_episode(task_id: str) -> dict:
                 temperature=TEMPERATURE,
                 max_tokens=MAX_TOKENS
             )
-            response_text = completion.choices[0].message.content or "{}"
-        except Exception:
-            response_text = "{}"
+            response_text = completion.choices[0].message.content or ""
+        except Exception as e:
+            print(f"  API Error: {e}")
+            response_text = ""
             
         action = parse_action(response_text, obs.get("task_type", "write_query"))
         
-        res = requests.post(f"{ENV_URL}/step", json=action)
-        if res.status_code != 200:
-            break
-        obs_obj = res.json()
-        reward = obs_obj.get("reward", 0.05)
+        print(f"    Step {step}: {action.get('sql', '')[:70]}...")
         
+        try:
+            step_res = requests.post(f"{ENV_URL}/step", json=action)
+            step_res.raise_for_status()
+            step_data = step_res.json()
+        except Exception as e:
+            print(f"  Env Step Error: {e}")
+            break
+            
+        reward = step_data.get("reward", 0.05)
         best_reward = max(best_reward, reward)
-        print(f"Step {step+1}: SQL={action.get('sql', '')[:60]}... Reward={reward:.3f} Feedback={obs_obj.get('observation', {}).get('feedback', '')}")
+        obs = step_data.get("observation", {})
+        
+        history.append(f"Step {step}: reward={reward:.3f} | {obs.get('feedback', '')[:60]}")
+        print(f"    Reward: {reward:.3f} | Done: {obs.get('done', False)} | {obs.get('feedback', '')[:60]}")
         
-        if obs_obj.get("done", obs_obj.get("observation", {}).get("done", False)):
+        if obs.get("done", False):
             break
             
     return {"task_id": task_id, "best_reward": round(best_reward, 3)}
 
 def main():
-    print(f"--- SQL Agent Inference ---")
-    print(f"Model: {MODEL_NAME}")
-    print(f"Env URL: {ENV_URL}")
+    print("=" * 65)
+    print("SQL Debugger OpenEnv — Baseline Inference")
+    print(f"Model : {MODEL_NAME}")
+    print(f"Env   : {ENV_URL}")
+    print("=" * 65)
     
-    health = {"status": "unreachable"}
     try:
-        health = requests.get(f"{ENV_URL}/health").json()
-    except Exception:
-        pass
-    print(f"Health: {health}")
-    
+        health_res = requests.get(f"{ENV_URL}/health")
+        print(f"Health: {health_res.json()}")
+    except Exception as e:
+        print(f"Health Check Failed: {e}")
+        
     task_ids = ["easy_01", "easy_02", "medium_01", "medium_02", "hard_01", "hard_02"]
     results = []
     
-    for tid in task_ids:
-        print(f"\n=== Task: {tid} ===")
+    for i, tid in enumerate(task_ids):
+        print(f"\\n[{i+1}/6] Running {tid}...")
         try:
             res = run_episode(tid)
         except Exception as e:
+            print(f"  Error running {tid}: {e}")
             res = {"task_id": tid, "best_reward": 0.05}
-            print(f"Error running task {tid}: {e}")
         results.append(res)
         
-    print("\n=== FINAL SCORES ===")
-    easy_scores = []
-    medium_scores = []
-    hard_scores = []
+    print("\n" + "=" * 65)
+    print("BASELINE SCORES")
+    print("=" * 65)
     
     for r in results:
-        t = r["task_id"]
-        v = r["best_reward"]
-        print(f"{t}: {v}")
-        if t.startswith("easy"): easy_scores.append(v)
-        elif t.startswith("medium"): medium_scores.append(v)
-        elif t.startswith("hard"): hard_scores.append(v)
+        print(f"  {r['task_id']:15s}: {r['best_reward']:.3f}")
         
+    easy_scores = [r['best_reward'] for r in results if r['task_id'].startswith('easy')]
+    medium_scores = [r['best_reward'] for r in results if r['task_id'].startswith('medium')]
+    hard_scores = [r['best_reward'] for r in results if r['task_id'].startswith('hard')]
+    
     easy_avg = sum(easy_scores)/len(easy_scores) if easy_scores else 0.0
     medium_avg = sum(medium_scores)/len(medium_scores) if medium_scores else 0.0
     hard_avg = sum(hard_scores)/len(hard_scores) if hard_scores else 0.0
-    overall = (easy_avg + medium_avg + hard_avg) / 3.0
+    overall = sum(r['best_reward'] for r in results) / len(results) if results else 0.0
     
-    print("\n--- AVERAGES ---")
-    print(f"Easy Average:   {easy_avg:.3f}")
-    print(f"Medium Average: {medium_avg:.3f}")
-    print(f"Hard Average:   {hard_avg:.3f}")
-    print(f"Overall Score:  {overall:.3f}")
+    print(f"\\nEasy average    : {easy_avg:.3f}")
+    print(f"Medium average  : {medium_avg:.3f}")
+    print(f"Hard average    : {hard_avg:.3f}")
+    print(f"Overall average : {overall:.3f}")
+    print("=" * 65)
 
 if __name__ == "__main__":
     main()

openenv.yaml

@@ -1,9 +1,10 @@
 name: sql-debugger-env
 version: "1.0.0"
 description: >
-  An RL environment for training AI agents to write, fix, and optimize
-  SQL queries against a real SQLite database with employees, departments,
-  projects and project assignments tables.
+  An RL environment for training AI agents to write, debug, and optimize
+  SQL queries against a real SQLite database. Three task types cover
+  real data analyst work: writing queries from specs, fixing broken SQL,
+  and optimizing slow correlated subqueries.
 
 tasks:
   - id: easy_01
@@ -11,42 +12,54 @@ tasks:
     difficulty: easy
     type: write_query
     grader: true
-    description: "Find Engineering employees with salary above 90000, return name and salary ordered by salary DESC"
+    description: >
+      Find all employees in the Engineering department with salary above 
+      90000. Return name and salary ordered by salary descending.
 
   - id: easy_02
     name: Count by department
     difficulty: easy
     type: write_query
     grader: true
-    description: "Count employees per department, return department and count ordered by count DESC"
+    description: >
+      Count how many employees are in each department. Return department 
+      name and count ordered by count descending.
 
   - id: medium_01
     name: Fix broken JOIN
     difficulty: medium
     type: fix_query
     grader: true
-    description: "Fix the broken JOIN query missing ON keyword and using wrong table alias in WHERE"
+    description: >
+      Fix the broken query that returns employees on active projects. 
+      Bugs: missing ON keyword in JOIN, wrong table alias in WHERE clause.
 
   - id: medium_02
     name: Fix wrong GROUP BY
     difficulty: medium
     type: fix_query
     grader: true
-    description: "Fix the GROUP BY clause that incorrectly groups by id instead of department"
+    description: >
+      Fix the query that computes average salary per department but 
+      incorrectly groups by id instead of department.
 
   - id: hard_01
     name: Optimize correlated subquery
     difficulty: hard
     type: optimize_query
     grader: true
-    description: "Replace correlated subqueries with window functions or CTEs to find top earner per department"
+    description: >
+      Replace correlated subqueries with window functions or CTEs to 
+      find the top earner per department with total hours worked.
 
   - id: hard_02
     name: Eliminate N+1 problem
     difficulty: hard
     type: optimize_query
     grader: true
-    description: "Rewrite N+1 correlated subquery using LEFT JOIN and GROUP BY"
+    description: >
+      Rewrite the N+1 correlated subquery using LEFT JOIN and GROUP BY 
+      to get department stats in a single efficient query.
 
 action_space:
   type: object
@@ -54,28 +67,46 @@ action_space:
     action_type:
       type: string
       enum: [write_query, fix_query, optimize_query]
+      description: The type of SQL action being taken
     sql:
       type: string
-      description: SQL SELECT statement
+      description: A valid SQLite SELECT statement
     explanation:
       type: string
-      description: Optional reasoning
+      description: Optional explanation of reasoning
 
 observation_space:
   type: object
   properties:
-    task_id: {type: string}
-    task_type: {type: string}
-    task_description: {type: string}
-    schema_info: {type: string}
-    sample_data: {type: string}
-    last_sql: {type: string}
-    last_result: {type: string}
-    last_error: {type: string}
-    step_count: {type: integer}
-    done: {type: boolean}
-    reward: {type: number, minimum: 0.05, maximum: 0.95}
-    feedback: {type: string}
+    task_id:
+      type: string
+    task_type:
+      type: string
+    task_description:
+      type: string
+    schema_info:
+      type: string
+      description: DDL schema of all tables
+    sample_data:
+      type: string
+      description: First 3 rows of each table
+    last_sql:
+      type: string
+    last_result:
+      type: string
+    last_error:
+      type: string
+    step_count:
+      type: integer
+    done:
+      type: boolean
+    reward:
+      type: number
+      minimum: 0.05
+      maximum: 0.95
+      description: Strictly between 0 and 1, never exactly 0.0 or 1.0
+    feedback:
+      type: string
 
 reward_range: [0.05, 0.95]
 max_steps_per_episode: 8

server/app.py

@@ -1,12 +1,20 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
-from typing import Optional, Dict, Any
+from typing import Optional
+import json
+import asyncio
 
 from server.environment import SQLEnvironment
 
-app = FastAPI(title="SQL Debugger OpenEnv")
+app = FastAPI(title="SQL Debugger OpenEnv", version="1.0.0")
+
+# One environment instance per HTTP session (stateless endpoints)
+# WebSocket sessions get their own instance
 env = SQLEnvironment()
 
+# ── Pydantic models ───────────────────────────────────────────────────────────
+
 class ResetRequest(BaseModel):
     task_id: Optional[str] = None
     difficulty: Optional[str] = None
@@ -21,8 +29,10 @@ class GradeRequest(BaseModel):
     task_id: str
     action: dict
 
+# ── HTTP endpoints ────────────────────────────────────────────────────────────
+
 @app.get("/")
-def read_root():
+def root():
     return {
         "info": "SQL Agent Environment API",
         "version": "1.0.0",
@@ -31,28 +41,29 @@ def read_root():
     }
 
 @app.get("/health")
-def read_health():
+def health():
     return {"status": "healthy"}
 
-@app.get("/state")
-def read_state():
-    return env.state
-
 @app.post("/reset")
-def do_reset(req: Optional[ResetRequest] = None):
-    if req:
-        obs = env.reset(task_id=req.task_id, difficulty=req.difficulty)
-    else:
-        obs = env.reset()
-    return obs
+def reset(req: ResetRequest = ResetRequest()):
+    return env.reset(task_id=req.task_id, difficulty=req.difficulty)
 
 @app.post("/step")
-def do_step(req: StepRequest):
-    obs = env.step(req.model_dump() if hasattr(req, "model_dump") else req.dict())
-    return obs
+def step(req: StepRequest):
+    return env.step(req.dict())
+
+@app.get("/state")
+def state():
+    return env.state
+
+# ── CRITICAL MISSING ENDPOINTS — ADD THESE ───────────────────────────────────
 
 @app.get("/tasks")
 def list_tasks():
+    """
+    Validator uses this endpoint to find and enumerate all graders.
+    Every task MUST have grader: true.
+    """
     return {
         "tasks": [
             {
@@ -61,7 +72,7 @@ def list_tasks():
                 "difficulty": "easy",
                 "type": "write_query",
                 "grader": True,
-                "description": "Find all employees in Engineering with salary above 90000"
+                "description": "Find all employees in Engineering with salary above 90000, return name and salary ordered by salary DESC"
             },
             {
                 "id": "easy_02",
@@ -69,7 +80,7 @@ def list_tasks():
                 "difficulty": "easy",
                 "type": "write_query",
                 "grader": True,
-                "description": "Count employees per department ordered by count"
+                "description": "Count employees per department, return department and count ordered by count DESC"
             },
             {
                 "id": "medium_01",
@@ -77,7 +88,7 @@ def list_tasks():
                 "difficulty": "medium",
                 "type": "fix_query",
                 "grader": True,
-                "description": "Fix the broken JOIN query for active project assignments"
+                "description": "Fix the broken JOIN query with missing ON keyword and wrong table alias in WHERE"
             },
             {
                 "id": "medium_02",
@@ -85,7 +96,7 @@ def list_tasks():
                 "difficulty": "medium",
                 "type": "fix_query",
                 "grader": True,
-                "description": "Fix the GROUP BY clause to aggregate by department"
+                "description": "Fix the GROUP BY that incorrectly groups by id instead of department"
             },
             {
                 "id": "hard_01",
@@ -93,7 +104,7 @@ def list_tasks():
                 "difficulty": "hard",
                 "type": "optimize_query",
                 "grader": True,
-                "description": "Replace correlated subqueries with window functions"
+                "description": "Replace correlated subqueries with window functions or CTEs to find top earner per department"
             },
             {
                 "id": "hard_02",
@@ -101,24 +112,341 @@ def list_tasks():
                 "difficulty": "hard",
                 "type": "optimize_query",
                 "grader": True,
-                "description": "Rewrite N+1 query using LEFT JOIN and GROUP BY"
+                "description": "Rewrite N+1 correlated subquery using LEFT JOIN and GROUP BY"
             }
         ]
     }
 
 @app.post("/grade")
-def do_grade(req: GradeRequest):
-    env.reset(task_id=req.task_id)
-    step_result = env.step(req.action)
+def grade(req: GradeRequest):
+    """
+    Validator calls this to get a score for a specific task.
+    Score MUST be strictly between 0 and 1 (never 0.0, never 1.0).
+    This is enforced by _clamp() in the environment AND by the
+    absolute safety clamp below.
+    """
+    # Use a fresh environment instance for grading to avoid state pollution
+    grade_env = SQLEnvironment()
+    grade_env.reset(task_id=req.task_id)
+    result = grade_env.step(req.action)
+    
+    raw_score = result.get("reward", 0.05)
     
-    raw_score = step_result.get("reward", 0.05)
-    score = max(0.05, min(0.95, float(raw_score)))
-    score = round(score, 4)
+    # ABSOLUTE SAFETY CLAMP — even if environment has a bug,
+    # this line guarantees the validator never sees 0.0 or 1.0
+    score = round(max(0.05, min(0.95, float(raw_score))), 4)
     
     return {
         "task_id": req.task_id,
         "score": score,
-        "feedback": step_result["observation"]["feedback"],
-        "done": step_result["done"],
-        "info": step_result.get("info", {})
+        "feedback": result["observation"]["feedback"],
+        "done": result["done"],
+        "info": result.get("info", {})
     }
+
+# ── WebSocket endpoint — required by OpenEnv spec ─────────────────────────────
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """
+    Persistent WebSocket session. Each connection gets its own 
+    SQLEnvironment instance so sessions are isolated.
+    """
+    await websocket.accept()
+    ws_env = SQLEnvironment()
+    
+    try:
+        while True:
+            data = await websocket.receive_text()
+            try:
+                message = json.loads(data)
+            except json.JSONDecodeError:
+                await websocket.send_text(json.dumps({
+                    "error": "Invalid JSON"
+                }))
+                continue
+            
+            msg_type = message.get("type", "")
+            
+            if msg_type == "reset":
+                result = ws_env.reset(
+                    task_id=message.get("task_id"),
+                    difficulty=message.get("difficulty")
+                )
+                await websocket.send_text(json.dumps(result))
+            
+            elif msg_type == "step":
+                action = message.get("action", {})
+                result = ws_env.step(action)
+                await websocket.send_text(json.dumps(result))
+            
+            elif msg_type == "state":
+                await websocket.send_text(json.dumps(ws_env.state))
+            
+            else:
+                await websocket.send_text(json.dumps({
+                    "error": f"Unknown message type: {msg_type}"
+                }))
+    
+    except WebSocketDisconnect:
+        pass
+
+# ── Web UI ─────────────────────────────────────────────────────────────────────
+
+@app.get("/web", response_class=HTMLResponse)
+def web_ui():
+    """
+    Interactive web playground for testing the environment manually.
+    """
+    return """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<title>SQL data analyst playground</title>
+<style>
+  * { box-sizing: border-box; margin: 0; padding: 0; }
+  body { 
+    background: #0d1117; color: #e6edf3; 
+    font-family: 'Courier New', monospace; 
+    padding: 20px; 
+  }
+  .header { 
+    display: flex; justify-content: space-between; 
+    align-items: center; margin-bottom: 24px; 
+    border-bottom: 1px solid #30363d; padding-bottom: 16px;
+  }
+  .badge { 
+    background: #1f6feb; color: white; 
+    font-size: 11px; padding: 2px 8px; 
+    border-radius: 4px; margin-right: 8px;
+  }
+  .title { font-size: 28px; font-weight: bold; margin-top: 4px; }
+  .links a { color: #58a6ff; text-decoration: none; margin-left: 16px; }
+  .grid { display: grid; grid-template-columns: 1fr 1fr; gap: 16px; margin-bottom: 16px; }
+  .panel { 
+    background: #161b22; border: 1px solid #30363d; 
+    border-radius: 8px; padding: 16px; 
+  }
+  .panel h3 { 
+    font-size: 11px; letter-spacing: 1px; 
+    color: #8b949e; margin-bottom: 12px; 
+  }
+  .task-meta { 
+    display: flex; gap: 8px; margin-bottom: 8px; 
+  }
+  .tag { 
+    background: #21262d; border: 1px solid #30363d;
+    padding: 2px 8px; border-radius: 4px; font-size: 12px;
+  }
+  .task-desc { font-size: 14px; line-height: 1.6; color: #e6edf3; }
+  .schema-text { 
+    font-size: 12px; color: #8b949e; 
+    white-space: pre; overflow-x: auto; 
+  }
+  .sql-panel { 
+    background: #161b22; border: 1px solid #30363d; 
+    border-radius: 8px; padding: 16px; margin-bottom: 16px;
+  }
+  .sql-panel h3 { 
+    font-size: 11px; letter-spacing: 1px; 
+    color: #8b949e; margin-bottom: 4px;
+  }
+  .sql-hint { font-size: 12px; color: #8b949e; margin-bottom: 12px; }
+  textarea { 
+    width: 100%; height: 120px; 
+    background: #0d1117; color: #79c0ff;
+    border: 2px solid #238636; border-radius: 6px;
+    padding: 12px; font-family: 'Courier New', monospace;
+    font-size: 14px; resize: vertical;
+  }
+  .actions { display: flex; gap: 12px; margin-top: 12px; }
+  button { 
+    padding: 8px 20px; border-radius: 6px; 
+    border: none; cursor: pointer; 
+    font-family: 'Courier New', monospace; font-size: 14px;
+  }
+  .btn-primary { background: #238636; color: white; }
+  .btn-secondary { background: #21262d; color: #e6edf3; border: 1px solid #30363d; }
+  .btn-primary:hover { background: #2ea043; }
+  .result-panel { 
+    background: #161b22; border: 1px solid #30363d; 
+    border-radius: 8px; padding: 16px; margin-bottom: 16px;
+    min-height: 80px;
+  }
+  .result-panel h3 { 
+    font-size: 11px; letter-spacing: 1px; 
+    color: #8b949e; margin-bottom: 12px; 
+  }
+  .reward-bar { 
+    height: 6px; background: #21262d; 
+    border-radius: 3px; margin: 8px 0; 
+  }
+  .reward-fill { height: 100%; border-radius: 3px; background: #238636; transition: width 0.3s; }
+  .feedback { font-size: 13px; color: #8b949e; margin-top: 8px; }
+  .result-text { 
+    font-size: 12px; color: #e6edf3; 
+    white-space: pre; overflow-x: auto; margin-top: 8px;
+  }
+  .status { 
+    font-size: 13px; color: #8b949e; 
+    padding: 8px 0; 
+  }
+  .status.connected { color: #3fb950; }
+  .status.error { color: #f85149; }
+</style>
+</head>
+<body>
+<div class="header">
+  <div>
+    <div><span class="badge">OPENENV</span></div>
+    <div class="title">SQL data analyst playground</div>
+  </div>
+  <div class="links">
+    <a href="/docs">API docs</a>
+    <a href="/health">Health</a>
+    <a href="/tasks">Tasks</a>
+  </div>
+</div>
+
+<p style="font-size:13px;color:#8b949e;margin-bottom:20px;">
+  Interactive mode uses a persistent <code>WebSocket</code> session at <code>/ws</code> 
+  (OpenEnv HTTP <code>/step</code> is stateless). Connect, then run reset → SQL steps.
+</p>
+
+<div class="grid">
+  <div class="panel">
+    <h3>CURRENT TASK</h3>
+    <div class="task-meta">
+      <span class="tag" id="task-id">—</span>
+      <span class="tag" id="task-diff">—</span>
+    </div>
+    <div class="task-desc" id="task-desc">Click "Start episode" to load the first task.</div>
+    <div class="feedback" id="hints"></div>
+  </div>
+  <div class="panel">
+    <h3>SCHEMA</h3>
+    <div class="schema-text" id="schema-text">—</div>
+  </div>
+</div>
+
+<div class="sql-panel">
+  <h3>YOUR SQL</h3>
+  <p class="sql-hint">Type <strong>only valid SQLite</strong> here.</p>
+  <textarea id="sql-input" placeholder="Executable SQL only, e.g. SELECT department_id, AVG(salary) ..."></textarea>
+  <div class="actions">
+    <button class="btn-primary" onclick="startEpisode()">Start episode</button>
+    <button class="btn-primary" onclick="submitSQL()">Submit SQL</button>
+    <button class="btn-secondary" onclick="nextEpisode()">Next task</button>
+  </div>
+</div>
+
+<div class="result-panel">
+  <h3>LAST RESULT</h3>
+  <div id="reward-display" style="font-size:13px;color:#8b949e;">No result yet.</div>
+  <div class="reward-bar"><div class="reward-fill" id="reward-bar" style="width:0%"></div></div>
+  <div class="feedback" id="feedback-text"></div>
+  <div class="result-text" id="result-text"></div>
+</div>
+
+<div class="status" id="status">Not connected. Click "Start episode" to begin.</div>
+
+<script>
+let ws = null;
+let connected = false;
+
+function connect(callback) {
+  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
+  ws = new WebSocket(proto + '//' + location.host + '/ws');
+  
+  ws.onopen = () => {
+    connected = true;
+    document.getElementById('status').textContent = 'Connected via WebSocket.';
+    document.getElementById('status').className = 'status connected';
+    if (callback) callback();
+  };
+  
+  ws.onmessage = (event) => {
+    const data = JSON.parse(event.data);
+    handleResult(data);
+  };
+  
+  ws.onerror = () => {
+    document.getElementById('status').textContent = 'WebSocket error.';
+    document.getElementById('status').className = 'status error';
+  };
+  
+  ws.onclose = () => {
+    connected = false;
+    document.getElementById('status').textContent = 'Disconnected.';
+    document.getElementById('status').className = 'status';
+  };
+}
+
+function handleResult(data) {
+  const obs = data.observation || data;
+  if (!obs) return;
+  
+  // Update task panel
+  if (obs.task_id) document.getElementById('task-id').textContent = obs.task_id;
+  if (obs.task_type) document.getElementById('task-diff').textContent = obs.task_type;
+  if (obs.task_description) document.getElementById('task-desc').textContent = obs.task_description;
+  if (obs.expected_description) document.getElementById('hints').textContent = obs.expected_description;
+  if (obs.schema_info) {
+    const lines = obs.schema_info.split('\\n').slice(0, 20).join('\\n');
+    document.getElementById('schema-text').textContent = lines;
+  }
+  
+  // Update reward
+  const reward = data.reward || obs.reward || 0;
+  const pct = Math.round(reward * 100);
+  document.getElementById('reward-display').textContent = 'Reward: ' + reward.toFixed(4) + ' (' + pct + '%)';
+  document.getElementById('reward-bar').style.width = pct + '%';
+  
+  // Update feedback
+  if (obs.feedback) document.getElementById('feedback-text').textContent = obs.feedback;
+  if (obs.last_result) document.getElementById('result-text').textContent = obs.last_result;
+  if (obs.last_error) document.getElementById('result-text').textContent = 'ERROR: ' + obs.last_error;
+  
+  // Show starter SQL if this is a fix/optimize task
+  if (obs.last_sql && document.getElementById('sql-input').value === '') {
+    document.getElementById('sql-input').value = obs.last_sql;
+  }
+}
+
+function startEpisode() {
+  if (!connected) {
+    connect(() => {
+      ws.send(JSON.stringify({ type: 'reset' }));
+    });
+  } else {
+    ws.send(JSON.stringify({ type: 'reset' }));
+  }
+}
+
+function submitSQL() {
+  const sql = document.getElementById('sql-input').value.trim();
+  if (!sql) { alert('Please enter a SQL query.'); return; }
+  if (!connected) { alert('Click Start episode first.'); return; }
+  ws.send(JSON.stringify({ 
+    type: 'step', 
+    action: { action_type: 'write_query', sql: sql } 
+  }));
+}
+
+function nextEpisode() {
+  if (!connected) {
+    connect(() => {
+      ws.send(JSON.stringify({ type: 'reset' }));
+    });
+  } else {
+    ws.send(JSON.stringify({ type: 'reset' }));
+  }
+  document.getElementById('sql-input').value = '';
+  document.getElementById('result-text').textContent = '';
+  document.getElementById('feedback-text').textContent = '';
+}
+</script>
+</body>
+</html>
+"""

server/environment.py

@@ -1,328 +1,483 @@
-import sqlite3
-import re
-from typing import Dict, Any, List, Optional
-import json
+import sqlite3, uuid, random, re
+from typing import Optional, Tuple, List, Dict, Any
 
 SCHEMA_SQL = """
-CREATE TABLE employees (id INTEGER PRIMARY KEY, name TEXT, department TEXT, salary REAL, hire_date TEXT, manager_id INTEGER);
-CREATE TABLE departments (id INTEGER PRIMARY KEY, name TEXT, budget REAL, location TEXT);
-CREATE TABLE projects (id INTEGER PRIMARY KEY, name TEXT, department_id INTEGER, start_date TEXT, status TEXT);
-CREATE TABLE project_assignments (employee_id INTEGER, project_id INTEGER, hours_worked REAL, PRIMARY KEY(employee_id, project_id));
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department TEXT NOT NULL,
+    salary REAL NOT NULL,
+    hire_date TEXT NOT NULL,
+    manager_id INTEGER REFERENCES employees(id)
+);
+
+CREATE TABLE departments (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    budget REAL NOT NULL,
+    location TEXT NOT NULL
+);
+
+CREATE TABLE projects (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department_id INTEGER REFERENCES departments(id),
+    start_date TEXT NOT NULL,
+    status TEXT NOT NULL
+);
+
+CREATE TABLE project_assignments (
+    employee_id INTEGER REFERENCES employees(id),
+    project_id INTEGER REFERENCES projects(id),
+    hours_worked REAL NOT NULL,
+    PRIMARY KEY (employee_id, project_id)
+);
 """
 
-SEED_DATA_SQL = """
-INSERT INTO departments VALUES (1,'Engineering',500000,'New York'), (2,'Marketing',200000,'Chicago'), (3,'Sales',300000,'Los Angeles');
-INSERT INTO employees VALUES (1,'Alice Chen','Engineering',95000,'2020-01-15',NULL), (2,'Bob Smith','Engineering',85000,'2021-03-20',1), (3,'Carol Davis','Marketing',72000,'2019-07-01',NULL), (4,'David Lee','Sales',68000,'2022-11-01',NULL), (5,'Eve Turner','Engineering',110000,'2018-05-10',1), (6,'Frank White','Marketing',65000,'2023-01-15',3);
-INSERT INTO projects VALUES (1,'Data Platform',1,'2023-01-01','active'), (2,'Website Redesign',2,'2023-03-15','completed'), (3,'CRM Migration',3,'2023-06-01','active');
-INSERT INTO project_assignments VALUES (1,1,120),(2,1,80),(5,1,200), (3,2,160),(6,2,40),(4,3,100);
+SEED_SQL = """
+INSERT INTO departments VALUES
+(1,'Engineering',500000,'New York'),
+(2,'Marketing',200000,'Chicago'),
+(3,'Sales',300000,'Los Angeles');
+
+INSERT INTO employees VALUES
+(1,'Alice Chen','Engineering',95000,'2020-01-15',NULL),
+(2,'Bob Smith','Engineering',85000,'2021-03-20',1),
+(3,'Carol Davis','Marketing',72000,'2019-07-01',NULL),
+(4,'David Lee','Sales',68000,'2022-11-01',NULL),
+(5,'Eve Turner','Engineering',110000,'2018-05-10',1),
+(6,'Frank White','Marketing',65000,'2023-01-15',3);
+
+INSERT INTO projects VALUES
+(1,'Data Platform',1,'2023-01-01','active'),
+(2,'Website Redesign',2,'2023-03-15','completed'),
+(3,'CRM Migration',3,'2023-06-01','active');
+
+INSERT INTO project_assignments VALUES
+(1,1,120),(2,1,80),(5,1,200),
+(3,2,160),(6,2,40),(4,3,100);
 """
 
 TASKS = {
   "easy_01": {
     "type": "write_query",
     "difficulty": "easy",
-    "description": "Find all employees in the Engineering department with salary above 90000. Return name and salary ordered by salary descending.",
+    "description": "Find all employees in the Engineering department with salary above 90000. Return their name and salary, ordered by salary descending.",
     "expected_columns": ["name", "salary"],
     "expected_row_count": 2,
-    "hints": ["Filter department = Engineering", "salary > 90000", "ORDER BY salary DESC"]
+    "expected_rows": [
+      {"name": "Eve Turner", "salary": 110000.0},
+      {"name": "Alice Chen", "salary": 95000.0}
+    ],
+    "hints": ["Filter WHERE department = 'Engineering'", "AND salary > 90000", "ORDER BY salary DESC"]
   },
   "easy_02": {
-    "type": "write_query", 
+    "type": "write_query",
     "difficulty": "easy",
-    "description": "Count how many employees are in each department. Return department and employee count ordered by count descending.",
+    "description": "Count how many employees are in each department. Return department name and employee count, ordered by count descending.",
     "expected_columns": ["department", "count"],
     "expected_row_count": 3,
-    "hints": ["Use GROUP BY department", "Use COUNT(*)"]
+    "hints": ["Use GROUP BY department", "Use COUNT(*)", "ORDER BY count DESC"]
   },
   "medium_01": {
     "type": "fix_query",
     "difficulty": "medium",
-    "description": "Fix this broken query that returns employees on active projects with project name and hours worked.",
-    "broken_sql": "SELECT e.name, p.name, pa.hours_worked FROM employees e JOIN project_assignments pa e.id = pa.employee_id JOIN projects p ON pa.project_id = p.id WHERE projects.status = 'active'",
+    "description": "Fix this broken query that should return all employees working on active projects along with the project name and hours worked.",
+    "broken_sql": "SELECT e.name, p.name, pa.hours_worked\nFROM employees e\nJOIN project_assignments pa e.id = pa.employee_id\nJOIN projects p ON pa.project_id = p.id\nWHERE projects.status = 'active'\n",
     "expected_row_count": 3,
-    "bugs": ["Missing ON keyword in first JOIN", "Wrong alias 'projects' should be 'p' in WHERE"]
+    "bugs": ["Missing ON keyword between 'pa' and 'e.id' in first JOIN", "Wrong table reference: 'projects.status' should be 'p.status' in WHERE"]
   },
   "medium_02": {
     "type": "fix_query",
-    "difficulty": "medium", 
-    "description": "Fix this query that should return average salary per department but groups incorrectly.",
-    "broken_sql": "SELECT department, AVG(salary) FROM employees GROUP BY id ORDER BY AVG(salary) DESC",
+    "difficulty": "medium",
+    "description": "Fix this query that should return average salary per department but is producing wrong results because it groups by the wrong column.",
+    "broken_sql": "SELECT department, AVG(salary)\nFROM employees\nGROUP BY id\nORDER BY AVG(salary) DESC\n",
     "expected_row_count": 3,
-    "bugs": ["GROUP BY should use department not id"]
+    "bugs": ["GROUP BY should use 'department' not 'id'"]
   },
   "hard_01": {
     "type": "optimize_query",
     "difficulty": "hard",
-    "description": "Optimize this slow query: find top earner in each department with their total hours worked. Replace correlated subqueries with window functions or CTEs.",
-    "slow_sql": "SELECT e.name, e.department, e.salary, (SELECT SUM(hours_worked) FROM project_assignments pa WHERE pa.employee_id = e.id) as total_hours FROM employees e WHERE e.salary = (SELECT MAX(salary) FROM employees e2 WHERE e2.department = e.department) ORDER BY e.salary DESC",
+    "description": "This query works but is slow. Optimize it to find the top earner in each department along with their total hours worked across all projects. Replace correlated subqueries with window functions or CTEs.",
+    "slow_sql": "SELECT e.name, e.department, e.salary,\n       (SELECT SUM(hours_worked) \n        FROM project_assignments pa \n        WHERE pa.employee_id = e.id) as total_hours\nFROM employees e\nWHERE e.salary = (SELECT MAX(salary) \n                  FROM employees e2 \n                  WHERE e2.department = e.department)\nORDER BY e.salary DESC\n",
     "expected_row_count": 3,
-    "good_patterns": ["ROW_NUMBER", "RANK", "WITH ", "LEFT JOIN"],
-    "optimization_hints": ["Use window functions", "Use CTEs"]
+    "good_patterns": ["WITH ", "ROW_NUMBER", "RANK", "LEFT JOIN"],
+    "bad_patterns": ["SELECT.*SELECT.*SELECT"],
+    "optimization_hints": ["Use a CTE (WITH clause) to pre-aggregate hours", "Use ROW_NUMBER() or RANK() window function", "Replace nested SELECT with LEFT JOIN"]
   },
   "hard_02": {
     "type": "optimize_query",
     "difficulty": "hard",
-    "description": "Rewrite this N+1 query: return each department name, budget, employee count, and active project count using JOINs instead of subqueries.",
-    "slow_sql": "SELECT d.name, d.budget, (SELECT COUNT(*) FROM employees e WHERE e.department = d.name) as emp_count, (SELECT COUNT(*) FROM projects p WHERE p.department_id = d.id AND p.status='active') as active_projects FROM departments d",
+    "description": "Rewrite this slow N+1 query to return each department name, total budget, number of employees, and number of active projects. Use JOINs instead of correlated subqueries.",
+    "slow_sql": "SELECT d.name, d.budget,\n       (SELECT COUNT(*) FROM employees e \n        WHERE e.department = d.name) as emp_count,\n       (SELECT COUNT(*) FROM projects p \n        WHERE p.department_id = d.id \n        AND p.status='active') as active_projects\nFROM departments d\n",
     "expected_row_count": 3,
     "good_patterns": ["LEFT JOIN", "GROUP BY", "COUNT("],
-    "optimization_hints": ["Use LEFT JOIN with GROUP BY"]
+    "optimization_hints": ["Use LEFT JOIN employees ON department name", "Use LEFT JOIN projects ON department_id", "Use GROUP BY d.id or d.name", "Use COUNT() in SELECT"]
   }
 }
 
+def build_db() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    conn.row_factory = sqlite3.Row
+    conn.executescript(SCHEMA_SQL)
+    conn.executescript(SEED_SQL)
+    conn.commit()
+    return conn
+
+def get_schema_info() -> str:
+    return SCHEMA_SQL.strip()
+
+def get_sample_data() -> str:
+    conn = build_db()
+    lines = []
+    for table in ["employees", "departments", "projects", "project_assignments"]:
+        cursor = conn.execute(f"SELECT * FROM {table} LIMIT 3")
+        cols = [d[0] for d in cursor.description]
+        rows = cursor.fetchall()
+        lines.append(f"\\n-- {table} (first 3 rows)")
+        lines.append(", ".join(cols))
+        for row in rows:
+            lines.append(", ".join(str(v) for v in row))
+    conn.close()
+    return "\\n".join(lines)
+
 class SQLEnvironment:
     def __init__(self):
+        self.episode_id = ""
         self.step_count = 0
+        self.total_reward = 0.0
+        self.best_score = 0.0
         self.done = False
-        self.reward = 0.05
-        self.feedback = ""
-        self.last_sql = ""
-        self.last_result = ""
-        self.last_error = ""
-        self.task_id = "easy_01"
-        self.task = TASKS[self.task_id]
+        self.task = None
+        self.task_id = None
         self.conn = None
-        self._build_db()
-
-    def _build_db(self):
-        if self.conn:
-            self.conn.close()
-        self.conn = sqlite3.connect(":memory:")
-        self.conn.row_factory = sqlite3.Row
-        self.conn.executescript(SCHEMA_SQL)
-        self.conn.executescript(SEED_DATA_SQL)
-        self.conn.commit()
-        return self.conn
+        self.attempts = 0
 
     def _clamp(self, score: float) -> float:
-        return round(max(0.05, min(0.95, score)), 4)
+        return round(max(0.05, min(0.95, float(score))), 4)
 
-    def reset(self, task_id=None, difficulty=None):
+    def _reset_state(self):
+        self.episode_id = str(uuid.uuid4())
         self.step_count = 0
+        self.total_reward = 0.0
+        self.best_score = 0.0
         self.done = False
-        self.last_sql = ""
-        self.last_result = ""
-        self.last_error = ""
-        self.reward = self._clamp(0.15)
-        self.feedback = "Environment reset."
+        self.attempts = 0
+        if self.conn:
+            try:
+                self.conn.close()
+            except:
+                pass
+        self.conn = build_db()
 
+    def reset(self, task_id=None, difficulty=None) -> dict:
+        self._reset_state()
+        
+        available = list(TASKS.keys())
         if task_id and task_id in TASKS:
             self.task_id = task_id
-            self.task = TASKS[task_id]
         elif difficulty:
-            candidates = [k for k, v in TASKS.items() if v["difficulty"] == difficulty]
-            if candidates:
-                self.task_id = candidates[0]
-                self.task = TASKS[self.task_id]
+            pool = [k for k,v in TASKS.items() if v["difficulty"] == difficulty]
+            self.task_id = random.choice(pool) if pool else random.choice(available)
         else:
-            self.task_id = list(TASKS.keys())[0]
-            self.task = TASKS[self.task_id]
-
-        self.conn = self._build_db()
-
-        if self.task["type"] == "fix_query":
-            self.last_sql = self.task["broken_sql"]
-        elif self.task["type"] == "optimize_query":
-            self.last_sql = self.task["slow_sql"]
+            self.task_id = random.choice(available)
+            
+        self.task = TASKS[self.task_id]
+        
+        starter_sql = (self.task.get("broken_sql") or self.task.get("slow_sql") or "")
+        
+        return self._build_result(
+            reward=0.15,
+            feedback="New episode started. Study the schema and task carefully.",
+            last_sql=starter_sql.strip() if starter_sql else None,
+            last_result=None,
+            last_error=None
+        )
 
-        return self._build_observation(self.reward, self.feedback)
+    def step(self, action: dict) -> dict:
+        if self.done:
+            return self._build_result(
+                reward=self._clamp(0.0),
+                feedback="Episode complete. Call reset() to start a new episode."
+            )
+            
+        self.step_count += 1
+        self.attempts += 1
+        sql = (action.get("sql") or "").strip()
+        
+        if not sql:
+            return self._build_result(
+                reward=self._clamp(0.0),
+                feedback="Empty SQL submitted. Please write a SQL SELECT statement."
+            )
+            
+        result_text, error_text = self._execute_sql(sql)
+        
+        if error_text:
+            reward = self._clamp(0.04)
+            self.total_reward += reward
+            return self._build_result(
+                reward=reward,
+                feedback=f"SQL Error: {error_text}. Fix the syntax and try again.",
+                last_sql=sql,
+                last_result=None,
+                last_error=error_text
+            )
+            
+        task_type = self.task["type"]
+        if task_type == "write_query":
+            rows, cols = self._get_rows_cols(sql)
+            reward, feedback = self._grade_write(sql, rows, cols)
+        elif task_type == "fix_query":
+            rows, cols = self._get_rows_cols(sql)
+            reward, feedback = self._grade_fix(sql, rows, error_text)
+        elif task_type == "optimize_query":
+            rows, cols = self._get_rows_cols(sql)
+            reward, feedback = self._grade_optimize(sql, rows)
+        else:
+            reward, feedback = self._clamp(0.0), "Unknown task type."
+            
+        self.total_reward += reward
+        self.best_score = max(self.best_score, reward)
+        
+        if reward >= 0.90 or self.step_count >= 8:
+            self.done = True
+            
+        return self._build_result(
+            reward=reward,
+            feedback=feedback,
+            last_sql=sql,
+            last_result=result_text,
+            last_error=None
+        )
 
-    def _execute_sql(self, sql: str):
-        upper_sql = sql.upper()
-        blocked = ["DROP", "DELETE", "INSERT", "UPDATE", "CREATE", "ALTER", "TRUNCATE"]
-        for b in blocked:
-            if re.search(rf"\b{b}\b", upper_sql):
-                return None, "Forbidden operation"
+    def _execute_sql(self, sql: str) -> Tuple[Optional[str], Optional[str]]:
+        forbidden = ["DROP", "DELETE", "INSERT", "UPDATE", "CREATE", "ALTER", "TRUNCATE"]
+        sql_upper = sql.upper()
+        for kw in forbidden:
+            if kw in sql_upper:
+                return None, f"Forbidden: {kw} not allowed. Only SELECT."
                 
         try:
-            cursor = self.conn.cursor()
-            cursor.execute(sql)
+            cursor = self.conn.execute(sql)
+            if cursor.description is None:
+                return "Query executed (no rows returned).", None
+            cols = [d[0] for d in cursor.description]
             rows = cursor.fetchall()
             if not rows:
-                return json.dumps([{}]) if cursor.description else "[]", None
-            
-            res_dict = [dict(r) for r in rows[:10]]
-            return json.dumps(res_dict), None
-        except Exception as e:
+                return "Query ran but returned 0 rows.", None
+            lines = [", ".join(cols)]
+            for row in rows[:10]:
+                lines.append(", ".join(str(v) for v in row))
+            if len(rows) > 10:
+                lines.append(f"... ({len(rows)} total rows, showing first 10)")
+            return "\\n".join(lines), None
+        except sqlite3.Error as e:
             return None, str(e)
 
-    def _grade_write(self, sql, rows, cols) -> tuple[float, str]:
+    def _get_rows_cols(self, sql: str):
+        forbidden = ["DROP","DELETE","INSERT","UPDATE","CREATE","ALTER","TRUNCATE"]
+        if any(kw in sql.upper() for kw in forbidden):
+            return [], []
+        try:
+            cursor = self.conn.execute(sql)
+            if cursor.description is None:
+                return [], []
+            cols = [d[0] for d in cursor.description]
+            rows = cursor.fetchall()
+            return rows, cols
+        except:
+            return [], []
+
+    def _grade_write(self, sql: str, rows: list, cols: list) -> Tuple[float, str]:
         score = 0.0
-        feedback_parts = []
-        lower_sql = sql.lower()
+        parts = []
+        sql_lower = sql.lower()
         
-        if "employees" in lower_sql or "departments" in lower_sql:
+        task_tables = ["employees", "departments", "projects"]
+        if any(t in sql_lower for t in task_tables):
             score += 0.15
-            feedback_parts.append("Correct table used")
+            parts.append("Correct table used (+0.15)")
+        else:
+            parts.append("Wrong or missing table")
             
-        if "where" in lower_sql:
+        if "where" in sql_lower:
             score += 0.10
-            feedback_parts.append("WHERE clause present")
+            parts.append("WHERE clause present (+0.10)")
+        else:
+            parts.append("No WHERE clause")
             
-        expected_cols = [c.lower() for c in self.task["expected_columns"]]
+        expected_cols = [c.lower() for c in self.task.get("expected_columns", [])]
         actual_cols = [c.lower() for c in cols]
-        matched = sum(1 for ec in expected_cols if ec in actual_cols)
-        if expected_cols:
-            score += (matched / len(expected_cols)) * 0.25
-            feedback_parts.append(f"Columns: {matched}/{len(expected_cols)} matched")
+        if expected_cols and actual_cols:
+            matched = sum(1 for ec in expected_cols if any(ec in ac or ac in ec for ac in actual_cols))
+            col_score = (matched / len(expected_cols)) * 0.25
+            score += col_score
+            parts.append(f"Columns {matched}/{len(expected_cols)} matched (+{col_score:.2f})")
+        elif not expected_cols:
+            score += 0.15
+            parts.append("Column check skipped (+0.15)")
+        else:
+            parts.append("Wrong columns returned")
             
-        expected_rows = self.task["expected_row_count"]
-        if rows is not None:
-            if len(rows) == expected_rows:
-                score += 0.25
-                feedback_parts.append(f"Row count correct: {len(rows)}")
-            elif abs(len(rows) - expected_rows) == 1:
-                score += 0.12
-                feedback_parts.append(f"Row count close: {len(rows)} vs {expected_rows}")
-            else:
-                feedback_parts.append(f"Row count wrong: {len(rows)} vs {expected_rows}")
-                
-        if "order by" in lower_sql:
+        expected_count = self.task.get("expected_row_count", 0)
+        actual_count = len(rows)
+        if actual_count == expected_count:
+            score += 0.25
+            parts.append(f"Row count correct: {actual_count} (+0.25)")
+        elif abs(actual_count - expected_count) == 1:
+            score += 0.12
+            parts.append(f"Row count close: {actual_count} vs {expected_count} (+0.12)")
+        else:
+            parts.append(f"Row count wrong: {actual_count} vs {expected_count}")
+            
+        if "order by" in sql_lower:
             score += 0.15
-            feedback_parts.append("ORDER BY present")
+            parts.append("ORDER BY present (+0.15)")
+        else:
+            parts.append("No ORDER BY")
             
-        return self._clamp(score), " | ".join(feedback_parts)
+        if "desc" in sql_lower and "order by" in sql_lower:
+            score += 0.10
+            parts.append("DESC ordering (+0.10)")
+            
+        return self._clamp(score), " | ".join(parts)
 
-    def _grade_fix(self, sql, rows, error) -> tuple[float, str]:
+    def _grade_fix(self, sql: str, rows: list, error: Optional[str]) -> Tuple[float, str]:
         score = 0.0
-        feedback_parts = []
-        lower_sql = sql.lower()
+        parts = []
+        sql_lower = sql.lower()
         
         if error is None and rows is not None:
             score += 0.35
-            feedback_parts.append("Query executes successfully")
+            parts.append("Query executes successfully (+0.35)")
         else:
-            feedback_parts.append(f"Still broken: {error}")
+            parts.append(f"Still broken: {error}")
             
-        expected_rows = self.task["expected_row_count"]
-        if rows is not None and len(rows) == expected_rows:
-            score += 0.35
-            feedback_parts.append(f"Correct rows: {len(rows)}")
+        expected_count = self.task.get("expected_row_count", 0)
+        actual_count = len(rows) if rows else 0
+        if actual_count == expected_count:
+            score += 0.30
+            parts.append(f"Correct rows: {actual_count} (+0.30)")
+        elif actual_count > 0:
+            score += 0.10
+            parts.append(f"Returns data but wrong count: {actual_count} vs {expected_count} (+0.10)")
+        else:
+            parts.append("Returns no rows")
             
-        if "join" in lower_sql and " on " in lower_sql:
-            score += 0.15
-            feedback_parts.append("JOIN...ON syntax correct")
+        if "join" in sql_lower and " on " in sql_lower:
+            score += 0.20
+            parts.append("JOIN...ON syntax present (+0.20)")
+        elif "join" in sql_lower:
+            parts.append("JOIN present but missing ON keyword")
             
-        if "group by" in lower_sql or "where" in lower_sql:
+        broken_sql = self.task.get("broken_sql", "")
+        if "projects.status" in broken_sql and "projects.status" not in sql:
             score += 0.10
-            feedback_parts.append("Filter/grouping present")
-            
-        return self._clamp(score), " | ".join(feedback_parts)
+            parts.append("Fixed table alias in WHERE (+0.10)")
+        elif "group by" in sql_lower:
+            if "group by department" in sql_lower or "group by e.department" in sql_lower:
+                score += 0.10
+                parts.append("GROUP BY correct column (+0.10)")
+                
+        return self._clamp(score), " | ".join(parts)
 
-    def _grade_optimize(self, sql, rows) -> tuple[float, str]:
+    def _grade_optimize(self, sql: str, rows: list) -> Tuple[float, str]:
         score = 0.0
-        feedback_parts = []
+        parts = []
         sql_upper = sql.upper()
+        sql_lower = sql.lower()
         
-        if "WITH " in sql_upper or "ROW_NUMBER" in sql_upper or "RANK" in sql_upper:
+        uses_cte = "WITH " in sql_upper
+        uses_window = any(fn in sql_upper for fn in ["ROW_NUMBER", "RANK(", "DENSE_RANK"])
+        if uses_cte or uses_window:
             score += 0.30
-            feedback_parts.append("Uses CTE or window function")
+            label = "CTE" if uses_cte else "window function"
+            parts.append(f"Uses {label} (+0.30)")
+        else:
+            parts.append("No CTE or window function found")
             
-        if "JOIN" in sql_upper:
-            score += 0.25
-            feedback_parts.append("Uses JOIN")
+        if "join" in sql_lower:
+            score += 0.20
+            parts.append("Uses JOIN (+0.20)")
+        else:
+            parts.append("No JOIN found")
             
         select_count = sql_upper.count("SELECT")
         if select_count == 1:
             score += 0.20
-            feedback_parts.append("No nested SELECT (no correlated subqueries)")
+            parts.append("No nested SELECT — no correlated subqueries (+0.20)")
+        elif select_count == 2:
+            score += 0.10
+            parts.append(f"Reduced to {select_count} SELECTs (+0.10)")
         else:
-            feedback_parts.append(f"Still has {select_count - 1} nested SELECT(s)")
+            parts.append(f"Still has {select_count} SELECTs (correlated subqueries remain)")
             
-        expected_rows = self.task["expected_row_count"]
-        if rows is not None and len(rows) == expected_rows:
-            score += 0.20
-            feedback_parts.append(f"Correct rows: {len(rows)}")
+        expected_count = self.task.get("expected_row_count", 0)
+        actual_count = len(rows) if rows else 0
+        if actual_count == expected_count:
+            score += 0.25
+            parts.append(f"Correct row count: {actual_count} (+0.25)")
+        elif actual_count > 0:
+            score += 0.10
+            parts.append(f"Returns data, wrong count: {actual_count} vs {expected_count} (+0.10)")
+        else:
+            parts.append("Returns no rows")
             
-        return self._clamp(score), " | ".join(feedback_parts)
-
-    def get_schema_info(self) -> str:
-        return SCHEMA_SQL.strip()
-        
-    def get_sample_data(self) -> str:
-        samples = []
-        tables = ["departments", "employees", "projects", "project_assignments"]
-        cursor = self.conn.cursor()
-        for t in tables:
-            cursor.execute(f"SELECT * FROM {t} LIMIT 3")
-            rows = cursor.fetchall()
-            cols = [desc[0] for desc in cursor.description] if cursor.description else []
-            samples.append(f"Table {t}:\nCols: {', '.join(cols)}")
-            for r in rows:
-                samples.append(str(dict(r)))
-        return "\n".join(samples)
+        good_patterns = self.task.get("good_patterns", [])
+        found_good = [p for p in good_patterns if p.upper() in sql_upper]
+        if len(found_good) >= 2:
+            score += 0.05
+            parts.append(f"Good patterns found: {found_good} (+0.05)")
+            
+        return self._clamp(score), " | ".join(parts)
 
-    def _build_observation(self, reward, feedback, last_sql=None, last_result=None, last_error=None):
-        reward = self._clamp(reward)
-        obs = {
-            "task_id": self.task_id,
-            "task_type": self.task["type"],
-            "task_description": self.task["description"],
-            "schema_info": self.get_schema_info(),
-            "sample_data": self.get_sample_data(),
-            "last_sql": last_sql or self.last_sql,
-            "last_result": last_result or self.last_result,
-            "last_error": last_error or self.last_error,
+    def _build_result(self, reward, feedback, last_sql=None, 
+                      last_result=None, last_error=None) -> dict:
+        safe_reward = self._clamp(reward)
+        observation = {
+            "task_id": self.task_id or "",
+            "task_type": self.task["type"] if self.task else "",
+            "task_description": self.task["description"] if self.task else "",
+            "schema_info": get_schema_info(),
+            "sample_data": get_sample_data(),
+            "expected_description": self._get_hints(),
+            "last_sql": last_sql,
+            "last_result": last_result,
+            "last_error": last_error,
             "step_count": self.step_count,
             "done": self.done,
-            "reward": reward,
+            "reward": safe_reward,
             "feedback": feedback,
-            "metadata": self.task.get("hints", []) + self.task.get("good_patterns", [])
+            "metadata": {}
         }
-        return {"observation": obs, "reward": reward, "done": self.done, "info": {}}
+        return {
+            "observation": observation,
+            "reward": safe_reward,
+            "done": self.done,
+            "info": {
+                "episode_id": self.episode_id,
+                "total_reward": round(self.total_reward, 3),
+                "best_score": self.best_score,
+                "attempts": self.attempts
+            }
+        }
+
+    def _get_hints(self) -> str:
+        if not self.task:
+            return ""
+        hints = (self.task.get("hints") or self.task.get("optimization_hints") or self.task.get("bugs") or [])
+        return "Hints: " + "; ".join(hints) if hints else ""
 
     @property
-    def state(self):
+    def state(self) -> dict:
         return {
-            "episode_id": "ep_" + str(self.step_count),
-            "task_type": self.task["type"],
-            "task_id": self.task_id,
+            "episode_id": self.episode_id,
+            "task_type": self.task["type"] if self.task else "",
+            "task_id": self.task_id or "",
             "step_count": self.step_count,
-            "total_reward": self.reward,
-            "best_score_so_far": self.reward,
-            "attempts": self.step_count
+            "total_reward": round(self.total_reward, 3),
+            "best_score_so_far": self.best_score,
+            "attempts": self.attempts,
+            "done": self.done
         }
-
-    def step(self, action: dict):
-        self.step_count += 1
-        sql = action.get("sql", "").strip()
-        self.last_sql = sql
-        
-        if not sql:
-            self.reward = self._clamp(0.0)
-            self.last_error = "No SQL provided"
-            self.last_result = ""
-            self.feedback = "Missing SQL"
-        else:
-            res_text, err = self._execute_sql(sql)
-            if err:
-                self.last_error = err
-                self.last_result = ""
-                self.reward = self._clamp(0.08)
-                self.feedback = err
-            else:
-                self.last_error = ""
-                self.last_result = res_text
-                
-                rows = []
-                cols = []
-                if res_text and res_text != "[]" and res_text != "[{}]":
-                    try:
-                        rows = json.loads(res_text)
-                        if rows and isinstance(rows[0], dict):
-                            cols = list(rows[0].keys())
-                    except:
-                        pass
-                
-                ttype = self.task["type"]
-                if ttype == "write_query":
-                    self.reward, self.feedback = self._grade_write(sql, rows, cols)
-                elif ttype == "fix_query":
-                    self.reward, self.feedback = self._grade_fix(sql, rows, err)
-                elif ttype == "optimize_query":
-                    self.reward, self.feedback = self._grade_optimize(sql, rows)
-
-        if self.reward > 0.90 or self.step_count >= 8:
-            self.done = True
-
-        return self._build_observation(self.reward, self.feedback)