diff --git a/AGENTS.md b/AGENTS.md index ab6571a..851f7bc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -141,6 +141,17 @@ For an LLM agent driving `atif-sql`, the discovery loop is three commands: prefer `atif-sql search 'query'` — it embeds the text first, then runs the same `semantic_search` kNN. + What the query sandbox will and won't do: `SELECT`, `EXPLAIN`, `SET + TimeZone`, and in-memory DDL/DML run; `COPY`, `EXPORT`, `ATTACH`, + `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 with kind + `sandbox_refused` before anything executes, so emit results on stdout + rather than writing files. Running as uid 0 exits 77 (`root_refused`) + unless `ATIF_SQL_ALLOW_ROOT=1`. The connection is sized to the host + before registration (`ATIF_SQL_QUERY_MEMORY_LIMIT` / `ATIF_SQL_QUERY_THREADS` + override), and no extension is ever installed at query time: `atif-sql + status` reports `vector_search` as `ready`, `no_store`, or + `extension_missing` (fix the last with `atif-sql embed --install-extension`). + Adding a view/macro? The drift tests force: a `DESCRIPTIONS` entry in `atif_duck/domain/catalog.py`, an `ARG_EXEMPLARS` entry for any new parameter name, `TABLE_MACRO_NAMES` membership if the DDL is `AS TABLE`, and diff --git a/README.md b/README.md index 6fb6966..d8ad495 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,15 @@ not in a public issue. Supported versions, the disclosure expectations, and what vulnerability in a tool that reads local transcripts are in **[SECURITY.md](SECURITY.md)**. +`atif-sql query` runs agent-composed SQL, so it runs it in a box: sized to the host before the +corpus is registered (override with `ATIF_SQL_QUERY_MEMORY_LIMIT` and `ATIF_SQL_QUERY_THREADS`), +a private spill directory outside the corpus that's removed on exit, no extension installs at +query time (`atif-sql embed --install-extension` is where the lance extension comes from), file +facing statements (`COPY`, `EXPORT`, `ATTACH`, `INSTALL`, `LOAD`, `PREPARE`, `EXECUTE`) refused +before they run, and a refusal to run as root unless `ATIF_SQL_ALLOW_ROOT=1` says so. The details +and the one accepted disclosure (`duckdb_settings()` lists the granted paths) are in +[docs/reference/cli.md](docs/reference/cli.md#query). + ## License [Apache License 2.0](LICENSE). Each of the seven module directories carries the same license diff --git a/SECURITY.md b/SECURITY.md index 08a2140..acc88e2 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -55,4 +55,12 @@ If your transcripts may not be sent to a third-party model provider, do not run DuckDB can read and write local files. The tool assumes the caller already owns the shell and the data; do not treat it as a sandbox, and do not wire it behind an interface that lets an untrusted party choose the SQL, the corpus root, or -the environment. +the environment. What the hardened connection does do is keep an injected +statement from writing the corpus or reaching the network: file-facing +statement kinds (`COPY`, `EXPORT`, `ATTACH`, `INSTALL`, `LOAD`, `PREPARE`, +`EXECUTE`) are refused before execution, the only granted directory is a +private per-process spill dir, no extension is ever installed at query time, +and `query`, `search` and `analyze` refuse to run as root (file modes don't +bind uid 0) unless `ATIF_SQL_ALLOW_ROOT=1` is set. Caller SQL can still read +`duckdb_settings()`, which lists the granted paths and so every session id; +that's the local user's own corpus listing, not a leak across a boundary. diff --git a/docs/CONTRACT.md b/docs/CONTRACT.md index e4e3dfe..adc437f 100644 --- a/docs/CONTRACT.md +++ b/docs/CONTRACT.md @@ -159,6 +159,15 @@ atif-sql materialize [--force] [--quiesce-seconds N] [--agent ...] [--workers N] # report adds rejected / rejected_session_ids atif-sql status [--agent ...] # corpus freshness, counts, watermark age atif-sql query 'SQL' [--format auto|json|csv] + # sized to the host before registration + # (ATIF_SQL_QUERY_MEMORY_LIMIT / _THREADS + # override); private mkdtemp spill dir; + # refuses uid 0 unless ATIF_SQL_ALLOW_ROOT=1; + # COPY/EXPORT/ATTACH/INSTALL/LOAD/PREPARE/EXECUTE + # exit 70 sandbox_refused before execution; + # never installs an extension +atif-sql embed --install-extension # the ONE place the lance DuckDB extension + # is downloaded (also done by a real embed run) atif-sql schema # static, <50ms, no duckdb bind ## Parity oracle (satisfied and retired) @@ -181,6 +190,9 @@ repo neither declares nor provides. source_root (default CLAUDE_CONFIG_DIR~/.claude /projects), corpus_root, quiesce_seconds=300, agent=claude-code (every --agent command reads it), materialize_workers=min(8, cpu_count) (materialize's pool size; 1 = single process). +ATIF_SQL_QUERY_MEMORY_LIMIT (DuckDB size literal, e.g. 6GB) and ATIF_SQL_QUERY_THREADS +override query's host-derived cap and thread count; ATIF_SQL_ALLOW_ROOT=1 lets +query/search/analyze run as uid 0 (a warning is logged). _default_*() factories read env at call time. With agent=codex the two roots re-derive to $CODEX_HOME (default ~/.codex)/sessions and ~/.atif-sql/corpus/codex; an explicitly set ATIF_SQL_SOURCE_ROOT or ATIF_SQL_CORPUS_ROOT always wins over that diff --git a/docs/architecture/data-flow.md b/docs/architecture/data-flow.md index f077c9a..f870706 100644 --- a/docs/architecture/data-flow.md +++ b/docs/architecture/data-flow.md @@ -103,12 +103,19 @@ sequenceDiagram written by a different provider or width raises instead of binding, because vectors from different models live in incompatible spaces and would return numerically valid but meaningless cosine scores — `packages/atif-duck/src/atif_duck/domain/embedding_guard.py:46`. -7. The fully-registered connection is then sandboxed: spill directory, memory cap, a directory - allowlist holding only the spill area, a path allowlist holding the individual analytics - parquets, `enable_external_access=false`, and `lock_configuration` last so caller SQL cannot - widen any of it — `packages/atif-cli/src/atif_cli/app.py:138`. -8. The caller's statement executes against the locked connection (`:606`) and the cursor drains in - batches to stdout — a JSON array of row objects on a pipe, a width-aligned table on a TTY — `packages/atif-cli/src/atif_cli/output.py:154`. +7. Before any of that, the connection was sized to the host (`_configure_query_resources`: a memory + cap from available RAM, a thread count from that cap, a private `mkdtemp` spill directory, and + extension auto-install and auto-load off), because registration is what needs the cap. The + fully-registered connection is then sandboxed: a directory allowlist holding only the private + spill area, a path allowlist holding the individual parquets the views read lazily, + `enable_external_access=false`, and `lock_configuration` last so caller SQL cannot widen any of + it (`packages/atif-cli/src/atif_cli/app.py`, `_harden_query_connection`). +8. The statement's kinds are checked with DuckDB's own parser on the locked connection; `COPY`, + `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 + (`sandbox_refused`) before anything runs. Then the caller's statement executes and the cursor + drains in batches to stdout: a JSON array of row objects on a pipe, a width-aligned table on a + TTY (`packages/atif-cli/src/atif_cli/output.py:154`). The spill directory is removed when the + process exits. ```mermaid sequenceDiagram @@ -119,6 +126,7 @@ sequenceDiagram participant Lance as Lance store CLI->>DB: duckdb.connect() + CLI->>DB: threads, memory_limit, private temp_directory, autoinstall off CLI->>Duck: register(con, corpus_root, expected model + dim) Duck->>DB: CREATE TEMP TABLE raw readers over corpus globs DB->>Disk: read_json sessions/*/meta.json then the rest @@ -130,6 +138,7 @@ sequenceDiagram Duck->>Duck: ensure_store_matches, then bind the view Duck-->>CLI: views and macros registered CLI->>DB: allowlists, external access off, lock_configuration + CLI->>DB: extract_statements(caller SQL), refuse file-facing kinds CLI->>DB: execute(caller SQL) DB-->>CLI: cursor CLI->>CLI: drain in batches to stdout diff --git a/docs/behavior/processes.md b/docs/behavior/processes.md index 7f7eb40..2b814eb 100644 --- a/docs/behavior/processes.md +++ b/docs/behavior/processes.md @@ -270,14 +270,18 @@ Entry point: `packages/atif-cli/src/atif_cli/app.py:529` 5. Create the macros, then the analytics views and analytics macros over the analytics parquets, which bind against both the parquets and the base views — `:1005`, `packages/atif-duck/src/atif_duck/infrastructure/analytics.py:124`. -6. Harden the connection in a fixed order: temp directory, memory cap, a - directory allowlist holding only the spill area, a file allowlist of the - analytics parquets, the config exemption list, then - `enable_external_access=false` and `lock_configuration=true` last — - `packages/atif-cli/src/atif_cli/app.py:138`. -7. Execute the caller's statement and stream the cursor: a plain table on a - TTY, a JSON array of row objects on a pipe — `:606`, - `packages/atif-cli/src/atif_cli/output.py:154`. +6. Harden the connection in a fixed order (the memory cap, thread count and + private spill directory were set before step 3, since registration is what + needs them): a directory allowlist holding only the spill area, a file + allowlist of the parquets the views read lazily, the config exemption + list, then `enable_external_access=false` and `lock_configuration=true` + last (`packages/atif-cli/src/atif_cli/app.py`, `_harden_query_connection`). +7. Check the statement's kinds with DuckDB's parser and refuse `COPY`, + `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` + before anything runs (exit 70, `sandbox_refused`); then execute the + caller's statement and stream the cursor: a plain table on a TTY, a JSON + array of row objects on a pipe (`packages/atif-cli/src/atif_cli/output.py:154`). + The spill directory is removed on exit. 8. Classify any DuckDB failure into parse, catalog, or runtime — or an embedding-provider mismatch — and exit 64, 65, or 70 with a JSON error envelope — `packages/atif-cli/src/atif_cli/duck_errors.py:34`, `:66`. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 5c40cdb..d484ecf 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -100,13 +100,17 @@ Flags: - `--corpus-root` — override the materialized corpus root. `:504` - `--format` — `table` on a TTY, a JSON array of row objects on a pipe. `:505` -The statement runs against a hardened connection: reads reach the registered views and nothing else, and the only writable path is the query engine's own spill directory `/.duckdb_tmp`. `:601` +The statement runs against a hardened connection: reads reach the registered views and nothing else, and nothing under the corpus root is writable. Before registration the connection is sized to the host: a memory cap derived from available RAM (half of physical RAM or 8 GiB, whichever is larger, never above 80% of what's available) and a thread count of one per 2 GiB of that cap, capped at the CPUs the process may use. `ATIF_SQL_QUERY_MEMORY_LIMIT` (a DuckDB size such as `6GB`) and `ATIF_SQL_QUERY_THREADS` override both; a malformed value exits 64. The spill directory is a private `mkdtemp` (mode 0700) under the system temp dir, the only directory the sandbox grants, and it's removed when the process exits. Extension auto-install and auto-load are off, and the lance extension is loaded only when it's already installed, so registration never reaches the network (`packages/atif-cli/src/atif_cli/app.py`, `_configure_query_resources`). + +Two layers keep caller SQL from writing the corpus. DuckDB's file grants are read-write and it has no read-only grant, so `COPY ... TO (USE_TMP_FILE false)` would overwrite one; the CLI therefore refuses every statement kind that names a file before executing anything, using DuckDB's own parser: `COPY`, `EXPORT`, `ATTACH`, `DETACH`, `INSTALL`, `LOAD`, `PREPARE` and `EXECUTE` exit 70 with kind `sandbox_refused`, for any uid, and a batch containing one of them runs nothing. And because a `0444` file mode doesn't bind root, `query` refuses to run as uid 0 (exit 77, kind `root_refused`) unless `ATIF_SQL_ALLOW_ROOT=1` is set, which logs a warning. `search` and `analyze` refuse root the same way. + +What caller SQL can still see: `duckdb_settings()` and `current_setting(...)` return the sandbox's own configuration, including the corpus root, the spill directory, the memory cap and every granted parquet path, which names every session id. DuckDB can't hide a setting from SQL and the grants have to be per file, so this is accepted: the caller is the local user, who can list the corpus and `SELECT session_id FROM sessions` anyway, and `query` isn't a privilege boundary (see SECURITY.md). The views themselves carry no corpus path as statement text. The registry hands its globs and file lists to `read_json(?)` as bound parameters and builds the parquet readers through DuckDB's relation API, so a corpus root such as `o'brien ?; --$1` and transcript content carrying SQL text both register as data (`packages/atif-duck/src/atif_duck/infrastructure/registry.py`). A session directory whose name fails the session id boundary (`packages/atif-duck/src/atif_duck/domain/session_id.py`) registers nothing and is logged once. Sessions that carry current columnar artifacts are served from their parquet files, so no JSON is parsed for them at query time; the rest are read from `trajectory.json`, and the views union the two. The per-session parquet files the registry bound are granted to the sandbox the same way the analytics parquets are (as individual `allowed_paths` entries, `packages/atif-cli/src/atif_cli/app.py:221`), and they're written read-only (`0444`), so a `COPY ... TO` at one of them fails at the filesystem even though DuckDB's grant is read-write. `atif-sql status` says which path a corpus takes. -Exit codes: `64` parse error, `65` catalog error, `65` embedding mismatch, `70` runtime error. `:550` +Exit codes: `64` parse error or a malformed `ATIF_SQL_QUERY_*` override, `65` catalog error, `65` embedding mismatch, `70` runtime error, `70` `sandbox_refused` (a statement kind the sandbox never runs), `77` `root_refused`. `:550` ## analyze diff --git a/packages/atif-cli/src/atif_cli/app.py b/packages/atif-cli/src/atif_cli/app.py index 79c99a4..457d135 100644 --- a/packages/atif-cli/src/atif_cli/app.py +++ b/packages/atif-cli/src/atif_cli/app.py @@ -54,7 +54,12 @@ import json import os +import re +import shutil import sys +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Annotated, Any @@ -72,7 +77,7 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Generator, Sequence from atif_corpus.application.materialize import MaterializationReport @@ -121,25 +126,269 @@ def _sql_str(value: str) -> str: return f"'{escaped}'" -#: Lower bound for ``query``'s DuckDB memory cap. Registration alone parses -#: the corpus through 1 GiB-per-object JSON buffers on every thread, and the -#: base views (``tool_calls``, ``tool_rank``) need several GiB more on a -#: multi-GB corpus — a tighter cap turns working queries into OutOfMemory. -_QUERY_MEMORY_FLOOR_BYTES = 8 * 1024**3 +#: Env override for the DuckDB memory cap ``query`` runs under: any size +#: literal DuckDB's ``SET memory_limit`` accepts (``6GB``, ``512MiB``, ``2000000000B``). +QUERY_MEMORY_LIMIT_ENV = "ATIF_SQL_QUERY_MEMORY_LIMIT" + +#: Env override for the DuckDB thread count ``query`` runs under. +QUERY_THREADS_ENV = "ATIF_SQL_QUERY_THREADS" + +#: Set to ``1`` to let ``query``, ``search`` and ``analyze`` run as uid 0. +ALLOW_ROOT_ENV = "ATIF_SQL_ALLOW_ROOT" + +#: Target for ``query``'s DuckDB memory cap when the host can afford it: the +#: base views (``tool_calls``, ``tool_rank``) need several GiB on a multi-GB +#: corpus, and a tighter cap turns working queries into OutOfMemory. A +#: target, not a floor: it never exceeds what the host has (the old 8 GiB +#: floor set a "limit" above physical RAM on an 8 GiB guest). +_QUERY_MEMORY_TARGET_BYTES = 8 * 1024**3 + +#: Never cap below this. DuckDB needs room to open the readers at all, and a +#: host with less available than this fails to register whatever the cap says. +_QUERY_MEMORY_MIN_BYTES = 512 * 1024**2 + +#: Cap budgeted per DuckDB thread when deriving the thread count. The JSON +#: readers reserve about twice their ``maximum_object_size`` per thread, and +#: that bound is now sized from the largest file present (the largest +#: trajectory.json seen is 436 MB), so 2 GiB a thread keeps a corpus on the +#: JSON path registering under the cap. +_QUERY_BYTES_PER_THREAD = 2 * 1024**3 + + +def _host_memory() -> tuple[int, int]: + """``(physical, available)`` bytes on this host. + + Physical comes from ``os.sysconf``. Available is Linux's ``MemAvailable`` + from ``/proc/meminfo`` (what a new allocation can really get, reclaimable + page cache included); where that file is absent (macOS) available is + taken as physical. + """ + physical = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") + available = physical + try: + with Path("/proc/meminfo").open(encoding="ascii") as handle: + for line in handle: + if line.startswith("MemAvailable:"): + available = int(line.split()[1]) * 1024 + break + except (OSError, ValueError, IndexError): + pass + return physical, min(physical, available) def _query_memory_limit_bytes() -> int: - """Bytes to cap ``query``'s DuckDB heap at. + """Bytes to cap ``query``'s DuckDB heap at, derived from the host. + + Half of physical RAM or :data:`_QUERY_MEMORY_TARGET_BYTES`, whichever is + larger, but never above 80% of physical RAM (DuckDB's own default) and + never above 80% of the memory available right now, floored at + :data:`_QUERY_MEMORY_MIN_BYTES`. On a 124 GiB host that is 62 GiB; on an + 8 GiB guest with 6 GiB free it is 4.8 GiB, where the previous rule said + 8 GiB and DuckDB's default said 6.4 GiB, neither of which the guest had. + """ + physical, available = _host_memory() + target = max(int(physical * 0.5), _QUERY_MEMORY_TARGET_BYTES) + ceiling = min(int(physical * 0.8), int(available * 0.8)) + return max(_QUERY_MEMORY_MIN_BYTES, min(target, ceiling)) + + +def _query_threads(memory_limit_bytes: int) -> int: + """DuckDB threads for ``query``: one per :data:`_QUERY_BYTES_PER_THREAD` of cap. - Half of physical RAM, but never below :data:`_QUERY_MEMORY_FLOOR_BYTES` - and never above DuckDB's own 80%-of-RAM default — on a small host the - default is already the tighter of the two, and raising it would make an - unbounded query worse rather than better. + Capped at the CPUs this process may run on (``sched_getaffinity``, which + sees a container's cpuset where ``cpu_count`` does not) and floored at + one. DuckDB's own default is the core count, which is what made a 4 vCPU + guest reserve four readers' worth of memory it did not have. """ - import os + affinity = getattr(os, "sched_getaffinity", None) + cpus = len(affinity(0)) if affinity is not None else (os.cpu_count() or 1) + return max(1, min(cpus, memory_limit_bytes // _QUERY_BYTES_PER_THREAD)) + + +_SIZE_UNITS: dict[str, int] = { + "B": 1, + "KB": 10**3, + "MB": 10**6, + "GB": 10**9, + "TB": 10**12, + "KIB": 1024, + "MIB": 1024**2, + "GIB": 1024**3, + "TIB": 1024**4, +} + +_SIZE_RE = re.compile(r"\s*(\d+(?:\.\d+)?)\s*([A-Za-z]*)\s*") + + +def _parse_size(text: str) -> int: + """Bytes for a DuckDB-style size literal (``6GB``, ``512 MiB``, ``123B``).""" + match = _SIZE_RE.fullmatch(text) + if match is None: + problem = f"{text!r} is not a size (expected e.g. 6GB, 512MiB, 2000000000B)" + raise ValueError(problem) + number, unit = match.group(1), (match.group(2) or "B").upper() + if unit not in _SIZE_UNITS: + problem = f"{text!r} has an unknown unit (use B, KB, MB, GB, TB, KiB, MiB, GiB, TiB)" + raise ValueError(problem) + return int(float(number) * _SIZE_UNITS[unit]) + + +@dataclass(frozen=True, slots=True) +class QueryResources: + """What ``query`` hands DuckDB before it registers anything.""" + + memory_limit_bytes: int + threads: int + + +def _query_resources() -> QueryResources: + """Resolve the memory cap and thread count for one ``query`` process. + + :data:`QUERY_MEMORY_LIMIT_ENV` and :data:`QUERY_THREADS_ENV` override the + host-derived values; the thread default follows whichever cap is in + force, so a caller who lowers the cap gets fewer threads for free. + Raises ``ValueError`` on a malformed override; the command turns that + into exit 64. + """ + memory_env = os.environ.get(QUERY_MEMORY_LIMIT_ENV, "").strip() + threads_env = os.environ.get(QUERY_THREADS_ENV, "").strip() + if memory_env: + memory = _parse_size(memory_env) + if memory <= 0: + problem = f"{QUERY_MEMORY_LIMIT_ENV} must be a positive size" + raise ValueError(problem) + else: + memory = _query_memory_limit_bytes() + if threads_env: + threads = int(threads_env) + if threads < 1: + problem = f"{QUERY_THREADS_ENV} must be a positive integer" + raise ValueError(problem) + else: + threads = _query_threads(memory) + return QueryResources(memory_limit_bytes=memory, threads=threads) - physical = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") - return min(int(physical * 0.8), max(int(physical * 0.5), _QUERY_MEMORY_FLOOR_BYTES)) + +def _configure_query_resources(con: Any, resources: QueryResources, spill_dir: Path) -> None: + """Size the connection to the host BEFORE registration. + + Registration is the heaviest thing ``query`` does (the eager JSON readers + reserve memory per thread), so the cap and thread count must be in force + before it runs, not after, which is where they used to be applied. The + spill directory is set here for the same reason: a registration that + spills must spill into the private directory. Extension auto-install and + auto-load are switched off so nothing during registration can reach the + network; the lance extension is LOADed only when it is already installed + (:func:`atif_duck.infrastructure.registry.load_lance_extension`). + """ + con.execute(f"SET threads={int(resources.threads)}") + con.execute(f"SET memory_limit='{int(resources.memory_limit_bytes)}B'") + con.execute(f"SET temp_directory={_sql_str(str(spill_dir))}") + con.execute("SET autoinstall_known_extensions=false") + con.execute("SET autoload_known_extensions=false") + + +@contextmanager +def _private_spill_dir() -> Generator[Path]: + """A per-process spill directory outside the corpus, gone when the process is. + + ``mkdtemp`` creates it mode 0700 under the system temp dir. It is the + ONLY directory the sandbox grants, so it is where DuckDB spills a query + that exceeds ``memory_limit``; it used to be ``/.duckdb_tmp``, + inside the tree ``materialize`` scans, and files written there by caller + SQL persisted between runs. Removed on every exit path, error included, + after the connection is closed. + """ + path = Path(tempfile.mkdtemp(prefix="atif-sql-query-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors=True) + + +def _refuse_root(command: str, fmt: OutputFormat) -> None: + """Exit 77 when running as uid 0, unless :data:`ALLOW_ROOT_ENV` is ``1``. + + The corpus artifacts are written mode 0444, which is what stops a + ``COPY ... TO`` at a granted parquet from succeeding at the filesystem. + Root ignores file modes, so as uid 0 that protection is gone and caller + SQL could overwrite the corpus (measured: probe d2 of the MicroVM review + rewrote ``steps.parquet`` as a one-row file). The override exists for + containers that have no other user; it logs a warning so the run is on + record. + """ + geteuid = getattr(os, "geteuid", None) + if geteuid is None or geteuid() != 0: + return + if os.environ.get(ALLOW_ROOT_ENV, "").strip() == "1": + from loguru import logger + + logger.warning( + "atif-sql {} is running as root because {}=1; the 0444 file modes no longer " + "protect the corpus from the SQL this process runs", + command, + ALLOW_ROOT_ENV, + ) + return + err = ClassifiedError( + kind="root_refused", + exit_code=EXIT_CODES["root_refused"], + message=f"atif-sql {command} refuses to run as root (uid 0)", + hint=f"run it as an unprivileged user, or set {ALLOW_ROOT_ENV}=1 to override " + "(a warning is logged; root can then overwrite corpus files from SQL)", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) + + +#: Statement kinds the query sandbox executes. Everything else is refused +#: before execution, by name, using DuckDB's own parser on the hardened +#: connection. The refused kinds are the ones that name a file or defer a +#: statement past this check: COPY and COPY_DATABASE (DuckDB's ``allowed_paths`` +#: grants are read-write, so ``COPY ... TO (USE_TMP_FILE +#: false)`` would overwrite corpus data, and root ignores the 0444 mode that +#: used to stop it), EXPORT, ATTACH, DETACH, LOAD, EXTENSION (INSTALL), +#: PREPARE and EXECUTE (a prepared COPY parses as PREPARE). An allowlist, so a +#: statement kind a future DuckDB adds is refused until someone reads what it +#: does. +_QUERY_STATEMENT_KINDS: frozenset[str] = frozenset( + { + "SELECT", + "EXPLAIN", + "SET", + "VARIABLE_SET", + "CREATE", + "CREATE_FUNC", + "DROP", + "ALTER", + "INSERT", + "UPDATE", + "DELETE", + "MERGE_INTO", + "CALL", + "PRAGMA", + "TRANSACTION", + "VACUUM", + "ANALYZE", + } +) + + +def _refused_statement_kinds(con: Any, sql: str) -> list[str]: + """Statement kinds in ``sql`` outside :data:`_QUERY_STATEMENT_KINDS`, in order, once each. + + ``extract_statements`` is the parser ``execute`` uses, so there is no + second grammar to disagree with it, and it runs on the hardened + connection so a PRAGMA that reads a file while parsing meets the same + allowlist the query would. A parse error propagates as DuckDB's own + ``ParserException`` (exit 64), exactly as ``execute`` would have raised it. + """ + refused: list[str] = [] + for statement in con.extract_statements(sql): + kind = str(statement.type.name) + if kind not in _QUERY_STATEMENT_KINDS and kind not in refused: + refused.append(kind) + return refused #: Config options caller SQL may still change after the sandbox locks. Only @@ -169,7 +418,7 @@ def _harden_query_connection( con: Any, *, corpus_root: Path, - temp_dir: Path, + spill_dir: Path, columnar_paths: Sequence[Path] = (), ) -> None: """Sandbox a fully-registered connection before it runs caller SQL. @@ -180,25 +429,35 @@ def _harden_query_connection( path, ``ATTACH`` of unrelated databases, and ``INSTALL httpfs`` for network egress. + The memory cap, thread count and spill directory are NOT set here any + more: :func:`_configure_query_resources` applies them before + ``register`` runs, because registration is what needs them. This + function arms the allowlists and freezes the configuration. + Reach is granted at two granularities, and the split is what keeps injected SQL from overwriting the corpus it was read from: - * ``allowed_directories`` holds ONLY ``temp_dir`` — the spill area, which - must accept writes for a query that exceeds ``memory_limit`` to - complete at all. It sits inside the corpus root but holds no corpus - data. + * ``allowed_directories`` holds ONLY ``spill_dir``, the per-process + private directory from :func:`_private_spill_dir`, which must accept + writes for a query that exceeds ``memory_limit`` to complete at all. + It lives outside the corpus and is removed when the process exits, so + nothing written there persists and nothing under the corpus root is + writable. (DuckDB grants a directory read-write, so a ``COPY`` into it + would succeed at this layer; the statement gate in :func:`query` + refuses COPY before it gets here.) * ``allowed_paths`` holds the individual files the views read lazily: the analytics parquets from :func:`_lazy_read_paths` and the per-session columnar parquets in ``columnar_paths`` (what ``register`` bound this connection to). A file grant lets - ``read_parquet`` open that exact path; because ``COPY`` writes through - a sibling ``tmp_`` file, which is a different path, the grant - does not carry a plain write. Nothing else under the corpus is named, - so ``COPY`` over a ``trajectory.json`` or into a new ``pwned.csv`` is - refused. The columnar parquets are additionally written read-only - (``0o444``) by their producer, which closes the ``USE_TMP_FILE false`` - residual hole documented on :func:`query` for them: the open for - writing fails at the filesystem for any non-root user. + ``read_parquet`` open that exact path. Nothing else under the corpus + is named, so ``read_text`` of a ``trajectory.json`` is refused. DuckDB + 1.5.5 has no read-only grant (measured: a ``read_parquet`` relation + bound before ``enable_external_access=false`` is refused at query time + without a grant, and ``duckdb_settings()`` lists no write switch), so + each grant is read-write and ``COPY ... TO + (USE_TMP_FILE false)`` would overwrite it. Two things close that: the + statement gate refuses COPY for any uid, and :func:`_refuse_root` + keeps the producer's 0444 modes meaningful by refusing uid 0. Four ordering constraints, each one required: @@ -214,9 +473,7 @@ def _harden_query_connection( including itself). Without it the memory cap is decorative: caller SQL can just ``SET memory_limit`` back up. """ - con.execute(f"SET temp_directory={_sql_str(str(temp_dir))}") - con.execute(f"SET memory_limit='{_query_memory_limit_bytes()}B'") - con.execute(f"SET allowed_directories=[{_sql_str(str(temp_dir))}]") + con.execute(f"SET allowed_directories=[{_sql_str(str(spill_dir))}]") lazy_paths = ", ".join( _sql_str(str(path)) for path in (*columnar_paths, *_lazy_read_paths(corpus_root)) ) @@ -612,6 +869,44 @@ def materialize( # --------------------------------------------------------------------------- +def _vector_surface(corpus_root: Path) -> dict[str, Any]: + """How ``query`` and ``search`` will see the embeddings store, without installing anything. + + The extension check reads ``duckdb_extensions()`` on a throwaway + connection (the local extension directory, a few milliseconds, no + network); the store check is the directory ``embed`` writes, resolved the + way ``query`` and ``search`` resolve it (``ATIF_SQL_LANCE_URI`` wins). + """ + import duckdb + + from atif_duck.infrastructure.registry import lance_extension_installed + from atif_embed.infrastructure.settings import EmbedSettings + + store = EmbedSettings().resolve_lance_uri(corpus_root) + con = duckdb.connect() + try: + installed = lance_extension_installed(con) + finally: + con.close() + store_present = store.is_dir() + if store_present and installed: + state, note = "ready", f"store at {store}" + elif store_present: + state = "extension_missing" + note = ( + "store present but the lance DuckDB extension is not installed; " + "run `atif-sql embed --install-extension`" + ) + else: + state, note = "no_store", "no embeddings store; run `atif-sql embed --all --no-dry-run`" + return { + "lance_extension_installed": installed, + "embeddings_store_present": store_present, + "vector_search": state, + "note": note, + } + + def _dir_bytes(root: Path) -> int: """Total bytes of every file under ``root`` (0 if absent).""" if not root.is_dir(): @@ -647,6 +942,14 @@ def status( ``materialize --force`` brings them over). It applies the same per-session predicate the registry does, so it cannot say ``columnar`` while ``query`` silently parses JSON. + + ``vector search`` says whether ``semantic_search`` and ``search`` can + reach an embeddings store: ``ready`` (store present, lance extension + installed), ``no_store`` (run ``atif-sql embed``), or + ``extension_missing`` (a store exists but the DuckDB lance extension is + not installed; ``query`` binds ``message_embeddings`` empty rather than + downloading it, so run ``atif-sql embed --install-extension``). Read + from DuckDB's local extension listing; never installs anything. """ import time @@ -687,6 +990,7 @@ def status( "live": len(plan.skipped_live), } coverage = columnar_coverage(settings.corpus_root) + vector = _vector_surface(settings.corpus_root) if resolve_format(fmt) is OutputFormat.TABLE: print(f"agent: {settings.agent.value}") print(f"source root: {settings.source_root}") @@ -706,6 +1010,7 @@ def status( f"({coverage.columnar_sessions} of {coverage.total_sessions} complete sessions " "carry typed columnar artifacts)" ) + print(f"vector search: {vector['vector_search']} ({vector['note']})") else: emit_json( { @@ -721,6 +1026,9 @@ def status( "columnar_sessions": coverage.columnar_sessions, "json_sessions": coverage.json_sessions, "query_path": coverage.query_path, + "lance_extension_installed": vector["lance_extension_installed"], + "embeddings_store_present": vector["embeddings_store_present"], + "vector_search": vector["vector_search"], }, fmt, ) @@ -755,28 +1063,55 @@ def query( Sandbox ------- - The statement runs against a hardened connection (see + Before anything is registered the connection is sized to the host + (:func:`_query_resources`: a memory cap derived from available RAM and a + thread count derived from that cap, overridable with + ``ATIF_SQL_QUERY_MEMORY_LIMIT`` and ``ATIF_SQL_QUERY_THREADS``), its spill + directory is a private ``mkdtemp`` (mode 0700) outside the corpus that is + removed when the process exits, and extension auto-install and auto-load + are off. The lance extension is loaded only when it is already installed; + when a store exists and it is not, the vector surface binds empty with a + warning and ``atif-sql status`` says so. Registration therefore reaches + neither the network nor the corpus for writing, and it fits an 8 GiB + guest, where it used to exit 70 out of memory before the cap applied. + + The statement then runs against the hardened connection (see :func:`_harden_query_connection`). Reads reach the registered views and - nothing else; the only writable path is the query engine's own spill - directory ``/.duckdb_tmp``. ``read_text`` outside the - corpus, ``COPY`` anywhere in it (including over a ``trajectory.json``), - ``ATTACH`` of unrelated databases, and extension installs all fail with - exit 70 rather than reading credentials, corrupting the corpus, or - reaching the network. The embedding-store guard binds here exactly as it - does for ``search``, so a store written by a different provider refuses - to bind (exit 65) instead of scoring garbage. - - Two residual holes remain, both over recomputable derived data rather - than the transcript artifacts: - - * ``COPY ... TO '' (USE_TMP_FILE false)`` - overwrites it. A plain ``COPY`` to the same path is refused, because - DuckDB stages it through a sibling ``tmp_`` that carries no - grant; ``USE_TMP_FILE false`` writes the granted path directly. Re-run - ``atif-sql analyze`` to rebuild. - * ``DELETE``/``INSERT`` against ``lance_store.main.embeddings`` reach the - ATTACHed store, which the filesystem allowlist does not cover. Re-run - ``atif-sql embed --all --no-dry-run`` to rebuild. + nothing else. ``read_text`` outside the corpus, ``ATTACH`` of unrelated + databases and extension installs fail with exit 70 rather than reading + credentials or reaching the network, and nothing under the corpus root + is writable. Two layers keep caller SQL from writing the corpus: + + * DuckDB's file grants are read-write, and ``COPY ... TO (USE_TMP_FILE false)`` used to overwrite one, so every + statement kind that names a file is refused BEFORE execution, for any + uid, by DuckDB's own parser: COPY, EXPORT, ATTACH, DETACH, INSTALL, + LOAD, and PREPARE/EXECUTE (which could defer one). Exit 70, kind + ``sandbox_refused``. + * A 0444 file mode does not bind root, so ``query`` refuses to run as + uid 0 (exit 77) unless ``ATIF_SQL_ALLOW_ROOT=1`` is set, which logs a + warning. + + The embedding-store guard binds here exactly as it does for ``search``, + so a store written by a different provider refuses to bind (exit 65) + instead of scoring garbage. + + One residual hole remains, over recomputable derived data rather than + the transcript artifacts: ``DELETE``/``INSERT`` against + ``lance_store.main.embeddings`` reach the ATTACHed store, which the + filesystem allowlist does not cover. Re-run + ``atif-sql embed --all --no-dry-run`` to rebuild. + + What caller SQL can still see, and why that is accepted: + ``duckdb_settings()`` and ``current_setting(...)`` return the sandbox's + own configuration, including the corpus root, the spill directory, the + memory cap and every granted parquet path, which names every session id. + DuckDB cannot hide a setting from SQL, and the grants have to be per-file + for the lazy reads to work at all, so the inventory is a property of the + design. The caller is the local user, who can list the corpus directory + and ``SELECT session_id FROM sessions`` anyway; ``query`` is not a + privilege boundary (SECURITY.md) and must not be exposed to a caller who + may not know the session ids. ``lock_configuration`` freezes the connection's settings, so ``SET`` is refused for everything except ``TimeZone``, which stays open for @@ -787,10 +1122,15 @@ def query( Exit codes ---------- * 64 parse_error malformed SQL (or no SQL and no --examples) + * 64 invalid_input a malformed ``ATIF_SQL_QUERY_MEMORY_LIMIT`` or + ``ATIF_SQL_QUERY_THREADS`` * 65 catalog_error unknown view/macro/column (try ``atif-sql schema``) * 65 embedding_mismatch the Lance store was written by another provider + * 70 sandbox_refused a statement kind the sandbox never runs (COPY, + EXPORT, ATTACH, INSTALL, LOAD, PREPARE, ...) * 70 runtime_error everything else (an unmaterialized corpus, or SQL the sandbox refused) + * 77 root_refused running as uid 0 without ``ATIF_SQL_ALLOW_ROOT=1`` Examples -------- @@ -821,6 +1161,20 @@ def query( ) raise SystemExit(EXIT_CODES["parse_error"]) + _refuse_root("query", fmt) + try: + resources = _query_resources() + except ValueError as exc: + err = ClassifiedError( + kind="invalid_input", + exit_code=EXIT_CODES["invalid_input"], + message=str(exc), + hint=f"{QUERY_MEMORY_LIMIT_ENV} takes a size such as 6GB or 512MiB; " + f"{QUERY_THREADS_ENV} takes a positive integer", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + import duckdb from atif_cli.duck_errors import REGISTRATION_ERRORS, classify_registration_error @@ -832,35 +1186,49 @@ def query( expected_model, expected_dim = embed_settings.expected_embedding_identity() lance_uri = embed_settings.resolve_lance_uri(settings.corpus_root) - con = duckdb.connect() - try: - try: - sources = register( - con, - settings.corpus_root, - lance_uri=lance_uri, - expected_model=expected_model, - expected_dim=expected_dim, - ) - _harden_query_connection( - con, - corpus_root=settings.corpus_root, - temp_dir=settings.corpus_root / ".duckdb_tmp", - columnar_paths=sources.lazy_read_paths, - ) - cursor = con.execute(sql) - except REGISTRATION_ERRORS as exc: - err = classify_registration_error(exc) - emit_error(err, fmt) - raise SystemExit(err.exit_code) from exc + with _private_spill_dir() as spill_dir: + con = duckdb.connect() try: - emit_cursor(cursor, fmt) - except REGISTRATION_ERRORS as exc: - err = classify_registration_error(exc) - emit_error(err, fmt) - raise SystemExit(err.exit_code) from exc - finally: - con.close() + try: + _configure_query_resources(con, resources, spill_dir) + sources = register( + con, + settings.corpus_root, + lance_uri=lance_uri, + expected_model=expected_model, + expected_dim=expected_dim, + ) + _harden_query_connection( + con, + corpus_root=settings.corpus_root, + spill_dir=spill_dir, + columnar_paths=sources.lazy_read_paths, + ) + refused = _refused_statement_kinds(con, sql) + if refused: + err = ClassifiedError( + kind="sandbox_refused", + exit_code=EXIT_CODES["sandbox_refused"], + message=f"the query sandbox does not run {', '.join(refused)} statements", + hint="query runs SELECT and in-memory DDL/DML only; COPY, EXPORT, " + "ATTACH, DETACH, INSTALL, LOAD, PREPARE and EXECUTE are refused " + "because they reach files or defer a statement past this check", + ) + emit_error(err, fmt) + raise SystemExit(err.exit_code) + cursor = con.execute(sql) + except REGISTRATION_ERRORS as exc: + err = classify_registration_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + try: + emit_cursor(cursor, fmt) + except REGISTRATION_ERRORS as exc: + err = classify_registration_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + finally: + con.close() # --------------------------------------------------------------------------- @@ -930,6 +1298,8 @@ def analyze( fmt Summary format; ``auto`` = JSON on a pipe. """ + _refuse_root("analyze", fmt) + from atif_analytics.application.analyze import run_analyze from atif_analytics.infrastructure.settings import AnalyticsSettings @@ -975,17 +1345,67 @@ def analyze( # --------------------------------------------------------------------------- +def _install_lance_extension(fmt: OutputFormat, *, quiet: bool = False) -> None: + """``INSTALL lance; LOAD lance`` on a throwaway connection. + + The one network reach outside Bedrock, so it lives behind ``embed``. + With ``quiet`` (a real embed run) a failure is a warning rather than an + exit: the backfill itself still runs, and ``status`` will keep saying + the extension is missing until a later attempt succeeds. Without + ``quiet`` (``--install-extension``) the result is emitted as JSON and a + failure exits 70. + """ + import duckdb + + from atif_cli.duck_errors import classify_duckdb_error + from atif_duck.infrastructure.registry import install_lance_extension + + con = duckdb.connect() + try: + try: + install_path = install_lance_extension(con) + except duckdb.Error as exc: + if quiet: + from loguru import logger + + logger.warning( + "Could not install the lance DuckDB extension ({}); vector search stays " + "unavailable until `atif-sql embed --install-extension` succeeds", + str(exc).splitlines()[0], + ) + return + err = classify_duckdb_error(exc) + emit_error(err, fmt) + raise SystemExit(err.exit_code) from exc + finally: + con.close() + if not quiet: + emit_json({"extension": "lance", "installed": True, "install_path": install_path}, fmt) + + @app.command def embed( *, limit: int | None = None, all_steps: Annotated[bool, cyclopts.Parameter(name="--all")] = False, dry_run: bool = False, + install_extension: Annotated[bool, cyclopts.Parameter(name="--install-extension")] = False, corpus_root: Path | None = None, fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO, ) -> None: """Embed unembedded corpus steps with Cohere Embed v4 and append to LanceDB. + Extension + --------- + ``query`` and ``search`` read the store through DuckDB's lance extension + and never install it themselves (installing is a 242 MB download from + the extension repository, and ``query`` runs unattended). It is + installed here, where the network is already a deliberate act: every + REAL run installs it first, and ``--install-extension`` installs it and + exits without touching Bedrock or the store (``{"extension": "lance", + "installed": true, "install_path": ...}``). ``atif-sql status`` reports + whether it is present. + Cost ---- Calls Bedrock (``global.cohere.embed-v4:0``) on every unembedded step @@ -1001,6 +1421,7 @@ def embed( --limit N Cap the number of steps embedded this run. --all Explicitly embed EVERY unembedded step (full backfill). --dry-run Preview only; emit plan JSON, no embedding calls. + --install-extension Install the lance DuckDB extension and exit (no Bedrock). --corpus-root Override the materialized corpus root. Output @@ -1014,6 +1435,10 @@ def embed( state (the store or its config requires operator action; retrying without intervention cannot succeed, so unattended lanes suppress retries on 78). """ + if install_extension: + _install_lance_extension(fmt) + return + import asyncio from atif_embed.application.embed import run_backfill @@ -1034,6 +1459,8 @@ def embed( settings = _corpus_settings(None, corpus_root) embed_settings = EmbedSettings() + if not dry_run: + _install_lance_extension(fmt, quiet=True) try: result = asyncio.run( run_backfill( @@ -1107,8 +1534,13 @@ def search( Sorted by cosine distance ascending — highest sim first. Exit codes: 0 success, 2 no_embeddings, 65 embedding_mismatch (the store - was written by another provider), 70 runtime. + was written by another provider), 70 runtime, 77 root_refused (uid 0 + without ``ATIF_SQL_ALLOW_ROOT=1``), 78 extension_missing (a store exists + but the lance DuckDB extension is not installed; run + ``atif-sql embed --install-extension``). """ + _refuse_root("search", fmt) + import duckdb from atif_cli.duck_errors import ( @@ -1116,7 +1548,7 @@ def search( classify_duckdb_error, classify_registration_error, ) - from atif_duck.infrastructure.registry import register + from atif_duck.infrastructure.registry import lance_extension_installed, register from atif_embed.application.embed import embed_query from atif_embed.infrastructure.settings import EmbedSettings @@ -1140,6 +1572,19 @@ def search( emit_error(err, fmt) raise SystemExit(err.exit_code) from exc + if lance_uri.is_dir() and not lance_extension_installed(con): + emit_error( + ClassifiedError( + kind="extension_missing", + exit_code=EXIT_CODES["extension_missing"], + message="the lance DuckDB extension is not installed, so the embeddings " + f"store at {lance_uri} cannot be read", + hint="run: atif-sql embed --install-extension (a one-time download)", + ), + fmt, + ) + raise SystemExit(EXIT_CODES["extension_missing"]) + row = con.execute("SELECT count(*) FROM message_embeddings").fetchone() if not row or int(row[0]) == 0: emit_error( diff --git a/packages/atif-cli/src/atif_cli/errors.py b/packages/atif-cli/src/atif_cli/errors.py index ca9cac5..6686960 100644 --- a/packages/atif-cli/src/atif_cli/errors.py +++ b/packages/atif-cli/src/atif_cli/errors.py @@ -35,6 +35,15 @@ "validation_error": 65, # convert: trajectory failed TrajectoryValidator "embedding_mismatch": 65, # query/search: Lance store written by another provider "runtime_error": 70, # everything else duckdb (or the adapter) raises + "sandbox_refused": 70, # query: a statement kind the sandbox never runs (COPY, EXPORT, + # ATTACH, INSTALL, LOAD, PREPARE, ...). Same number as runtime_error, which is what + # every other sandbox refusal exits with; the distinct kind names the reason. + "root_refused": 77, # query/search/analyze: running as uid 0 without ATIF_SQL_ALLOW_ROOT=1. + # EX_NOPERM by convention: the process has too MUCH permission for the file modes + # that protect the corpus to mean anything. + "extension_missing": 78, # search: the Lance store exists but the lance DuckDB extension + # is not installed; `atif-sql embed --install-extension` is the operator action, so the + # same EX_CONFIG contract as terminal_state. "terminal_state": 78, # embed: store/config state requires OPERATOR action (EX_CONFIG); # retrying without intervention cannot succeed — unattended lanes suppress # retries on this code instead of burning identical ticks diff --git a/packages/atif-cli/tests/cli_fixtures.py b/packages/atif-cli/tests/cli_fixtures.py index f1bd4b6..73f2373 100644 --- a/packages/atif-cli/tests/cli_fixtures.py +++ b/packages/atif-cli/tests/cli_fixtures.py @@ -201,3 +201,21 @@ def write_analytics_parquets(corpus_root: Path, session_id: str) -> list[Path]: finally: scratch.close() return sorted(written) + + +# ``register_vss`` no longer installs the lance extension (that download at +# query time was finding 4 of the MicroVM review), so the suite installs it +# once up front. A no-op where it is already present; a one-time download on a +# fresh runner, exactly what every register() call used to do implicitly. +# A public name on purpose: ``conftest.py`` re-exports this module with +# ``import *``, which skips underscore names, so an underscore here would leave +# the fixture unregistered (it did, on a runner with no extension cached). +@pytest.fixture(scope="session", autouse=True) +def lance_extension_present() -> None: + import duckdb + + con = duckdb.connect() + try: + con.execute("INSTALL lance") + finally: + con.close() diff --git a/packages/atif-cli/tests/test_app.py b/packages/atif-cli/tests/test_app.py index 1de40a0..611108c 100644 --- a/packages/atif-cli/tests/test_app.py +++ b/packages/atif-cli/tests/test_app.py @@ -4,8 +4,12 @@ from __future__ import annotations +import hashlib import json +import os import re +import shutil +import subprocess import sys from datetime import UTC, datetime from pathlib import Path @@ -13,19 +17,26 @@ import pytest from cli_fixtures import write_analytics_parquets +from loguru import logger from atif_cli import app as app_module from atif_cli.app import ( + ALLOW_ROOT_ENV, DEFAULT_LOG_LEVEL, LOG_LEVEL_ENV, + QUERY_MEMORY_LIMIT_ENV, + QUERY_THREADS_ENV, _stderr_log_level, + analyze, app, convert, + embed, examples, materialize as materialize_cmd, query, schema, search, + status, ) from atif_cli.errors import EXIT_CODES from atif_cli.output import ( @@ -326,16 +337,107 @@ def duckdb_default_memory_limit() -> str: con.close() -def _expect_refusal(sql: str, corpus_root: Path, capsys: pytest.CaptureFixture[str]) -> str: - """Run ``sql`` expecting the sandbox to refuse it; return the message.""" +def _expect_refusal( + sql: str, + corpus_root: Path, + capsys: pytest.CaptureFixture[str], + *, + kind: str = "runtime_error", +) -> str: + """Run ``sql`` expecting the sandbox to refuse it; return the message. + + ``kind`` is ``runtime_error`` for a refusal DuckDB itself raises + (allowlist, locked configuration) and ``sandbox_refused`` for a + statement kind the CLI refuses by name before executing anything. + """ with pytest.raises(SystemExit) as excinfo: _run_query(sql, corpus_root) - assert excinfo.value.code == EXIT_CODES["runtime_error"] + assert excinfo.value.code == EXIT_CODES[kind] payload = json.loads(capsys.readouterr().err) - assert payload["error"]["kind"] == "runtime_error" + assert payload["error"]["kind"] == kind return str(payload["error"]["message"]) +def _tree_digest(root: Path) -> str: + """One hash over every path, mode and byte under ``root``.""" + digest = hashlib.sha256() + for path in sorted(p for p in root.rglob("*")): + rel = path.relative_to(root).as_posix() + digest.update(f"{rel}\0{oct(path.stat().st_mode)}\0".encode()) + if path.is_file(): + digest.update(path.read_bytes()) + digest.update(b"\n") + return digest.hexdigest() + + +def _capture_warnings() -> tuple[list[str], int]: + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + return warnings, sink_id + + +def _write_two_session_corpus(root: Path) -> Path: + """The contract corpus plus a second session: two ``edges.jsonl`` for the glob reader. + + One file registers under any limit because the reader uses one thread + per file; the per-thread reservation that finding 1 is about needs at + least two. + """ + _write_contract_corpus(root) + second = "22222222-2222-2222-2222-222222222222" + src = root / "sessions" / SESSION_ID + dst = root / "sessions" / second + shutil.copytree(src, dst) + for name in ("trajectory.json", "meta.json"): + path = dst / name + path.write_text(path.read_text().replace(SESSION_ID, second)) + return root + + +class _RecordingConnection: + """Wrap a DuckDB connection so every statement text is observable. + + ``execute`` is recorded; everything else (``read_parquet``, + ``extract_statements``, ``close``) delegates. Registration only ever + touches those, so the wrapper is transparent to the code under test. + """ + + def __init__(self, con: Any, statements: list[str]) -> None: + self._con = con + self._statements = statements + + def execute(self, sql: str, *args: Any, **kwargs: Any) -> Any: + self._statements.append(sql) + return self._con.execute(sql, *args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._con, name) + + +@pytest.fixture +def empty_extension_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, list[str]]: + """Every ``duckdb.connect()`` in the command bodies sees NO installed extensions. + + The connection is pointed at an empty ``extension_directory`` and + wrapped so the statements it runs are recorded. Returns the directory + (still empty afterwards proves nothing was installed) and the recording. + """ + import duckdb + + ext_dir = tmp_path / "no-extensions" + ext_dir.mkdir() + statements: list[str] = [] + real_connect = duckdb.connect + + def connect(*args: Any, **kwargs: Any) -> Any: + con = real_connect(*args, **kwargs) + con.execute(f"SET extension_directory='{ext_dir}'") + return _RecordingConnection(con, statements) + + monkeypatch.setattr(duckdb, "connect", connect) + return ext_dir, statements + + class TestMaterializeReportOutput: """An unreadable session is counted nowhere else, so hiding it reads as success.""" @@ -403,19 +505,24 @@ def test_copy_out_of_tree_is_refused( ) -> None: target = tmp_path / "exfil.csv" - _expect_refusal(f"COPY (SELECT 42 AS x) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42 AS x) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() def test_attach_unrelated_database_is_refused( self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] ) -> None: - _expect_refusal(f"ATTACH '{tmp_path / 'other.db'}' AS o", query_corpus, capsys) + _expect_refusal( + f"ATTACH '{tmp_path / 'other.db'}' AS o", query_corpus, capsys, kind="sandbox_refused" + ) - def test_extension_install_is_refused( + def test_extension_install_and_load_are_refused( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: """No INSTALL/LOAD means no httpfs, which means no network egress.""" - _expect_refusal("INSTALL httpfs", query_corpus, capsys) + _expect_refusal("INSTALL httpfs", query_corpus, capsys, kind="sandbox_refused") + _expect_refusal("LOAD httpfs", query_corpus, capsys, kind="sandbox_refused") def test_injected_sql_cannot_reopen_the_sandbox( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] @@ -426,6 +533,8 @@ def test_injected_sql_cannot_reopen_the_sandbox( "SET allowed_directories=['/']", "SET memory_limit='512GB'", "SET temp_directory='/tmp'", + "SET autoinstall_known_extensions=true", + "SET threads=64", ): _expect_refusal(statement, query_corpus, capsys) @@ -465,6 +574,180 @@ def test_corpus_root_containing_a_quote_still_binds( assert json.loads(capsys.readouterr().out) == [{"n": 1}] +class TestQueryStatementGate: + """Statement kinds that name a file are refused by name before anything executes. + + DuckDB's ``allowed_paths`` grants are read-write and it has no read-only + grant (pinned in :class:`TestHardenedConnectionLayer`), so this gate is + what stops ``COPY ... TO (USE_TMP_FILE false)`` for + any uid. + """ + + @pytest.mark.parametrize( + ("statement", "kind"), + [ + ("COPY (SELECT 1) TO '{tmp}/a.csv'", "COPY"), + ("COPY (SELECT 1) TO '{tmp}/a.csv' (USE_TMP_FILE false)", "COPY"), + ("EXPORT DATABASE '{tmp}/exported'", "EXPORT"), + ("PREPARE p AS COPY (SELECT 1) TO '{tmp}/a.csv'", "PREPARE"), + ("EXECUTE p", "EXECUTE"), + ("ATTACH '{tmp}/other.db' AS o", "ATTACH"), + ("DETACH o", "DETACH"), + # DuckDB parses INSTALL and LOAD to the same statement kind. + ("INSTALL httpfs", "LOAD"), + ("LOAD httpfs", "LOAD"), + ], + ) + def test_file_statement_kinds_are_refused_by_name( + self, + query_corpus: Path, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + statement: str, + kind: str, + ) -> None: + message = _expect_refusal( + statement.format(tmp=tmp_path), query_corpus, capsys, kind="sandbox_refused" + ) + assert kind in message + assert not (tmp_path / "a.csv").exists() + + def test_a_refused_statement_stops_the_whole_batch( + self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + """Nothing in a batch runs when any statement in it is refused.""" + target = tmp_path / "late.csv" + message = _expect_refusal( + f"SELECT 1; COPY (SELECT 1) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) + assert "COPY" in message + assert not target.exists() + + def test_select_batches_and_in_memory_ddl_still_run( + self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + _run_query("SELECT 1 AS a; SELECT 2 AS b", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"b": 2}] + _run_query("CREATE TEMP TABLE t AS SELECT 1 AS x; SELECT * FROM t", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"x": 1}] + _run_query("EXPLAIN SELECT count(*) FROM sessions", query_corpus) + assert json.loads(capsys.readouterr().out) + + def test_parse_errors_still_exit_64( + self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + with pytest.raises(SystemExit) as excinfo: + _run_query("SELEC 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["parse_error"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "parse_error" + + def test_allowlist_names_real_statement_kinds_and_omits_the_file_facing_ones(self) -> None: + import duckdb + + known = {member for member in dir(duckdb.StatementType) if member.isupper()} + assert known >= app_module._QUERY_STATEMENT_KINDS + for refused in ( + "COPY", + "COPY_DATABASE", + "EXPORT", + "ATTACH", + "DETACH", + "LOAD", + "EXTENSION", + "PREPARE", + "EXECUTE", + ): + assert refused in known + assert refused not in app_module._QUERY_STATEMENT_KINDS + + +class TestQueryRefusesRoot: + """A 0444 file mode does not bind uid 0, so the commands that run SQL refuse it.""" + + @staticmethod + def _as_root(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "geteuid", lambda: 0) + monkeypatch.delenv(ALLOW_ROOT_ENV, raising=False) + + def test_query_exits_77_before_opening_duckdb( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + import duckdb + + self._as_root(monkeypatch) + + def refuse_connect(*args: object, **kwargs: object) -> Any: + del args, kwargs + pytest.fail("DuckDB was opened as root") + + monkeypatch.setattr(duckdb, "connect", refuse_connect) + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["root_refused"] == 77 + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "root_refused" + assert "root" in err["message"] + assert ALLOW_ROOT_ENV in err["hint"] + + def test_search_and_analyze_refuse_root_too( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + with pytest.raises(SystemExit) as excinfo: + search("anything", corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["root_refused"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "root_refused" + with pytest.raises(SystemExit) as excinfo: + analyze(corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["root_refused"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "root_refused" + + def test_the_escape_hatch_runs_with_a_warning( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + monkeypatch.setenv(ALLOW_ROOT_ENV, "1") + warnings, sink_id = _capture_warnings() + try: + _run_query("SELECT count(*) AS n FROM sessions", query_corpus) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out) == [{"n": 1}] + assert any("root" in w and ALLOW_ROOT_ENV in w for w in warnings) + + def test_only_the_value_one_opens_the_hatch( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + self._as_root(monkeypatch) + monkeypatch.setenv(ALLOW_ROOT_ENV, "yes") + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["root_refused"] + capsys.readouterr() + + def test_an_unprivileged_uid_is_untouched( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + monkeypatch.setattr(os, "geteuid", lambda: 1000) + _run_query("SELECT count(*) AS n FROM sessions", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 1}] + + class TestQuerySandboxFileAllowlist: """The allowlist gates LAZY reads only — the surface a corpus-only test misses. @@ -521,10 +804,12 @@ def test_copy_over_a_corpus_artifact_is_refused( trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" before = trajectory.read_text() - message = _expect_refusal( - f"COPY (SELECT 'destroyed' AS x) TO '{trajectory}'", query_corpus, capsys + _expect_refusal( + f"COPY (SELECT 'destroyed' AS x) TO '{trajectory}'", + query_corpus, + capsys, + kind="sandbox_refused", ) - assert "Permission" in message assert trajectory.read_text() == before def test_copy_new_file_into_the_corpus_is_refused( @@ -532,7 +817,9 @@ def test_copy_new_file_into_the_corpus_is_refused( ) -> None: target = query_corpus / "pwned.csv" - _expect_refusal(f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() def test_copy_into_the_lance_store_is_refused( @@ -541,62 +828,63 @@ def test_copy_into_the_lance_store_is_refused( _write_lance_store(query_corpus / "embeddings_lance") target = query_corpus / "embeddings_lance" / "exfil.csv" - _expect_refusal(f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys) + _expect_refusal( + f"COPY (SELECT 42) TO '{target}'", query_corpus, capsys, kind="sandbox_refused" + ) assert not target.exists() - def test_traversal_out_of_the_spill_directory_is_refused( + def test_copy_into_the_legacy_spill_directory_is_refused_and_nothing_persists( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: - """``.duckdb_tmp`` is writable and sits inside the corpus — ``..`` must not escape. + """Probe d3 of the review: ``/.duckdb_tmp`` used to be the one writable dir. - The spill dir is created up front on purpose: DuckDB reports a - missing parent directory as an IO error, which would make this pass - with no sandbox at all. + It is pre-created here, as the review did, so a success would be the + old grant and not a missing parent; it is no longer granted at all. """ - (query_corpus / ".duckdb_tmp").mkdir(exist_ok=True) - trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" - before = trajectory.read_text() - escape = query_corpus / ".duckdb_tmp" / ".." / "sessions" / SESSION_ID / "trajectory.json" + legacy = query_corpus / ".duckdb_tmp" + legacy.mkdir() + before = _tree_digest(query_corpus) - message = _expect_refusal( - f"COPY (SELECT 'destroyed' AS x) TO '{escape}'", query_corpus, capsys + _expect_refusal( + f"COPY (SELECT 1) TO '{legacy}/probe.csv'", query_corpus, capsys, kind="sandbox_refused" ) - assert "Permission" in message - assert trajectory.read_text() == before + _expect_refusal(f"SELECT * FROM read_text('{legacy}/probe.csv')", query_corpus, capsys) + assert _tree_digest(query_corpus) == before + assert list(legacy.iterdir()) == [] - def test_plain_copy_over_an_analytics_parquet_is_refused( - self, query_corpus: Path, capsys: pytest.CaptureFixture[str] + def test_use_tmp_file_false_over_a_granted_parquet_is_refused_for_any_uid( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: - """The read grant is per-file; DuckDB stages ``COPY`` through an ungranted sibling.""" + """Probe d2 of the review, the one that overwrote ``steps.parquet`` as root. + + The analytics parquet is a granted path and writable by this uid at + the filesystem, so a refusal here is the statement gate and not a + file mode. Repeated as uid 0 with the escape hatch open: still + refused, because the gate does not look at the uid. + """ parquet = query_corpus / "analytics" / "clusters.parquet" before = parquet.read_bytes() + statement = f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)" - _expect_refusal( - f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET)", query_corpus, capsys - ) + _expect_refusal(statement, query_corpus, capsys, kind="sandbox_refused") + assert parquet.read_bytes() == before + + monkeypatch.setattr(os, "geteuid", lambda: 0) + monkeypatch.setenv(ALLOW_ROOT_ENV, "1") + _expect_refusal(statement, query_corpus, capsys, kind="sandbox_refused") assert parquet.read_bytes() == before - def test_documented_residual_holes_still_behave_as_documented( + def test_documented_residual_hole_still_behaves_as_documented( self, query_corpus: Path, capsys: pytest.CaptureFixture[str] ) -> None: - """Pins the two holes ``query``'s docstring admits, so the promise cannot rot. + """Pins the one hole ``query``'s docstring admits, so the promise cannot rot. - Both are recomputable derived data. If a future DuckDB closes either, - this test fails and the docstring's Sandbox section must lose the - corresponding paragraph. + Recomputable derived data. If a future DuckDB closes it, this test + fails and the docstring's Sandbox section must lose the paragraph. """ - parquet = query_corpus / "analytics" / "clusters.parquet" - before = parquet.read_bytes() - _run_query( - f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)", - query_corpus, - ) - capsys.readouterr() - assert parquet.read_bytes() != before, ( - "USE_TMP_FILE false no longer overwrites a granted parquet — " - "drop that paragraph from query's Sandbox docstring" - ) - _write_lance_store(query_corpus / "embeddings_lance") _run_query("DELETE FROM lance_store.main.embeddings", query_corpus) capsys.readouterr() @@ -607,67 +895,470 @@ def test_documented_residual_holes_still_behave_as_documented( ) -class TestQueryMemoryCap: - """``_query_memory_limit_bytes`` is pure arithmetic over ``os.sysconf``.""" +def _hardened_connection(corpus_root: Path, spill_dir: Path) -> Any: + """What ``query`` builds, minus the statement gate: DuckDB's own layer alone.""" + import duckdb - @staticmethod - def _limit_for(physical_bytes: int, monkeypatch: pytest.MonkeyPatch) -> int: - import os + from atif_duck.infrastructure.registry import register - from atif_cli import app as app_mod + con = duckdb.connect() + app_module._configure_query_resources(con, app_module._query_resources(), spill_dir) + sources = register(con, corpus_root) + app_module._harden_query_connection( + con, + corpus_root=corpus_root, + spill_dir=spill_dir, + columnar_paths=sources.lazy_read_paths, + ) + return con - page = 4096 - def _sysconf(name: str | int) -> int: - return page if name == "SC_PAGE_SIZE" else physical_bytes // page +class TestHardenedConnectionLayer: + """DuckDB's allowlist, exercised WITHOUT the statement gate. - monkeypatch.setattr(os, "sysconf", _sysconf) - return app_mod._query_memory_limit_bytes() + The CLI tests above hit the gate first, so this class is what proves + the second layer still holds on its own, and pins the two facts the + gate exists for: a directory grant and a file grant are both read-write. + """ - def test_ceiling_wins_on_a_small_host(self, monkeypatch: pytest.MonkeyPatch) -> None: - """4 GiB of RAM: 50% is 2 GiB, but the 80% ceiling (3.2 GiB) is tighter than the floor.""" - from atif_cli.app import _QUERY_MEMORY_FLOOR_BYTES + @pytest.fixture + def layer(self, query_corpus: Path, tmp_path: Path) -> Any: + spill = tmp_path / "spill" + spill.mkdir(mode=0o700) + con = _hardened_connection(query_corpus, spill) + yield con, spill + con.close() - physical = 4 * 1024**3 - assert self._limit_for(physical, monkeypatch) == int(physical * 0.8) - assert int(physical * 0.8) < _QUERY_MEMORY_FLOOR_BYTES + def test_duckdb_refuses_every_ungranted_path( + self, layer: Any, query_corpus: Path, tmp_path: Path + ) -> None: + import duckdb - def test_half_of_ram_wins_on_a_large_host(self, monkeypatch: pytest.MonkeyPatch) -> None: - physical = 128 * 1024**3 - assert self._limit_for(physical, monkeypatch) == int(physical * 0.5) + con, spill = layer + trajectory = query_corpus / "sessions" / SESSION_ID / "trajectory.json" + parquet = query_corpus / "analytics" / "clusters.parquet" + for statement in ( + f"COPY (SELECT 1) TO '{tmp_path / 'out.csv'}'", + f"COPY (SELECT 1) TO '{trajectory}'", + f"COPY (SELECT 1) TO '{trajectory}' (USE_TMP_FILE false)", + f"COPY (SELECT 1) TO '{query_corpus / 'pwned.csv'}'", + f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET)", + f"COPY (SELECT 1) TO '{spill}/../escape.csv'", + f"SELECT * FROM read_text('{spill}/../../etc/passwd')", + "SELECT * FROM read_text('/etc/passwd')", + f"ATTACH '{tmp_path / 'other.db'}' AS o", + "INSTALL httpfs", + ): + with pytest.raises(duckdb.Error) as excinfo: + con.execute(statement) + assert "Permission" in str(excinfo.value) or "disabled" in str(excinfo.value), statement + assert not (tmp_path / "out.csv").exists() + assert not (tmp_path / "escape.csv").exists() + + def test_grants_are_read_write_which_is_why_the_gate_exists( + self, layer: Any, query_corpus: Path + ) -> None: + """DuckDB 1.5.5 has no read-only grant; if one appears, this fails and the design can simplify.""" + con, spill = layer + con.execute(f"COPY (SELECT 1 AS x) TO '{spill}/probe.csv'") + assert (spill / "probe.csv").is_file() - def test_floor_beats_half_between_the_two(self, monkeypatch: pytest.MonkeyPatch) -> None: - """12 GiB: half (6 GiB) is under the floor, and the floor is under the 80% ceiling.""" - from atif_cli.app import _QUERY_MEMORY_FLOOR_BYTES + parquet = query_corpus / "analytics" / "clusters.parquet" + before = parquet.read_bytes() + con.execute( + f"COPY (SELECT 'x' AS uuid) TO '{parquet}' (FORMAT PARQUET, USE_TMP_FILE false)" + ) + assert parquet.read_bytes() != before, ( + "DuckDB now refuses a write to a granted path: the statement gate is " + "belt-and-braces and query's Sandbox docstring can say so" + ) - physical = 12 * 1024**3 - assert self._limit_for(physical, monkeypatch) == _QUERY_MEMORY_FLOOR_BYTES - assert int(physical * 0.5) < _QUERY_MEMORY_FLOOR_BYTES < int(physical * 0.8) + def test_resources_and_extension_settings_are_in_force(self, layer: Any) -> None: + con, spill = layer + resources = app_module._query_resources() + settings = dict( + con.execute( + "SELECT name, value FROM duckdb_settings() WHERE name IN " + "('threads', 'temp_directory', 'autoinstall_known_extensions', " + "'autoload_known_extensions', 'enable_external_access', 'lock_configuration')" + ).fetchall() + ) + assert int(settings["threads"]) == resources.threads + assert settings["temp_directory"] == str(spill) + assert settings["autoinstall_known_extensions"] == "false" + assert settings["autoload_known_extensions"] == "false" + assert settings["enable_external_access"] == "false" + assert settings["lock_configuration"] == "true" - def test_never_exceeds_duckdbs_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: - for gib in (1, 4, 16, 64, 256): - physical = gib * 1024**3 - assert self._limit_for(physical, monkeypatch) <= int(physical * 0.8) - def test_connection_carries_the_cap_and_a_corpus_local_spill_dir( - self, query_corpus: Path, capsys: pytest.CaptureFixture[str] +class TestQuerySpillDirectory: + """The spill directory is private, outside the corpus, and gone when the process is.""" + + @pytest.fixture + def recorded_mkdtemp(self, monkeypatch: pytest.MonkeyPatch) -> list[tuple[Path, int]]: + import tempfile + + created: list[tuple[Path, int]] = [] + real = tempfile.mkdtemp + + def mkdtemp(*args: Any, **kwargs: Any) -> str: + path = str(real(*args, **kwargs)) + created.append((Path(path), Path(path).stat().st_mode & 0o777)) + return path + + # ``app`` reads ``tempfile.mkdtemp`` at call time, so the module attribute is the seam. + monkeypatch.setattr(tempfile, "mkdtemp", mkdtemp) + return created + + def test_private_dir_outside_the_corpus_removed_on_success( + self, + query_corpus: Path, + recorded_mkdtemp: list[tuple[Path, int]], + capsys: pytest.CaptureFixture[str], ) -> None: - """The settings are read back through caller SQL — the only honest witness.""" - from atif_cli.app import _query_memory_limit_bytes + _run_query( + "SELECT current_setting('temp_directory') AS tmp, " + "current_setting('allowed_directories') AS dirs", + query_corpus, + ) + row = json.loads(capsys.readouterr().out)[0] + [(spill, mode)] = recorded_mkdtemp + assert row["tmp"] == str(spill) + assert [Path(d) for d in row["dirs"]] == [spill] + assert spill.name.startswith("atif-sql-query-") + assert mode == 0o700 + assert not spill.is_relative_to(query_corpus) + assert not spill.exists() + assert not (query_corpus / ".duckdb_tmp").exists() + + def test_removed_on_error_too( + self, + query_corpus: Path, + recorded_mkdtemp: list[tuple[Path, int]], + capsys: pytest.CaptureFixture[str], + ) -> None: + with pytest.raises(SystemExit): + _run_query("SELEC 1", query_corpus) + capsys.readouterr() + [(spill, _)] = recorded_mkdtemp + assert not spill.exists() + + def test_a_query_process_leaves_the_corpus_tree_unchanged( + self, query_corpus: Path, tmp_path: Path + ) -> None: + """Through a real process: the review's probe d3, then a plain read. + + Before the fix the COPY exited 0 and ``probe.csv`` persisted inside + the corpus. Now the tree digest (paths, modes, bytes) is unchanged, + the pre-created legacy dir stays empty, no spill dir is left under + the process's TMPDIR, and a fresh corpus gains no ``.duckdb_tmp``. + """ + legacy = query_corpus / ".duckdb_tmp" + legacy.mkdir() + scratch = tmp_path / "process-tmp" + scratch.mkdir() + env: dict[str, str] = dict(os.environ) + env.update( + TMPDIR=str(scratch), + NO_COLOR="1", + PYTHONHASHSEED="0", + LITELLM_LOCAL_MODEL_COST_MAP="true", + ) + for var in ("ATIF_SQL_CORPUS_ROOT", "ATIF_SQL_LANCE_URI", "ATIF_SQL_EMBED_MODEL_ID"): + env.pop(var, None) + before = _tree_digest(query_corpus) + + def run(sql: str) -> subprocess.CompletedProcess[str]: + argv = [sys.executable, "-m", "atif_cli", "query", "--format", "json"] + argv += ["--corpus-root", str(query_corpus), sql] + return subprocess.run( # noqa: S603 - argv is built from constants and tmp paths + argv, capture_output=True, text=True, env=env, timeout=300, check=False + ) + + refused = run(f"COPY (SELECT 1) TO '{legacy}/probe.csv'") + assert refused.returncode == EXIT_CODES["sandbox_refused"], refused.stderr + assert json.loads(refused.stderr.strip().splitlines()[-1])["error"]["kind"] == ( + "sandbox_refused" + ) + read = run("SELECT count(*) AS n FROM sessions") + assert read.returncode == 0, read.stderr + assert json.loads(read.stdout) == [{"n": 1}] + + assert _tree_digest(query_corpus) == before + assert list(legacy.iterdir()) == [] + assert list(scratch.iterdir()) == [], "a spill dir outlived its process" + + legacy.rmdir() + fresh = run("SELECT count(*) AS n FROM steps") + assert fresh.returncode == 0, fresh.stderr + assert not (query_corpus / ".duckdb_tmp").exists() + + +class TestQueryResources: + """The cap and thread count are derived from the host and applied BEFORE registration.""" + + GIB = 1024**3 + + @staticmethod + def _host(monkeypatch: pytest.MonkeyPatch, physical: int, available: int) -> None: + monkeypatch.setattr(app_module, "_host_memory", lambda: (physical, available)) + def test_large_host_gets_half_of_ram(self, monkeypatch: pytest.MonkeyPatch) -> None: + self._host(monkeypatch, 124 * self.GIB, 89 * self.GIB) + assert app_module._query_memory_limit_bytes() == 62 * self.GIB + + def test_eight_gib_guest_is_capped_below_its_own_ram( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The old rule said 8 GiB here; the guest had 8 GiB in total and 6 free.""" + self._host(monkeypatch, 8 * self.GIB, 6 * self.GIB) + limit = app_module._query_memory_limit_bytes() + assert limit == int(6 * self.GIB * 0.8) + assert limit < 8 * self.GIB + + def test_target_wins_between_half_and_eighty_percent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + self._host(monkeypatch, 12 * self.GIB, 12 * self.GIB) + assert app_module._query_memory_limit_bytes() == 8 * self.GIB + + def test_never_exceeds_eighty_percent_of_physical_or_available( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + for gib in (1, 4, 8, 16, 64, 256): + for free_fraction in (0.1, 0.5, 1.0): + physical = gib * self.GIB + available = int(physical * free_fraction) + self._host(monkeypatch, physical, available) + limit = app_module._query_memory_limit_bytes() + assert limit <= max(int(physical * 0.8), app_module._QUERY_MEMORY_MIN_BYTES) + assert limit <= max(int(available * 0.8), app_module._QUERY_MEMORY_MIN_BYTES) + + def test_floor_when_almost_nothing_is_free(self, monkeypatch: pytest.MonkeyPatch) -> None: + self._host(monkeypatch, 2 * self.GIB, 100 * 1024**2) + assert app_module._query_memory_limit_bytes() == app_module._QUERY_MEMORY_MIN_BYTES + + @staticmethod + def _sixteen_cpus(_pid: int) -> set[int]: + return set(range(16)) + + def test_threads_follow_the_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "sched_getaffinity", self._sixteen_cpus, raising=False) + assert app_module._query_threads(int(4.8 * self.GIB)) == 2 + assert app_module._query_threads(62 * self.GIB) == 16 + assert app_module._query_threads(1 * self.GIB) == 1 + + def test_host_memory_reads_this_host(self) -> None: + physical, available = app_module._host_memory() + assert physical > 0 + assert 0 < available <= physical + + @pytest.mark.parametrize( + ("text", "expected"), + [ + ("6GB", 6 * 10**9), + ("512MiB", 512 * 1024**2), + ("123B", 123), + ("1.5GiB", int(1.5 * 1024**3)), + (" 6 gb ", 6 * 10**9), + ("4096", 4096), + ], + ) + def test_parse_size(self, text: str, expected: int) -> None: + assert app_module._parse_size(text) == expected + + @pytest.mark.parametrize("text", ["", "abc", "6XB", "GB", "-1GB"]) + def test_parse_size_rejects_garbage(self, text: str) -> None: + with pytest.raises(ValueError, match=r"size|unit"): + app_module._parse_size(text) + + def test_env_overrides(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "sched_getaffinity", self._sixteen_cpus, raising=False) + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, "6GB") + monkeypatch.setenv(QUERY_THREADS_ENV, "4") + assert app_module._query_resources() == app_module.QueryResources(6 * 10**9, 4) + monkeypatch.delenv(QUERY_THREADS_ENV) + assert app_module._query_resources().threads == 2, "threads follow the overridden cap" + + @pytest.mark.parametrize( + ("var", "value"), + [(QUERY_MEMORY_LIMIT_ENV, "lots"), (QUERY_THREADS_ENV, "0"), (QUERY_THREADS_ENV, "many")], + ) + def test_malformed_override_exits_64( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + var: str, + value: str, + ) -> None: + monkeypatch.setenv(var, value) + with pytest.raises(SystemExit) as excinfo: + _run_query("SELECT 1", query_corpus) + assert excinfo.value.code == EXIT_CODES["invalid_input"] + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "invalid_input" + assert var in err["hint"] + + def test_connection_carries_the_cap_and_threads( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + """The settings are read back through caller SQL — the only honest witness.""" + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, "3GiB") + monkeypatch.setenv(QUERY_THREADS_ENV, "3") _run_query( "SELECT current_setting('memory_limit') AS mem, " - "current_setting('temp_directory') AS tmp", + "current_setting('threads') AS threads, " + "current_setting('autoinstall_known_extensions') AS autoinstall, " + "current_setting('autoload_known_extensions') AS autoload", query_corpus, ) row = json.loads(capsys.readouterr().out)[0] - assert row["tmp"] == str(query_corpus / ".duckdb_tmp") + assert row["mem"] == "3.0 GiB" + assert row["mem"] != duckdb_default_memory_limit() + assert int(row["threads"]) == 3 + assert row["autoinstall"] is False + assert row["autoload"] is False + + @pytest.mark.parametrize(("threads", "memory"), [(4, "6GB"), (16, "8GB")]) + def test_registration_fits_under_limits_that_used_to_oom( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + threads: int, + memory: str, + ) -> None: + """Finding 1 of the MicroVM review, as the review reproduced it on a workstation. + + ``SET threads=N; SET memory_limit=...`` on every connection before the + CLI touches it (the review's shim), plus the same values through the + env so the CLI applies them too. Before the fix the edges reader + reserved 2 GiB a thread and every query exited 70 out of memory once + the glob matched two or more files (measured on main: one file + registers, two do not); the readers are now sized from the files. + """ + import duckdb + + for var in ("ATIF_SQL_CORPUS_ROOT", "ATIF_SQL_LANCE_URI", "ATIF_SQL_EMBED_MODEL_ID"): + monkeypatch.delenv(var, raising=False) + corpus = _write_two_session_corpus(tmp_path / "corpus") + monkeypatch.setenv(QUERY_MEMORY_LIMIT_ENV, memory) + monkeypatch.setenv(QUERY_THREADS_ENV, str(threads)) + real_connect = duckdb.connect + + def limited(*args: Any, **kwargs: Any) -> Any: + con = real_connect(*args, **kwargs) + con.execute(f"SET threads={threads}; SET memory_limit='{memory}'") + return con + + monkeypatch.setattr(duckdb, "connect", limited) + _run_query("SELECT count(*) AS n FROM sessions", corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 2}] + + +class TestQueryNeverInstallsExtensions: + """Registration must not reach the network, even when an embeddings store exists.""" + + def test_store_present_and_extension_absent_binds_empty_without_installing( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + ext_dir, statements = empty_extension_dir + _write_lance_store(query_corpus / "embeddings_lance") + warnings, sink_id = _capture_warnings() + try: + _run_query("SELECT count(*) AS n FROM message_embeddings", query_corpus) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out) == [{"n": 0}] + assert not any(re.match(r"\s*INSTALL\b", s, re.IGNORECASE) for s in statements) + assert list(ext_dir.rglob("*")) == [], "something was installed into the extension dir" + assert any("--install-extension" in w for w in warnings) + + def test_no_store_means_no_load_at_all( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _, statements = empty_extension_dir + _run_query("SELECT count(*) AS n FROM message_embeddings", query_corpus) + assert json.loads(capsys.readouterr().out) == [{"n": 0}] + assert not any(re.match(r"\s*(INSTALL|LOAD)\b", s, re.IGNORECASE) for s in statements) + + def test_status_reports_the_missing_extension( + self, + query_corpus: Path, + tmp_path: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _write_lance_store(query_corpus / "embeddings_lance") + source = tmp_path / "source" + source.mkdir() + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + payload = json.loads(capsys.readouterr().out) + assert payload["lance_extension_installed"] is False + assert payload["embeddings_store_present"] is True + assert payload["vector_search"] == "extension_missing" + + def test_search_exits_78_instead_of_pretending_the_store_is_empty( + self, + query_corpus: Path, + empty_extension_dir: tuple[Path, list[str]], + capsys: pytest.CaptureFixture[str], + ) -> None: + _write_lance_store(query_corpus / "embeddings_lance") + with pytest.raises(SystemExit) as excinfo: + search("anything", corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["extension_missing"] == 78 + err = json.loads(capsys.readouterr().err)["error"] + assert err["kind"] == "extension_missing" + assert "--install-extension" in err["hint"] + + def test_status_reports_ready_and_no_store( + self, query_corpus: Path, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + source = tmp_path / "source" + source.mkdir() + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + payload = json.loads(capsys.readouterr().out) + assert payload["lance_extension_installed"] is True + assert payload["vector_search"] == "no_store" + _write_lance_store(query_corpus / "embeddings_lance") + status(source_root=source, corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["vector_search"] == "ready" + + def test_embed_install_extension_installs_and_exits( + self, + query_corpus: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + """The explicit install path: no scope flag needed, no Bedrock, JSON receipt.""" + from atif_duck.infrastructure import registry as registry_mod + + calls: list[object] = [] - expected_gib = _query_memory_limit_bytes() / 1024**3 - reported = row["mem"] - assert reported.endswith("GiB"), reported - assert float(reported.removesuffix(" GiB")) == pytest.approx(expected_gib, abs=0.1) - assert reported != duckdb_default_memory_limit() + def fake_install(con: object) -> str: + calls.append(con) + return "/ext/lance.duckdb_extension" + + monkeypatch.setattr(registry_mod, "install_lance_extension", fake_install) + embed(install_extension=True, corpus_root=query_corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out) == { + "extension": "lance", + "installed": True, + "install_path": "/ext/lance.duckdb_extension", + } + assert len(calls) == 1 + assert not (query_corpus / "embeddings_lance").exists() class TestQueryLockedConfiguration: diff --git a/packages/atif-cli/tests/test_vss_commands.py b/packages/atif-cli/tests/test_vss_commands.py index efdb051..eabbf95 100644 --- a/packages/atif-cli/tests/test_vss_commands.py +++ b/packages/atif-cli/tests/test_vss_commands.py @@ -161,6 +161,79 @@ def test_dry_run_needs_no_scope(self, corpus: Path, capsys: pytest.CaptureFixtur payload = json.loads(capsys.readouterr().out) assert payload["dry_run"] is True + @pytest.mark.usefixtures("_fake_cohere") + def test_real_run_installs_the_lance_extension_first( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + """The network is already a deliberate act here, so this is where the download belongs.""" + from atif_duck.infrastructure import registry as registry_mod + + calls: list[object] = [] + + def recording_install(con: object) -> str: + calls.append(con) + return "/ext/lance" + + monkeypatch.setattr(registry_mod, "install_lance_extension", recording_install) + embed(all_steps=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["rows_processed"] == 1 + assert len(calls) == 1 + + def test_dry_run_installs_nothing( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + from atif_duck.infrastructure import registry as registry_mod + + def forbidden_install(con: object) -> str: + del con + pytest.fail("a dry run must not reach the extension repository") + + monkeypatch.setattr(registry_mod, "install_lance_extension", forbidden_install) + embed(dry_run=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert json.loads(capsys.readouterr().out)["dry_run"] is True + + @pytest.mark.usefixtures("_fake_cohere") + def test_a_failed_install_only_warns_and_the_backfill_still_runs( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + import duckdb + from loguru import logger + + from atif_duck.infrastructure import registry as registry_mod + + def failing(con: object) -> str: + del con + problem = "Failed to download extension" + raise duckdb.IOException(problem) + + monkeypatch.setattr(registry_mod, "install_lance_extension", failing) + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + try: + embed(all_steps=True, corpus_root=corpus, fmt=OutputFormat.JSON) + finally: + logger.remove(sink_id) + assert json.loads(capsys.readouterr().out)["rows_processed"] == 1 + assert any("--install-extension" in w for w in warnings) + + def test_install_extension_flag_exits_70_when_the_download_fails( + self, corpus: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + import duckdb + + from atif_duck.infrastructure import registry as registry_mod + + def failing(con: object) -> str: + del con + problem = "Failed to download extension" + raise duckdb.IOException(problem) + + monkeypatch.setattr(registry_mod, "install_lance_extension", failing) + with pytest.raises(SystemExit) as excinfo: + embed(install_extension=True, corpus_root=corpus, fmt=OutputFormat.JSON) + assert excinfo.value.code == EXIT_CODES["runtime_error"] + assert json.loads(capsys.readouterr().err)["error"]["kind"] == "runtime_error" + class TestSearchCommand: def test_empty_store_exits_2(self, corpus: Path, capsys: pytest.CaptureFixture[str]) -> None: diff --git a/packages/atif-corpus/tests/test_session_id.py b/packages/atif-corpus/tests/test_session_id.py index 5fbba60..83f0fab 100644 --- a/packages/atif-corpus/tests/test_session_id.py +++ b/packages/atif-corpus/tests/test_session_id.py @@ -220,3 +220,45 @@ def test_a_clean_pass_reports_no_rejections(self, source_root: Path, corpus_root report = _materialize(source_root, corpus_root) assert report.rejected_session_ids == () assert report.rejected_count == 0 + + +class TestFilenamesThatOnceDestroyedTheCorpus: + """Findings 6 and 7 of the MicroVM review, exactly as reproduced there. + + A transcript named ``..jsonl`` has the stem ``.``; ``sessions_dir / "."`` + IS ``sessions_dir``, so the atomic swap once renamed the whole sessions + tree aside and deleted it, with exit 0 and a report that said nothing. A + transcript named ``*.jsonl`` materialized fine and then multiplied every + other session's rows in the readers. Both names now stop at the boundary. + """ + + VICTIMS = ("victim-1", "victim-2", "victim-3") + + def test_dot_jsonl_and_star_jsonl_leave_the_materialized_tree_alone( + self, source_root: Path, corpus_root: Path + ) -> None: + for sid in self.VICTIMS: + write_session(source_root, sid, mtime_ns=STALE_NS) + first = _materialize(source_root, corpus_root) + assert first.materialized_count == len(self.VICTIMS) + layout = CorpusLayout(corpus_root=corpus_root) + before = sorted(p.name for p in layout.sessions_dir.iterdir()) + assert before == sorted(self.VICTIMS) + + # ``write_session`` names the file ``.jsonl``: these are ``..jsonl`` and ``*.jsonl``. + for hostile in (".", "*"): + write_session(source_root, hostile, mtime_ns=STALE_NS) + assert (source_root / "-proj-a" / "..jsonl").is_file() + assert (source_root / "-proj-a" / "*.jsonl").is_file() + + second = _materialize(source_root, corpus_root) + assert second.rejected_session_ids == ("*", ".") + assert second.rejected_count == 2 + assert second.materialized_count == 0 + assert second.failed_count == 0 + assert second.sessions_removed == 0 + assert sorted(p.name for p in layout.sessions_dir.iterdir()) == before + for sid in self.VICTIMS: + assert (layout.sessions_dir / sid / "meta.json").is_file() + # No loose artifact files at the sessions root, which is what the swap left behind. + assert all(p.is_dir() for p in layout.sessions_dir.iterdir()) diff --git a/packages/atif-duck/src/atif_duck/infrastructure/registry.py b/packages/atif-duck/src/atif_duck/infrastructure/registry.py index 58e4a96..04987d2 100644 --- a/packages/atif-duck/src/atif_duck/infrastructure/registry.py +++ b/packages/atif-duck/src/atif_duck/infrastructure/registry.py @@ -100,7 +100,7 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Iterable, Sequence from pathlib import Path import duckdb @@ -166,13 +166,51 @@ def coverage(self) -> ColumnarCoverage: ) -#: Inlined ``read_json`` upper bound. Live trajectory.json files reach 436 MB -#: because harbor inlines subagent sidechains and tool outputs, so 1 GiB is -#: roughly 2.3x headroom over the largest observed document — not a knob to -#: trim. Lowering it does not reduce registration memory: measured peak RSS is -#: unchanged at 512 MiB and rises at 128 MiB, because this bounds the parse -#: buffer rather than preallocating per thread. -_MAX_OBJECT_SIZE: int = 1_073_741_824 +#: Ceiling for a ``read_json`` ``maximum_object_size``. Live trajectory.json +#: files reach 436 MB because harbor inlines subagent sidechains and tool +#: outputs; 1 GiB is about 2.3x headroom over the largest observed document. +_OBJECT_SIZE_CAP: int = 1_073_741_824 + +#: Floor for the same bound: DuckDB's own default. +_OBJECT_SIZE_FLOOR: int = 16_777_216 + +#: Headroom over the largest file when sizing the bound: a quarter of the file +#: plus one MiB, so a file that grows a little between the stat and the read +#: still parses. +_OBJECT_SIZE_HEADROOM_DIVISOR: int = 4 +_OBJECT_SIZE_HEADROOM_BYTES: int = 1_048_576 + +#: The DuckDB extension that reads the Lance embeddings store. +LANCE_EXTENSION: str = "lance" + + +def _object_size_bound(paths: Iterable[Path]) -> int: + """``maximum_object_size`` for a ``read_json`` over ``paths``, sized from the files. + + The bound is not free. DuckDB's eager JSON reader reserves about twice + this many bytes PER THREAD before it parses anything, so the 1 GiB + constant this replaced cost 2 GiB a thread: over 300 ``edges.jsonl`` + files of under 8 MB each it exhausted a 6 GB ``memory_limit`` at four + threads and a 25 GB one at sixteen, while 16 MiB read the same rows in + 0.10 s (measured on the 300-session panel corpus). An earlier note here + claimed lowering it changed nothing; that measurement held the thread + count at one. + + A newline-delimited file's objects are its lines and a one-document file + IS its object, so the largest file present bounds every object either + reader meets. That size plus headroom, floored at DuckDB's default and + capped at :data:`_OBJECT_SIZE_CAP`, is the bound. A path that vanished + between the listing and the stat counts as zero. + """ + largest = 0 + for path in paths: + try: + largest = max(largest, path.stat().st_size) + except OSError: + continue + with_headroom = largest + largest // _OBJECT_SIZE_HEADROOM_DIVISOR + _OBJECT_SIZE_HEADROOM_BYTES + return max(_OBJECT_SIZE_FLOOR, min(_OBJECT_SIZE_CAP, with_headroom)) + # Explicit projection for ``v_raw_trajectories``. # @@ -618,7 +656,9 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource # the columnar sessions' trajectory.json too, which is the cost # this whole arrangement exists to avoid. The list is the # statement's one parameter. - json_files = [str(sessions_dir / sid / "trajectory.json") for sid in json_ids] + trajectory_paths = [sessions_dir / sid / "trajectory.json" for sid in json_ids] + json_files = [str(path) for path in trajectory_paths] + trajectory_bound = _object_size_bound(trajectory_paths) con.execute( f""" CREATE OR REPLACE TEMP TABLE {_RAW_TRAJECTORIES_JSON_TABLE} AS @@ -631,9 +671,9 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource format='auto', filename=true, columns={{{_render_columns_clause(_TRAJECTORY_COLUMNS)}}}, - maximum_object_size={_MAX_OBJECT_SIZE} + maximum_object_size={int(trajectory_bound)} ); - """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns/cap are constants + """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns are constants; the bound is an int [json_files], ) logger.debug( @@ -663,10 +703,10 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource # edges.jsonl is one line per RAW record -> newline_delimited. # The record's own `source_file` (the raw transcript path) is kept; - # the edges.jsonl path itself is aliased to `edges_path`. - # edges.jsonl is one line per RAW record -> newline_delimited. - # The record's own `source_file` (the raw transcript path) is kept; - # the edges.jsonl path itself is aliased to `edges_path`. + # the edges.jsonl path itself is aliased to `edges_path`. The glob + # reads every matching file (the meta gate filters rows afterwards), + # so the bound is sized over every file the glob can reach. + edges_bound = _object_size_bound(sessions_dir.glob("*/edges.jsonl")) con.execute( f""" CREATE OR REPLACE TEMP TABLE {_RAW_EDGES_TABLE} AS @@ -680,10 +720,10 @@ def register_raw(con: duckdb.DuckDBPyConnection, corpus_root: Path) -> RawSource format='newline_delimited', filename=true, columns={{{_render_columns_clause(_EDGE_COLUMNS)}}}, - maximum_object_size={_MAX_OBJECT_SIZE} + maximum_object_size={int(edges_bound)} ) ) WHERE {meta_gate}; - """, # noqa: S608 # nosec B608 - glob is a bound parameter; table/columns/gate are constants + """, # noqa: S608 # nosec B608 - glob is a bound parameter; table/columns/gate are constants; the bound is an int [edges_glob], ) logger.debug("Registered {} from glob {}", _RAW_EDGES_TABLE, edges_glob) @@ -1186,6 +1226,50 @@ def _lance_table_present(con: duckdb.DuckDBPyConnection) -> bool: return row is not None and int(row[0]) > 0 +def lance_extension_installed(con: duckdb.DuckDBPyConnection) -> bool: + """True when the lance extension is present in this connection's extension directory. + + Read from ``duckdb_extensions()``, which lists the local directory and + never reaches the network, so the answer costs a few milliseconds and no + download. + """ + row = con.execute( + "SELECT installed FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION] + ).fetchone() + return row is not None and bool(row[0]) + + +def load_lance_extension(con: duckdb.DuckDBPyConnection) -> bool: + """``LOAD`` the lance extension when it is installed; never install it. + + Returns ``False`` when the extension is absent, and executes no + ``INSTALL`` in that case whatever ``autoinstall_known_extensions`` says, + because the installed check runs first. Registration at query time must + never reach the network: the extension is 242 MB, fetched by name from + the default repository, and ``query`` runs unattended. + """ + if not lance_extension_installed(con): + return False + con.execute(f"LOAD {LANCE_EXTENSION};") + return True + + +def install_lance_extension(con: duckdb.DuckDBPyConnection) -> str: + """``INSTALL`` then ``LOAD`` the lance extension; return its install path. + + The one place atif-duck reaches the network, so it belongs to explicit + commands only (``atif-sql embed --install-extension``, or a real embed + run, which reaches Bedrock anyway). A query-time registration calls + :func:`load_lance_extension` instead. + """ + con.execute(f"INSTALL {LANCE_EXTENSION};") + con.execute(f"LOAD {LANCE_EXTENSION};") + row = con.execute( + "SELECT install_path FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION] + ).fetchone() + return str(row[0]) if row and row[0] is not None else "" + + def register_vss( con: duckdb.DuckDBPyConnection, *, @@ -1198,7 +1282,14 @@ def register_vss( LanceDB stores embeddings + its IVF_HNSW_SQ index in one place (written by atif-embed's backfill); reads come back through DuckDB via the lance - core extension (``INSTALL lance; LOAD lance; ATTACH (TYPE LANCE)``). + core extension (``LOAD lance; ATTACH (TYPE LANCE)``). The extension is + only LOADed, never INSTALLed, here: :func:`load_lance_extension` checks + the local extension directory first and a missing extension degrades to + the empty fallback table with a warning, so registration never reaches + the network. Installing is an explicit act + (:func:`install_lance_extension`, behind ``atif-sql embed``). When no + store directory exists the extension is not loaded at all, which is the + state of every corpus that has not run ``embed``. The store probe runs through DuckDB itself. atif-duck declares no lancedb dependency: lancedb belongs to atif-embed, which writes the store, and the independence contract forbids an import edge between the two packages — so @@ -1231,18 +1322,25 @@ def register_vss( ------- bool ``True`` when the Lance table is reachable through the - ``message_embeddings`` view; ``False`` when no embeddings exist yet - (the name is created as an empty TABLE with the right schema so a + ``message_embeddings`` view; ``False`` when no embeddings exist yet, + or the store exists but the lance extension is not installed (the + name is created as an empty TABLE with the right schema so a downstream ``CREATE MACRO semantic_search`` can still bind). """ dim_i = int(dim) - con.execute("INSTALL lance;") - con.execute("LOAD lance;") import duckdb as _duckdb attached = False - if lance_uri.is_dir(): + if lance_uri.is_dir() and not load_lance_extension(con): + logger.warning( + "Lance store at {} cannot be read: the {} extension is not installed, so " + "message_embeddings binds empty. Run `atif-sql embed --install-extension` " + "(a one-time download) to enable vector search.", + lance_uri, + LANCE_EXTENSION, + ) + elif lance_uri.is_dir(): try: # ATTACH is one of the statement kinds DuckDB will not prepare, so # this is the one corpus path that still enters statement text; it diff --git a/packages/atif-duck/tests/duck_fixtures.py b/packages/atif-duck/tests/duck_fixtures.py index 001cd5c..e8f11b5 100644 --- a/packages/atif-duck/tests/duck_fixtures.py +++ b/packages/atif-duck/tests/duck_fixtures.py @@ -707,3 +707,21 @@ def write_codex_session(root: Path) -> Path: ) ) return session_dir + + +# ``register_vss`` no longer installs the lance extension (that download at +# query time was finding 4 of the MicroVM review), so the suite installs it +# once up front. A no-op where it is already present; a one-time download on a +# fresh runner, exactly what every register() call used to do implicitly. +# A public name on purpose: ``conftest.py`` re-exports this module with +# ``import *``, which skips underscore names, so an underscore here would leave +# the fixture unregistered (it did, on a runner with no extension cached). +@pytest.fixture(scope="session", autouse=True) +def lance_extension_present() -> None: + import duckdb + + con = duckdb.connect() + try: + con.execute("INSTALL lance") + finally: + con.close() diff --git a/packages/atif-duck/tests/test_registration_limits.py b/packages/atif-duck/tests/test_registration_limits.py new file mode 100644 index 0000000..ee7bc3b --- /dev/null +++ b/packages/atif-duck/tests/test_registration_limits.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Registration fits the host, and never reaches the network. + +Two findings of the MicroVM review, pinned at the registry layer: + +* The eager ``read_json`` readers reserve about twice their + ``maximum_object_size`` per thread. At the old 1 GiB constant that was + 2 GiB a thread, so ``SET threads=4; SET memory_limit='6GB'`` (a 4 vCPU, + 8 GiB guest) failed to register a 131 MB corpus with an out-of-memory + error, and so did 16 threads under 25 GB. The bound is now sized from the + largest file the reader will open. +* ``register_vss`` used to ``INSTALL lance`` on every registration, a 242 MB + download from the extension repository. It now LOADs the extension only + when it is already installed and binds the vector surface empty otherwise. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import duckdb +import pytest +from duck_fixtures import SESSION_IDS +from loguru import logger +from test_vss import _write_lance + +from atif_duck.infrastructure import registry as registry_mod +from atif_duck.infrastructure.registry import ( + _OBJECT_SIZE_CAP, + _OBJECT_SIZE_FLOOR, + _object_size_bound, + install_lance_extension, + lance_extension_installed, + load_lance_extension, + register, + register_vss, +) + +MIB = 1024**2 + + +class TestRegistrationUnderHostLimits: + """The review's reproduction: limits set before ``register`` over the fixture corpus.""" + + @pytest.mark.parametrize(("threads", "memory"), [(4, "6GB"), (16, "8GB"), (16, "25GB")]) + def test_register_succeeds_under_limits_that_used_to_oom( + self, corpus_root: Path, threads: int, memory: str + ) -> None: + con = duckdb.connect() + try: + con.execute(f"SET threads={threads}; SET memory_limit='{memory}'") + sources = register(con, corpus_root) + assert set(sources.json_session_ids) == set(SESSION_IDS) + assert con.execute("SELECT count(*) FROM sessions").fetchone() == (2,) + messages = con.execute("SELECT count(*) FROM messages").fetchone() + assert messages is not None + assert int(messages[0]) > 0 + finally: + con.close() + + +class TestObjectSizeBound: + def test_no_files_gives_duckdbs_default(self) -> None: + assert _object_size_bound([]) == _OBJECT_SIZE_FLOOR == 16 * MIB + + def test_small_files_stay_at_the_floor(self, tmp_path: Path) -> None: + small = tmp_path / "edges.jsonl" + small.write_bytes(b"x" * 1000) + assert _object_size_bound([small]) == _OBJECT_SIZE_FLOOR + + def test_a_large_file_gets_a_quarter_plus_one_mib_of_headroom(self, tmp_path: Path) -> None: + big = tmp_path / "trajectory.json" + with big.open("wb") as handle: + handle.truncate(200 * MIB) # sparse: no bytes written + assert _object_size_bound([big]) == 200 * MIB + 50 * MIB + MIB + + def test_the_largest_file_wins(self, tmp_path: Path) -> None: + paths = [] + for i, size in enumerate((40 * MIB, 120 * MIB, 80 * MIB)): + path = tmp_path / f"{i}.json" + with path.open("wb") as handle: + handle.truncate(size) + paths.append(path) + assert _object_size_bound(paths) == 120 * MIB + 30 * MIB + MIB + + def test_capped_at_one_gib(self, tmp_path: Path) -> None: + huge = tmp_path / "huge.json" + with huge.open("wb") as handle: + handle.truncate(2 * 1024**3) + assert _object_size_bound([huge]) == _OBJECT_SIZE_CAP == 1024**3 + + def test_a_vanished_path_counts_as_zero(self, tmp_path: Path) -> None: + assert _object_size_bound([tmp_path / "gone.json"]) == _OBJECT_SIZE_FLOOR + + def test_the_bound_reaches_both_readers( + self, corpus_root: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The DDL carries the sized bound, not the old constant.""" + seen: list[int] = [] + real = registry_mod._object_size_bound + + def spy(paths: Any) -> int: + bound = real(paths) + seen.append(bound) + return bound + + monkeypatch.setattr(registry_mod, "_object_size_bound", spy) + con = duckdb.connect() + try: + register(con, corpus_root) + finally: + con.close() + assert len(seen) == 2, "one bound per eager JSON reader (trajectories, edges)" + assert all(bound == _OBJECT_SIZE_FLOOR for bound in seen) + + +class _Recording: + def __init__(self, con: duckdb.DuckDBPyConnection, statements: list[str]) -> None: + self._con = con + self._statements = statements + + def execute(self, sql: str, *args: Any, **kwargs: Any) -> Any: + self._statements.append(sql) + return self._con.execute(sql, *args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._con, name) + + +def _connection_without_extensions(tmp_path: Path) -> tuple[Any, Path, list[str]]: + ext_dir = tmp_path / "no-extensions" + ext_dir.mkdir() + statements: list[str] = [] + con = duckdb.connect() + con.execute(f"SET extension_directory='{ext_dir}'") + return _Recording(con, statements), ext_dir, statements + + +def _installs(statements: list[str]) -> list[str]: + return [s for s in statements if s.lstrip().upper().startswith("INSTALL")] + + +def _loads(statements: list[str]) -> list[str]: + return [s for s in statements if s.lstrip().upper().startswith("LOAD")] + + +class TestVssNeverInstalls: + def test_store_present_extension_absent_binds_empty_and_installs_nothing( + self, tmp_path: Path + ) -> None: + store = _write_lance(tmp_path / "store") + con, ext_dir, statements = _connection_without_extensions(tmp_path) + warnings: list[str] = [] + sink_id = logger.add(lambda message: warnings.append(str(message)), level="WARNING") + try: + bound = register_vss(con, lance_uri=store, expected_model="test-embedder:1") + finally: + logger.remove(sink_id) + con.close() + assert bound is False + assert _installs(statements) == [] + assert _loads(statements) == [] + assert list(ext_dir.rglob("*")) == [] + assert any("--install-extension" in w for w in warnings) + + def test_autoinstall_left_on_still_downloads_nothing(self, tmp_path: Path) -> None: + """The installed check runs before LOAD, so DuckDB's default autoinstall never triggers.""" + store = _write_lance(tmp_path / "store") + con, ext_dir, _ = _connection_without_extensions(tmp_path) + assert con.execute("SELECT current_setting('autoinstall_known_extensions')").fetchone() == ( + True, + ) + try: + register_vss(con, lance_uri=store) + assert con.execute("SELECT count(*) FROM message_embeddings").fetchone() == (0,) + finally: + con.close() + assert list(ext_dir.rglob("*")) == [] + + def test_no_store_means_the_extension_is_not_even_loaded(self, tmp_path: Path) -> None: + con, _, statements = _connection_without_extensions(tmp_path) + try: + assert register_vss(con, lance_uri=tmp_path / "absent") is False + finally: + con.close() + assert _loads(statements) == [] + assert _installs(statements) == [] + + def test_installed_check_reads_the_extension_directory(self, tmp_path: Path) -> None: + con, _, _ = _connection_without_extensions(tmp_path) + try: + assert lance_extension_installed(con) is False + assert load_lance_extension(con) is False + finally: + con.close() + default = duckdb.connect() + try: + assert lance_extension_installed(default) is True + assert load_lance_extension(default) is True + finally: + default.close() + + def test_register_over_a_corpus_with_a_store_still_binds_every_other_view( + self, corpus_root: Path, tmp_path: Path + ) -> None: + _write_lance(corpus_root / "embeddings_lance") + con, _, statements = _connection_without_extensions(tmp_path) + try: + register(con, corpus_root) + assert con.execute("SELECT count(*) FROM sessions").fetchone() == (2,) + assert con.execute("SELECT count(*) FROM message_embeddings").fetchone() == (0,) + assert ( + con.execute( + "SELECT count(*) FROM duckdb_functions() WHERE function_name = 'semantic_search'" + ).fetchone()[0] + >= 1 + ) + finally: + con.close() + assert _installs(statements) == [] + + def test_explicit_install_returns_the_path(self) -> None: + """A no-op where the extension is present; the one INSTALL atif-duck still owns.""" + con = duckdb.connect() + try: + path = install_lance_extension(con) + finally: + con.close() + assert "lance" in path + + def test_registry_source_has_no_install_outside_the_explicit_helper(self) -> None: + import inspect + + source = inspect.getsource(registry_mod) + helper = inspect.getsource(install_lance_extension) + assert source.count('"INSTALL') + source.count("INSTALL {LANCE_EXTENSION}") == ( + helper.count('"INSTALL') + helper.count("INSTALL {LANCE_EXTENSION}") + ) diff --git a/packages/atif-duck/tests/test_sql_text_boundaries.py b/packages/atif-duck/tests/test_sql_text_boundaries.py index f97c091..9e5dc2d 100644 --- a/packages/atif-duck/tests/test_sql_text_boundaries.py +++ b/packages/atif-duck/tests/test_sql_text_boundaries.py @@ -335,6 +335,53 @@ def test_a_bad_dir_with_columnar_artifacts_is_still_rejected( assert all(bad not in str(path) for path in sources.lazy_read_paths) +class TestGlobShapedNamesDoNotMultiplyRows: + """Finding 7 of the MicroVM review, exactly as reproduced there. + + Before the boundary a session dir named ``*`` or ``???`` landed inside + the readers' path lists, DuckDB expanded it as a glob, and every other + session's rows were read twice: ``sessions`` had seven rows for four + ids and ``count(*) FROM steps`` was 14 where 8 was true. The name is now + rejected at the boundary, so a registration over such a corpus counts + each session exactly once, on both read paths, and grants no glob. + """ + + @staticmethod + def _per_session_counts(con: duckdb.DuckDBPyConnection) -> dict[str, tuple[int, int]]: + rows = _rows( + con, + "SELECT s.session_id, s.step_count, (SELECT count(*) FROM steps t " + "WHERE t.session_id = s.session_id) FROM sessions s ORDER BY 1", + ) + return {str(r[0]): (int(r[1]), int(r[2])) for r in rows} + + @pytest.mark.parametrize("name", ["*", "???"]) + def test_each_session_is_counted_once_on_both_paths(self, tmp_path: Path, name: str) -> None: + clean = duckdb.connect(":memory:") + register(clean, build_corpus(tmp_path / "clean")) + expected = self._per_session_counts(clean) + assert set(expected) == set(SESSION_IDS) + + root = build_corpus(tmp_path / "hostile") + _write_session( + root / "sessions" / name, name, _adversarial_trajectory() | {"session_id": name} + ) + for path_label in ("json", "columnar"): + if path_label == "columnar": + add_columnar(root, tuple(SESSION_IDS)) + con = duckdb.connect(":memory:") + sources = register(con, root) + assert sources.rejected_session_ids == (name,), path_label + assert self._per_session_counts(con) == expected, path_label + assert _rows(con, "SELECT count(*) FROM sessions") == [(len(SESSION_IDS),)] + assert _rows(con, "SELECT count(*) FROM steps") == [ + (sum(steps for steps, _ in expected.values()),) + ] + assert all(not any(ch in str(p) for ch in "*?[") for p in sources.lazy_read_paths), ( + "a glob metacharacter reached the sandbox's file allowlist" + ) + + class TestRemainingSqlLiteralSites: """The two statements DuckDB will not prepare; each fails without ``sql_literal``."""