feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment - #1006
feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment#1006OfficialAbhinavSingh wants to merge 6 commits into
Conversation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Pull request overview
Adds a new community environment, sql_optim_env, where agents rewrite slow SQL and are rewarded based on execution-grounded measurements from an in-memory DuckDB (speedup + result correctness), plus auxiliary rubric components (issue detection, approval, summary, severity). The PR follows the standard OpenEnv env layout (models/client/server), ships Docker + openenv.yaml, and wires the environment into the docs catalog.
Changes:
- Introduces
sql_optim_envserver runtime (tasks, DuckDB executor, graders, FastAPI app) and client/models for typed transport. - Adds pytest coverage for reset/step/state and client parsing helpers.
- Adds documentation page + environments index/toctree entries for the new environment.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/envs/test_sql_optim_environment.py | Adds unit tests for env reset/step/state and client parsing helpers. |
| envs/sql_optim_env/server/tasks.py | Defines five SQL anti-pattern tasks and metadata used by the environment. |
| envs/sql_optim_env/server/sql_optim_environment.py | Implements the server-side Environment reset/step/state logic. |
| envs/sql_optim_env/server/scoring_models.py | Adds internal Reward model returned by the grader. |
| envs/sql_optim_env/server/graders.py | Implements execution-grounded grading and score breakdown. |
| envs/sql_optim_env/server/executor.py | Implements DuckDB-backed query execution, timing, and correctness checks. |
| envs/sql_optim_env/server/Dockerfile | Provides container build/run path for the environment server. |
| envs/sql_optim_env/server/app.py | Creates the FastAPI app via create_app for WebSocket sessions. |
| envs/sql_optim_env/server/init.py | Exposes SQLOptimEnvironment from the server package. |
| envs/sql_optim_env/README.md | Environment README with quick start, schema, reward breakdown, and usage. |
| envs/sql_optim_env/pyproject.toml | Declares the env package, dependencies (incl. duckdb), and server entrypoint. |
| envs/sql_optim_env/openenv.yaml | Registers the environment spec for deployment/runtime. |
| envs/sql_optim_env/models.py | Defines wire Action/Observation/State models. |
| envs/sql_optim_env/client.py | Adds SQLOptimEnv client with payload/parse helpers. |
| envs/sql_optim_env/init.py | Package exports and top-level docs/example. |
| docs/source/environments/sql_optim.md | Adds docs page for the environment. |
| docs/source/environments.md | Adds sql_optim_env to the environments catalog page. |
| docs/source/_toctree.yml | Adds the environment docs page to the docs toctree. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Real execution comparison, carried back to the agent for its next step. | ||
| opt_q = (action.optimized_query or "").strip() | ||
| if opt_q: | ||
| try: | ||
| self._last_execution = get_executor().compare( | ||
| self._task_data["sql_query"], opt_q | ||
| ) | ||
| except Exception: | ||
| self._last_execution = None | ||
|
|
| # Grade (runs DuckDB internally). | ||
| reward = grade(self._task_data, action) | ||
| self._cumulative_reward += reward.score | ||
|
|
||
| # Real execution comparison, carried back to the agent for its next step. | ||
| opt_q = (action.optimized_query or "").strip() | ||
| if opt_q: | ||
| try: | ||
| self._last_execution = get_executor().compare( | ||
| self._task_data["sql_query"], opt_q | ||
| ) |
| BIT_XOR is commutative+associative — order-independent fingerprint. | ||
| Falls back to count-only if the DuckDB version doesn't support the function. | ||
| """ | ||
| # Try BIT_XOR of a numeric hash (portable across DuckDB versions) |
| env.reset( | ||
| task_id="task_2_join_optimization" | ||
| if "task_2_join_optimization" in TASKS | ||
| else DEFAULT_TASK_ID | ||
| ) |
73116e3 to
34c1564
Compare
|
Thanks for the thorough review — all seven were real. Addressed in Execution safety / correctness
Task definitions
Verified: 21 passed (13 existing + 8 new regression tests), |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 18 out of 18 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (1)
envs/sql_optim_env/server/executor.py:230
_checksum()executes queries on the shared DuckDB connection without any locking. Even if_run()is serialized, concurrent calls to_checksum()(from another session) can still race with_run()/compare()and corrupt results or raise driver errors.
Wrap these conn.execute(...) calls with the same lock used elsewhere so the singleton connection is always accessed serially.
try:
wrapped = sql_template.format(query=query)
result = self.conn.execute(wrapped).fetchone()
return result[0], result[1], None
except Exception:
continue
# Final fallback: count only
try:
cnt = self.conn.execute(f"SELECT COUNT(*) FROM ({query}) t").fetchone()[0]
return cnt, None, None
| _FORBIDDEN_KEYWORDS = ( | ||
| "insert", | ||
| "update", | ||
| "delete", | ||
| "drop", | ||
| "create", | ||
| "alter", | ||
| "truncate", | ||
| "replace", | ||
| "attach", |
| if isolate: | ||
| self.conn.execute("BEGIN TRANSACTION") | ||
| try: | ||
| for _ in range(runs): | ||
| try: | ||
| t0 = time.perf_counter() | ||
| rows = self.conn.execute(query).fetchall() | ||
| timings.append((time.perf_counter() - t0) * 1000.0) | ||
| except Exception as exc: | ||
| return 99_999.0, None, str(exc) | ||
| finally: | ||
| if isolate: | ||
| try: | ||
| self.conn.execute("ROLLBACK") | ||
| except Exception: | ||
| pass |
| def _run( | ||
| self, query: str, runs: int = 3, isolate: bool = False | ||
| ) -> Tuple[float, Optional[List], Optional[str]]: |
34c1564 to
55432ec
Compare
|
Addressed the second round (Bugbot) plus @copilot's remaining points in
On @copilot's note about Verified: 23 passed (13 original + 10 regression), |
55432ec to
feeb579
Compare
|
Round 3 addressed in Execution engine
Scoring
Tasks
Verified: 28 passed (13 original + 15 regression), |
|
Addressed in Added a test that valid rewrites using Verified: 28 passed, |
feeb579 to
b1776b1
Compare
|
Addressed the remaining
The finding was measurable. For
Worth flagging the tradeoff rather than burying it: small-result queries now pay one extra profiled execution, so Behaviour is otherwise unchanged: writes are still rejected by the read-only connection ( New Env suite: One note in case it comes up in a future review of this file: the over-cap probe leaves a partially-consumed result, and it deliberately isn't closed — in duckdb-python |
|
Both round-5 findings addressed in 1. Untrusted SQL could reach the filesystem — confirmed, and broader than reported. Across earlier rounds I argued that the read-only connection was "the guarantee" and that no query-level checks were needed. That holds for database writes only. It is not a sandbox. On
So: arbitrary file read and arbitrary file write in the environment server process, driven by the model's own SQL. Both connections now open with 2. Checksum fallback could fake a result match — confirmed, but not by the mechanism described. The reported mechanism was that The conclusion was right anyway, through a different door: a trailing A million rows, every one of them different, awarded full correctness credit — and since correctness gates the speedup slice (added in round 3), that unlocked speedup credit as well. Two changes: Verification. New One inference rather than a verified fact, flagged as such: I have only tested this against duckdb 1.5.5 on Linux. The option names are stable across 1.x as far as I know, but I have not verified the behaviour on another version or platform. |
|
Addressed in DuckDB 1.5.5 has no statement-timeout setting, so the fix is a Two details worth surfacing, because both are easy to get wrong: A timeout must not be treated as "try a different strategy." My first attempt armed the deadline only in The cancel race is benign only because of lock ordering. The timer can fire between a statement finishing and Measured after the change: the cross join is cancelled at the deadline and reported as One thing this round broke and I want to be explicit about. In wiring the interrupt handler into
As with the last round: verified on duckdb 1.5.5 on Linux only. The 15s limit is a judgement call rather than a measured optimum — happy to change it if you'd prefer a different ceiling, or to make it configurable. |
…izing rows _run() fetched every result row into Python via fetchall(), three times over for the timing median. For a query like SELECT * FROM events (1M rows) that held the whole result in memory and made the measurement report tuple construction rather than query execution: fetchall() took 783.6ms of which DuckDB execution was 25.5ms, so 97% of the reported time was client-side conversion. Timing now comes from DuckDB's profiler (EXPLAIN ANALYZE -> Total Time), which executes the query for real and discards the result inside the engine, so no rows cross into Python. Rows are materialized only up to the 50k cap compare() already used to pick its precise row-by-row comparison; above that the exact row count comes from an in-engine COUNT(*) and correctness from the existing order-independent checksum. A wall-clock streaming drain is kept as a fallback for a DuckDB version whose profile footer cannot be parsed, so timing degrades instead of breaking. The execution lock becomes reentrant so a whole measurement stays serialized as before. Comparing SELECT * FROM events goes from a 867.7MB peak and 4212ms reported to 14.5MB and 18.5ms; task_5_window_functions compare() drops from 13.6s to 5.7s. Small-result queries pay one extra profiled execution (task_1 101ms -> 233ms); the total across all five tasks is still ~45% lower. Behaviour is otherwise unchanged: writes are still rejected by the read-only connection, row counts stay exact, and results under the cap are still compared row by row.
…verified match credit
Two High-severity review findings on the execution engine.
The read-only connection stops the agent rewrite from changing the database,
but not from reaching the filesystem. On the previous commit an agent-authored
query could run `COPY ... TO` to write an arbitrary file (verified: the file
was created), `read_csv('/etc/passwd')` to read one, plus `read_text`,
`read_blob`, `read_parquet`, `glob`, `COPY FROM`, `EXPORT DATABASE`, `ATTACH`
and `INSTALL`/`LOAD`. Both connections now open with
`enable_external_access=false`, and with `lock_configuration=true` so a
rewrite cannot `SET` or `PRAGMA` its way back out (DuckDB additionally refuses
to re-enable external access on a running database). The escape hatches are
blocked on the `EXPLAIN ANALYZE` timing path as well.
Second, a rewrite ending in a `--` comment silently bought full correctness
credit. The checksum wraps the query in a subquery, and on a single line the
trailing comment swallowed the closing parenthesis, so the checksum failed and
equal row counts alone were reported as a match: `SELECT id FROM events -- c`
against `SELECT id + 1 AS id FROM events -- c` matched on a million rows that
all differed, unlocking the speedup slice too. `_as_subquery()` now puts the
query on its own line, and an over-cap result with no usable checksum reports
`results_match=False` instead of guessing, as does the surrounding except
path.
Note that the review's stated mechanism for the second finding -- `str.format`
raising on braces in the query -- does not reproduce: `format()` does not
rescan substituted values, and a brace-containing query checksums normally.
The trailing-comment path is the one that was exploitable.
TestSandboxAndCorrectnessCredit covers both; five of its six tests fail on the
previous commit. 38 env tests pass, 1544 repo-wide, ruff clean.
… SQL The agent chooses the SQL, so `SELECT COUNT(*) FROM events a, events b` is a trillion-row cross join. It ran unbounded on the shared read-only connection while holding the execution lock, so a single rewrite could wedge the server and queue every other session behind it, and nothing capped memory. DuckDB has no statement-timeout setting, so `_deadline()` arms a timer that calls `interrupt()` on the connection, which cancels the running query and leaves the connection usable. It is armed by every helper that executes agent SQL and around the whole of `_run`, so a helper is bounded whoever calls it while a measurement is bounded as a whole. `memory_limit` and `max_temp_directory_size` cap memory and spill space. `duckdb.InterruptException` is handled explicitly everywhere rather than falling into the generic handlers, because a timeout must not be read as "that strategy failed, try another one": `_row_count` would have answered a cancelled count by starting a streaming drain of the same runaway query with the timer already spent, and `_engine_ms` would have fallen back to wall-clock timing the same way, so the deadline bought nothing. Trusted queries (`table_stats`) are deliberately left without a deadline, since arming one there would only add a window in which a stray interrupt could cancel them. Also fixes `explain()`, which this change had broken: the success return sat after an `except` clause as dead code, so the method returned None. Nothing in the repo covered `explain()`, which is why a full green run missed it; it now has tests. Verified: the cross join is cancelled at the deadline and reported as `Query cancelled: exceeded the 15s execution limit`, the connection stays healthy, and a second session gets the lock instead of waiting indefinitely. The five real tasks are unaffected -- their slowest engine time is 392ms, and an A/B over medians shows 10011ms before against 9950ms after across all five, so the limits cost nothing measurable. TestResourceLimits covers it; the runaway test hangs until pytest-timeout kills it on the previous commit. 44 env tests pass, 1550 repo-wide, ruff clean.
The environment landed before the per-env packaging conventions this repo now uses, so on current `main` it would take the shared test job down with it: `tests/envs/test_sql_optim_environment.py` imports `duckdb` at module scope, which the root `test` job does not install, and a collection error there fails the whole job rather than one file. Follow the pattern `thinkingbox_env` established: - Ignore the SQL tests in the root `test` job, alongside the other environments whose dependencies are not in the root install, and add a `test-sql-optim-env` job that syncs `envs/sql_optim_env` and runs them for real on 3.11 and 3.12. - Add `envs/sql_optim_env/uv.lock`. The `validate-env-locks` job runs `uv sync --frozen` for every environment whose `pyproject.toml` changes, which cannot succeed without a lockfile. - Raise the `openenv` floor to `>=0.4.1` and `requires-python` to `>=3.11` so the locked resolution matches the core release and the CI matrix. - Declare the `validation:` block from RFC 008, so `openenv validate envs/sql_optim_env --level static` reports PASS. The declared reward range, memory bound and query timeout are the values the executor actually enforces. - Add `.dockerignore` so the image build does not copy `.venv` or test artifacts.
b444c00 to
7d4562e
Compare
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 7d4562e. Configure here.
A read-only DuckDB database rejects persistent writes, but it still
accepts session-scoped DDL. `CREATE TEMP TABLE users AS SELECT 1` is
therefore accepted on the shared read-only connection, and the temporary
object shadows the base table for every later query in the process:
baseline: [('real',)]
ACCEPTED: CREATE TEMP TABLE users AS SELECT 2, 'HIJACKED'
after: [('HIJACKED',)]
That is a reward-hacking path, not just leakage between episodes. An
agent can shrink `orders` to a single row, and because the original
query is measured against the same shadowed table, its rewrite comes out
both instant and "correct".
This corrects an earlier claim in this PR. Previous rounds argued the
read-only connection was a complete guarantee, that "any write (DDL/DML,
in a CTE, or after a `;`) is rejected by the DuckDB engine", and dropped
the structural pre-check on that basis. The engine rejects persistent
writes only; the comments asserting otherwise have been corrected.
Queries now run on a short-lived `self.conn.cursor()`, whose temporary
catalog is discarded when it closes, so nothing an agent creates
survives its own statement. Measured cost is 0.15 ms per query.
The deadline resolves its target when the timer fires rather than when
it is armed, so a deadline armed around a whole measurement still
cancels the individual statement running underneath it.
Four regression tests; three of them fail without this change, the
fourth asserts that DDL is still never credited. DDL previously surfaced
as an error only as a side effect of the leak (the temp table survived
into the second timing run and collided with itself), so that behaviour
is deliberately not preserved: it is unnecessary, because the grader
gates speedup credit on `results_match`. Submitting DDL as a rewrite
scores 0.13 against 0.32 for an honest no-op rewrite.
|
Rebased onto main in The Bugbot finding is real, and it contradicts what I argued earlier in this PRIn rounds 3 to 5 I claimed the read-only connection was the guarantee, that "any write (DDL/DML, in a CTE, or after a The temporary object shadows the base table for every later query on that connection, so on a long-lived server it poisons subsequent episodes. It is also a reward-hacking path rather than only leakage: shrink Queries now run on a short-lived Four regression tests, three of which fail without the change: One deliberate behaviour change: DDL submitted as a rewrite no longer surfaces as an error. It only ever did so as a side effect of the leak, since the temp table survived into the second timing run and collided with itself. Preserving that is unnecessary, because the grader gates speedup credit on Packaging, in
|
|
Container proof for the session fix, since the earlier numbers were all from the test suite. Built both images with BuildKit (
The third row is the point. On the pre-fix image every later episode in that server process grades against a one-row One thing I did not expect: measurement variance narrowed. Five honest rewrites per image, same query, The reward breakdown is identical in both (
|
Summary
Adds
sql_optim_env, an execution-grounded SQL query optimization environment. The agent receives a slow SQL query plus its schema and returns a rewritten query; reward is not a keyword heuristic — the environment executes both the original and the optimized query against a DuckDB database that is built on a temporary file and then reopened read-only, seeded with synthetic data (10k users, 500k orders, 1k products, 1M events) and scores the agent on measured speedup and result-correctness across five anti-pattern tasks (basic anti-patterns → correlated subqueries → wildcard scans → implicit joins → window-function audits).Fills a gap in the current catalog: there are coding, reasoning, and finance-QA environments, but none that grounds reward in real database execution.
Type of Change
Alignment Checklist
Before submitting, verify:
.claude/docs/PRINCIPLES.mdand this PR aligns with our principles (reward is computed inside the environment from real execution; no domain logic leaks to the trainer).claude/docs/INVARIANTS.mdand no invariants are violated (Gym-likereset/step/state; simulation controls are training-orchestration only, not agent-facing; the client never imports fromserver/)bash .claude/hooks/lint.shand tests and addressed all issuesRFC Status
Test Plan
PYTHONPATH=src:envs uv run --frozen --project envs/sql_optim_env pytest tests/envs/test_sql_optim_environment.py— 48 passed (env driven against real DuckDB: reset loads each task, step scores speedup/correctness, state round-trips; model serialization; client parse helpers; executor sandboxing, timeout and session isolation).PYTHONPATH=src:envs uv run pytest tests/ --ignore=…(the roottestjob's command) — 1760 passed, 93 skipped.SQLOptimEnv(base_url=...).sync()→reset(task_id="task_1_basic_antipatterns")loads the query,step(...)returns a graded reward (0.68),state()round-trips. Round-trip passes.docker build -f envs/sql_optim_env/server/Dockerfile .onghcr.io/huggingface/openenv-base:latestbuilt successfully when this PR was opened. Addinguv.lockswitches the Dockerfile onto itsuv sync --frozenbranch, which I could not rebuild locally (nobuildxhere, and theRUN --mount=type=cachesteps need BuildKit); both frozen sync commands and importing the server from the resulting venv were verified directly instead.python scripts/sync_env_docs.py --check,ruff check,ruff format --check— all clean.Reviewers can drive it with:
Claude Code Review
Self-check against the two-tier model: no bugs, uninitialized state, or debug code; reward stays inside the environment (
PRINCIPLES.md), and the Gymreset/step/stateboundary, agent/infra separation, and client↔server import rule (INVARIANTS.md) are all preserved. No alignment flags.Context: I built the original for the Meta PyTorch OpenEnv hackathon and adapted it to the current
envs/conventions (models/client/server,connect4_envshape). Follows on from the sync-bootstrap work in #935 / #959.Note
Medium Risk
Agent-authored SQL runs server-side with substantial sandboxing, but misconfiguration could still affect availability or resource use; changes are additive under
envs/and do not alter core OpenEnv APIs.Overview
Introduces
sql_optim_env, a new OpenEnv environment where agents rewrite slow SQL and are scored by running both queries in DuckDB on seeded synthetic tables (not keyword-only rewards). The stack follows the usual models / client / FastAPI server layout, with five graded tasks, a composite reward (speedup, correctness, issue detection, approval, summary), Docker/openenv.yaml, and user-facing docs.The execution engine is the main implementation surface: read-only DB, disabled external access, per-query deadlines, memory/spill caps, short-lived sessions so temp DDL cannot poison later episodes, and bounded measurement for large result sets (profiler timing + checksum correctness).
CI excludes
test_sql_optim_environment.pyfrom the default env test job and adds a dedicatedtest-sql-optim-envmatrix (Python 3.11/3.12) that syncsenvs/sql_optim_envand runs those tests.Reviewed by Cursor Bugbot for commit 9b048fd. Bugbot is set up for automated code reviews on this repo. Configure here.