Fix analysis sandbox "Already borrowed" crash on VFS reads

This commit is contained in:
Yiorgis Gozadinos 2026-06-05 11:49:06 +03:00
parent a2def043cc
commit 7906688bf1
No known key found for this signature in database
8 changed files with 401 additions and 34 deletions

View file

@ -9,6 +9,7 @@
- Opening a database no longer writes to it: reads no longer rewrite the stored embedding settings or change the stored version, and the version is never downgraded. Embedding provider/name drift (matching `vector_dim`) warns on read-only opens and raises `ConfigMismatchError` on writable opens; reconcile with `rebuild --set-embedder`. - Opening a database no longer writes to it: reads no longer rewrite the stored embedding settings or change the stored version, and the version is never downgraded. Embedding provider/name drift (matching `vector_dim`) warns on read-only opens and raises `ConfigMismatchError` on writable opens; reconcile with `rebuild --set-embedder`.
- Read CLI verbs (`list`, `get`, `search`, `visualize`, `ask`, `analyze`, `inspect`, `chat`, `info`, `history`) open the database read-only. - Read CLI verbs (`list`, `get`, `search`, `visualize`, `ask`, `analyze`, `inspect`, `chat`, `info`, `history`) open the database read-only.
- Analysis-skill `execute_code` no longer crashes with LanceDB `Already borrowed` when sandboxed code reads the document VFS (`content.txt`, `items.jsonl`, `toc.json`); reads run on one background event loop with a single read-only connection for the sandbox's lifetime.
## [0.54.0] - 2026-06-04 ## [0.54.0] - 2026-06-04

View file

