Repository navigation
feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment - #1006
OfficialAbhinavSingh wants to merge 15 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.
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
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 |
dcdc763 to
058fb4d
Compare
|
Rebased onto current main, and the Bugbot finding addressed in The finding is right, in its narrower form. So anyone reading That is the same defect #1017 fixed in I also added Worth noting for that file's stated invariant: the new class does not On the rebased head:
|
…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.
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.
`serialize_observation` keeps `reward` and `done` on the response
envelope and strips them from the nested observation payload. The client
read the envelope for `StepResult`, so `result.reward` and `result.done`
were right, but it rebuilt the observation from the nested dict alone:
envelope reward/done: 0.75 True
nested has reward? False | done? False
StepResult.reward: 0.75 | StepResult.done: True
observation.reward: None | observation.done: False
So a caller reading `result.observation.reward` saw `None` on every
step, and `result.observation.done` stayed False after the episode
ended.
This is the same defect openenv#1017 fixed in `sophistry_bench_sprint_env`, so
the fix follows it: read both from the envelope, build the observation
once with them set rather than mutating it afterwards. `reset` still
leaves `reward` as `None` instead of coercing it to 0.0, since nothing
has been scored yet.
Adds `TestSQLOptimClient` to the contract file openenv#1017 introduced, driven
by the real serializer like the others. The first test fails without the
client change; the other two pin `reset` semantics and the surrounding
fields. These tests deliberately do not skip when duckdb is absent: the
client and models import only openenv and pydantic, so they run in the
root job, which is what keeps the client and the wire format from
drifting apart again.
058fb4d to
f9b1381
Compare
|
Rebased onto current One real conflict, in While resolving it I also corrected the card's description. It still said "in-memory DuckDB", which stopped being true in No functional changes in the rebase. The earlier Bugbot finding on the client is still addressed in Checks on the rebased head:
This one has not had a human review yet. Whenever someone has time, the fork workflow runs will need the same one-off approval that unblocked #1008 before exact-head CI can report here. |
`_engine_ms` prefixes `EXPLAIN ANALYZE ` onto the agent's query to time
it. A prefix only binds to the first statement, so a rewrite of the form
`SELECT 1; SELECT * FROM events` leaves the real query running as a
statement of its own, and the plan fetch returns that result set instead
of a plan. `_MAX_MATERIALIZED_ROWS` is bypassed entirely, and the
deadline does not help because `interrupt()` cannot stop rows already
being converted into Python objects.
Measured through `compare()` against the 1M-row `events` table:
honest rewrite peak= 0.0 MB wall= 14 ms
multi-statement peak= 343.8 MB wall= 15015 ms
The 15s is the query timeout being reached while holding the shared
execution lock, so this is a memory spike and a denial of the executor
to every other session, from a single submitted rewrite.
Agent SQL is now required to be exactly one statement, decided by
`duckdb.extract_statements` rather than by scanning for `;`. That
distinction matters: a leading `--` comment, a block comment, a `;`
inside a string literal, and a trailing terminator all stay valid, which
is what the keyword pre-check removed in an earlier round got wrong.
The check lives in `_probe`, which `_run` stops on, so nothing executes
before the rejection. `explain()` gets the same guard on its own path.
The plan fetch is also bounded now, as defence in depth rather than as
the fix: a plan is a couple of rows, so anything larger is not a plan.
After the change the same attack is `peak=0.0 MB wall=6 ms`, reported as
an error. Four parametrized cases pin the legitimate rewrites that must
keep working.
|
thanks for picking this up, and for dec1148. the half-the-envelope test is a better check than what i had. on the fingerprint gate: if the probabilistic part is the concern, the large-result path could use an exact multiset check in duckdb instead, |
|
Thanks — independently reproduced the useful part on pinned DuckDB 1.5.5: the materialized-CTE symmetric There is a concrete equality-contract decision before this can be a mechanical replacement. At
DuckDB documents combination casting for set operations. Aliases, decimal scale and signed zero also differ. The existing implementation is itself inconsistent across the threshold: My recommendation is to settle one schema/type/float equality rule for both paths first (including empty results), then implement the exact multiset check with coercion guarded according to that rule. Row order is already ignored; duplicate multiplicity must remain significant. Use explicit materialization for each operand so the two EXCEPT arms share one evaluation, keep parser validation and newline-safe wrapping, retain the locked session/deadline and fail closed on comparison errors. This still compares fresh executions after timing; it does not prove equivalence of volatile queries such as No branch changes in this round: the requested simple-fix scope does not authorize choosing these reward semantics. Focused suites pass 99 tests on each of Python 3.11/3.12. The container probe also verified interruption, subsequent connection reuse, session cleanup, multi-statement rejection and denied filesystem access. Its peak across all probes/tasks was 605,560,832 bytes; that is workload evidence, not a Python heap or hostile-query bound. The process/container availability policy remains a separate maintainer gate (@Darktex). |
|
agreed on settling the rule first. proposal, one rule for every result size:
volatile queries stay out of scope. if this works for you i'll implement it on the branch with a test per rule. |
|
Thanks — technically, I accept this as a coherent comparison policy for implementation. This is acceptance of the proposed equality contract, not approval of the current implementation or merge/resource-isolation gates. I independently ran DuckDB 1.5.5 probes: 66 cases across 0/2/50,002 rows, plus million-row reorder/shift and duplicate-multiplicity cases. The schema guard rejects the earlier lossy BIGINT/DOUBLE and INTEGER/VARCHAR examples; signed zero, NaN, NULL and exact float comparison behave as proposed. Strict types/names and no float tolerance are intentional compatibility choices; the old size-dependent behavior need not be preserved. Please encode these details in implementation/tests:
The exact comparison adds engine memory/spill work; the existing 256MiB/128MiB caps and fail-closed errors remain relevant. It does not bound |
Replace the row-string comparison below 50k rows and the row-hash fingerprint above it with one exact check for every result size. Output schemas are compared by position from a raw DESCRIBE: same column count, ASCII-only case-insensitive names and identical DuckDB type strings. Only then are the rows compared as a multiset, with one materialized CTE per operand and a symmetric EXCEPT ALL, so set operation type coercion cannot make true equal 1. Only the count of differing rows returns to Python. Schema inspection and comparison share one session, lock and deadline. Expiry is checked before each stage and before a match is reported, and any error or timeout is a mismatch with no fallback. Because an interrupt sent to an idle connection is not remembered, the deadline watchdog keeps re-sending it after expiry until the context exits.
|
implemented in 4cd0937.
unrelated heads-up: locally the root test job fails collecting |
|
Independently reviewed Validation with DuckDB 1.5.5:
This accepts the comparison-policy implementation, not overall merge readiness. The independent resource-isolation gate remains: the unchanged 256MiB/128MiB caps do not bound Python conversion or total process memory; cooperative interruption is not a hard wall-time/process boundary. Volatile SQL remains outside the contract. No new full-server/container containment validation was performed. I also reproduced the reported |
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.
Reviewed by Cursor Bugbot for commit 4cd0937. Configure here.
|
checking in since it's been a few days quiet after the terminator fix. is there anything i can do on the resource-isolation side, or someone else who could weigh in if darktex isn't available right now? no rush if this needs more time, just want to make sure it's not stuck waiting on me. |
Resolve docs index conflicts with the OR-Tools environment (openenv#1501) by keeping both entries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
I merged |
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— 54 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) — 2602 passed, 117 skipped.tests/envs/test_client_step_result_contract.py— 9 passed.sql_optim_envis registered in that file, and its cases run in the root job rather than skipping, since the client and models import only openenv and pydantic.task_1_basic_antipatterns:resetloads the query and a rewrite that drops theCASTand theyear()wrapper grades atoriginal_rows=26 optimized_rows=26 match=True. SubmittingCREATE TEMP TABLE orders AS SELECT ...as a rewrite scores below an honest one, and the next episode still measures 26 rows rather than the shadowed table, which is the session isolation in9dc28bf1holding in the image.docker build -f envs/sql_optim_env/server/Dockerfile envs/sql_optim_envwith BuildKit — builds clean onghcr.io/huggingface/openenv-base:latest.uv.lockputs the Dockerfile on itsuv sync --frozenbranch, confirmed in the build log (duckdb==1.5.5, CPython 3.11.16, 111 packages)./healthreturns{"status":"healthy"}.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 is executed with substantial sandboxing, but correctness and resource limits still depend on DuckDB/process isolation for untrusted workloads; changes are additive under
envs/and do not alter core OpenEnv APIs.Overview
Introduces
sql_optim_env, a new OpenEnv package where agents rewrite slow SQL and are scored from real DuckDB runs (measured speedup plus multiset result equality), not keyword-only checks. The server seeds a temp-file database (users/orders/products/events), grades composite rewards (execution, correctness, issue detection, approval, summary, severity), and exposes five escalating tasks via the usual client/models/FastAPI/Docker/openenv.yamllayout.The DuckDB executor is the main engineering surface: read-only reopen, locked sandbox (
enable_external_access=false), per-query deadlines, short-lived cursors so temp DDL cannot poison later episodes, single-statement validation, bounded materialization for large results, and in-engine correctness viaEXCEPT ALL.SQLOptimEnv._parse_resultmirrors envelopereward/doneonto the observation so clients match other envs.CI skips the heavy suite in the root pytest job and adds
test-sql-optim-env(Python 3.11/12, frozenenvs/sql_optim_envdeps). Docs add the catalog entry, toctree page, and syncedsql_optim.md. Tests include a largetest_sql_optim_environment.pyplusTestSQLOptimClientin the shared step-result contract file.Reviewed by Cursor Bugbot for commit 4e662a5. Bugbot is set up for automated code reviews on this repo. Configure here.