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`.
- 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

View file

@ -1,9 +1,9 @@
import asyncio
import atexit
import concurrent.futures
import json
import os
from collections.abc import Callable
import threading
import weakref
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from pathlib import Path
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:
from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@dataclass
class SandboxResult:
@ -29,13 +31,13 @@ class SandboxResult:
success: bool
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
atexit.register(_executor.shutdown, wait=False)
def _run_async(coro: Any) -> Any:
"""Run an async coroutine from a sync context (CallbackFile read)."""
return _executor.submit(asyncio.run, coro).result()
def _stop_loop(loop: asyncio.AbstractEventLoop, thread: threading.Thread) -> None:
"""Stop a background loop and join its thread. Safe if already stopped."""
if loop.is_closed():
return
loop.call_soon_threadsafe(loop.stop)
thread.join()
loop.close()
def _build_toc(
@ -127,6 +129,13 @@ class Sandbox:
sandbox = Sandbox(db_path, config, context)
result = await sandbox.execute("x = await search('query')")
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
@ -139,6 +148,10 @@ class Sandbox:
_toc_json_cache: dict[str, str]
_repl: MontyRepl | None
_vfs: OSAccess | None
_loop: asyncio.AbstractEventLoop | None
_loop_thread: threading.Thread | None
_vfs_rag: "HaikuRAG | None"
_finalizer: weakref.finalize | None
def __init__(
self,
@ -156,6 +169,60 @@ class Sandbox:
self._toc_json_cache = {}
self._repl = 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]:
"""Build async external functions for the Monty interpreter."""
@ -249,12 +316,12 @@ class Sandbox:
return cached
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:
return await rag.document_item_repository.get_all_items(did)
items = _run_async(_fetch())
items = sandbox._run_on_loop(_fetch())
sandbox._doc_items[did] = items
return items
@ -265,17 +332,13 @@ class Sandbox:
return cached
async def _fetch() -> dict[str, list[str]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
index = (
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
[did]
)
)
assert sandbox._vfs_rag is not None
index = await sandbox._vfs_rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
[did]
)
return index.get(did, {})
chunk_index = _run_async(_fetch())
chunk_index = sandbox._run_on_loop(_fetch())
sandbox._doc_chunk_index[did] = chunk_index
return chunk_index
@ -351,15 +414,13 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> 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(
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 sandbox._run_on_loop(_fetch())
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:
deps.rag = rag
deps.search_count = 0
deps.sandbox = Sandbox(
sandbox = Sandbox(
db_path=db_path,
config=config,
context=AnalysisContext(filter=doc_filter),
)
deps.sandbox = sandbox
_reset_invocation_state(deps.state)
yield
try:
yield
finally:
sandbox.close()
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 "Public Doc" 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()