@ -1,9 +1,9 @@
import asyncio import asyncio
import atexit
import concurrent.futures
import json import json
import os import os
from collections.abc import Callable import threading
import weakref
from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal from typing import TYPE_CHECKING, Any, Literal
@ -19,6 +19,8 @@ from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentIte
if TYPE_CHECKING: if TYPE_CHECKING:
from pathlib import PurePosixPath from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@dataclass @dataclass
class SandboxResult: class SandboxResult:
@ -29,13 +31,13 @@ class SandboxResult:
success: bool success: bool
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) def _stop_loop(loop: asyncio.AbstractEventLoop, thread: threading.Thread) -> None:
atexit.register(_executor.shutdown, wait=False) """Stop a background loop and join its thread. Safe if already stopped."""
if loop.is_closed():
return
def _run_async(coro: Any) -> Any: loop.call_soon_threadsafe(loop.stop)
"""Run an async coroutine from a sync context (CallbackFile read).""" thread.join()
return _executor.submit(asyncio.run, coro).result() loop.close()
def _build_toc( def _build_toc(
@ -127,6 +129,13 @@ class Sandbox:
sandbox = Sandbox(db_path, config, context) sandbox = Sandbox(db_path, config, context)
result = await sandbox.execute("x = await search('query')") result = await sandbox.execute("x = await search('query')")
result = await sandbox.execute("print(x[0]['content'])") # x persists result = await sandbox.execute("print(x[0]['content'])") # x persists
The virtual filesystem reads run on a dedicated background event loop with a
single read-only connection held for the sandbox's lifetime. Monty's file
callbacks are synchronous and cannot ``await``, so they bridge to that loop
via ``run_coroutine_threadsafe``. Call ``close()`` to release the loop and
connection; if it is not called they are released when the instance is
garbage collected.
""" """
_db_path: Path _db_path: Path
@ -139,6 +148,10 @@ class Sandbox:
_toc_json_cache: dict[str, str] _toc_json_cache: dict[str, str]
_repl: MontyRepl | None _repl: MontyRepl | None
_vfs: OSAccess | None _vfs: OSAccess | None
_loop: asyncio.AbstractEventLoop | None
_loop_thread: threading.Thread | None
_vfs_rag: "HaikuRAG | None"
_finalizer: weakref.finalize | None
def __init__( def __init__(
self, self,
@ -156,6 +169,60 @@ class Sandbox:
self._toc_json_cache = {} self._toc_json_cache = {}
self._repl = None self._repl = None
self._vfs = None self._vfs = None
self._loop = None
self._loop_thread = None
self._vfs_rag = None
self._finalizer = None
def _ensure_vfs_loop(self) -> None:
"""Start the background loop and open the read-only connection once."""
if self._loop is not None:
return
from haiku.rag.client import HaikuRAG
loop = asyncio.new_event_loop()
thread = threading.Thread(
target=loop.run_forever, daemon=True, name="sandbox-vfs"
)
thread.start()
self._loop = loop
self._loop_thread = thread
async def _open() -> "HaikuRAG":
rag = HaikuRAG(self._db_path, config=self._config, read_only=True)
await rag.__aenter__()
return rag
self._vfs_rag = self._run_on_loop(_open())
# GC backstop: capture only the loop + thread, never ``self`` (which
# would pin the instance and prevent collection).
self._finalizer = weakref.finalize(self, _stop_loop, loop, thread)
def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
"""Run a coroutine on the background loop from a synchronous caller."""
self._ensure_vfs_loop()
assert self._loop is not None
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
def close(self) -> None:
"""Release the background loop and read-only connection.
Idempotent and safe to call when no VFS read ever started the loop.
"""
loop, thread, rag = self._loop, self._loop_thread, self._vfs_rag
if loop is None or thread is None:
return
self._loop = None
self._loop_thread = None
self._vfs_rag = None
if self._finalizer is not None:
self._finalizer.detach()
self._finalizer = None
if rag is not None and not loop.is_closed():
asyncio.run_coroutine_threadsafe(
rag.__aexit__(None, None, None), loop
).result()
_stop_loop(loop, thread)
def _build_external_functions(self) -> dict[str, Any]: def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter.""" """Build async external functions for the Monty interpreter."""
@ -249,12 +316,12 @@ class Sandbox:
return cached return cached
async def _fetch() -> list[DocumentItem]: async def _fetch() -> list[DocumentItem]:
from haiku.rag.client import HaikuRAG assert sandbox._vfs_rag is not None
return await sandbox._vfs_rag.document_item_repository.get_all_items(
did
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag: items = sandbox._run_on_loop(_fetch())
return await rag.document_item_repository.get_all_items(did)
items = _run_async(_fetch())
sandbox._doc_items[did] = items sandbox._doc_items[did] = items
return items return items
@ -265,17 +332,13 @@ class Sandbox:
return cached return cached
async def _fetch() -> dict[str, list[str]]: async def _fetch() -> dict[str, list[str]]:
from haiku.rag.client import HaikuRAG assert sandbox._vfs_rag is not None
index = await sandbox._vfs_rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
async with HaikuRAG(db_path, config=config, read_only=True) as rag: [did]
index = ( )
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
[did]
)
)
return index.get(did, {}) return index.get(did, {})
chunk_index = _run_async(_fetch()) chunk_index = sandbox._run_on_loop(_fetch())
sandbox._doc_chunk_index[did] = chunk_index sandbox._doc_chunk_index[did] = chunk_index
return chunk_index return chunk_index
@ -351,15 +414,13 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]: ) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str: def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str: async def _fetch() -> str:
from haiku.rag.client import HaikuRAG assert sandbox._vfs_rag is not None
content = (
await sandbox._vfs_rag.document_repository.get_content(did)
)
return content or ""
async with HaikuRAG( return sandbox._run_on_loop(_fetch())
db_path, config=config, read_only=True
) as rag:
content = await rag.document_repository.get_content(did)
return content or ""
return _run_async(_fetch())
return read_content return read_content

View file

@ -67,12 +67,16 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
async with HaikuRAG(db_path, config=config, read_only=True) as rag: async with HaikuRAG(db_path, config=config, read_only=True) as rag:
deps.rag = rag deps.rag = rag
deps.search_count = 0 deps.search_count = 0
deps.sandbox = Sandbox( sandbox = Sandbox(
db_path=db_path, db_path=db_path,
config=config, config=config,
context=AnalysisContext(filter=doc_filter), context=AnalysisContext(filter=doc_filter),
) )
deps.sandbox = sandbox
_reset_invocation_state(deps.state) _reset_invocation_state(deps.state)
yield try:
yield
finally:
sandbox.close()
return lifespan return lifespan

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -412,3 +412,136 @@ class TestSandboxVFS:
assert "1" in result.stdout assert "1" in result.stdout
assert "Public Doc" in result.stdout assert "Public Doc" in result.stdout
assert "Private Doc" not in result.stdout assert "Private Doc" not in result.stdout
class TestSandboxHeldConnection:
"""VFS reads must work while another connection to the same DB stays open.
Mirrors the analysis-skill lifespan, which keeps a read-only connection open
for the whole turn while sandboxed code reads the document VFS.
"""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_reads_while_connection_held(self, temp_db_path):
"""All three VFS readers work while another connection to the DB is open."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs roam the quiet hills.",
uri="test://doc",
title="Doc",
)
assert doc.id
async with HaikuRAG(temp_db_path, config=config, read_only=True):
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
result = await sb.execute(
"from pathlib import Path\n"
f"d = Path('/documents/{doc.id}')\n"
"print('foxes' in (d / 'content.txt').read_text().lower())\n"
"print(len((d / 'items.jsonl').read_text().strip().split('\\n')))\n"
"print('tree' in (d / 'toc.json').read_text())"
)
assert result.success, result.stderr
lines = result.stdout.strip().split("\n")
assert lines[0] == "True"
assert int(lines[1]) > 0
assert lines[2] == "True"
finally:
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_loop_and_connection_are_reused(self, temp_db_path):
"""The background loop + read-only connection are created once and reused."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
assert sb._loop is None and sb._vfs_rag is None
read = (
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert (await sb.execute(read)).success
loop, rag = sb._loop, sb._vfs_rag
assert loop is not None and rag is not None
assert (await sb.execute(read)).success
assert sb._loop is loop
assert sb._vfs_rag is rag
finally:
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_variables_persist_across_executes_after_vfs_read(self, temp_db_path):
"""REPL state persists across execute() calls, including after a VFS read."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
first = await sb.execute(
"from pathlib import Path\n"
f"x = len(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert first.success, first.stderr
second = await sb.execute("print(x)")
assert second.success, second.stderr
assert int(second.stdout.strip()) > 0
finally:
sb.close()
@pytest.mark.asyncio
async def test_close_is_safe_without_vfs_read(self, temp_db_path):
"""close() is a no-op (and safe to call twice) when no VFS read happened."""
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
assert sb._loop is None
sb.close()
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_close_stops_loop_thread(self, temp_db_path):
"""close() tears down the background loop thread and is idempotent."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert result.success
thread = sb._loop_thread
assert thread is not None and thread.is_alive()
sb.close()
assert not thread.is_alive()
sb.close()