b374d5eb83
In Claude Code the client is the model, so the server no longer runs one. execute_code runs a Python program per call in the analysis sandbox over the selected documents and returns what it printed; the sandbox is created and closed per call so Monty's cumulative budget and a frozen mount never outlive a program. --no-agents goes with the two tools, and format_citations in haiku.rag.utils goes with its only caller. The sandbox exposes chunk metadata to code: chunk_meta on search results, metadata on list_documents rows and in metadata.json, and chunks.jsonl per document. A host-side failure inside a program, a document read or an in-code search raising, reaches the program by exception type only and is logged with its traceback. recovery_hint moves to haiku.rag.sandbox. Closes #604.
1185 lines
45 KiB
Python
1185 lines
45 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.sandbox import AnalysisContext, Sandbox, SandboxResult
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_cassette_dir():
|
|
return str(Path(__file__).parent.parent / "cassettes" / "test_sandbox")
|
|
|
|
|
|
class TestSandboxBasics:
|
|
"""Test basic sandbox functionality."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_simple_code(self, sandbox):
|
|
"""Test executing simple code in the sandbox."""
|
|
result = await sandbox.execute("print('hello world')")
|
|
assert isinstance(result, SandboxResult)
|
|
assert result.success
|
|
assert "hello world" in result.stdout
|
|
assert result.stderr == ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_expression_output(self, sandbox):
|
|
"""Test that expression values are captured."""
|
|
result = await sandbox.execute("1 + 2")
|
|
assert result.success
|
|
assert "3" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_print_and_expression(self, sandbox):
|
|
"""Test print output combined with expression value."""
|
|
result = await sandbox.execute("print('hello')\n42")
|
|
assert result.success
|
|
assert "hello" in result.stdout
|
|
assert "42" in result.stdout
|
|
|
|
|
|
class TestSandboxErrors:
|
|
"""Test error handling in sandbox."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_syntax_error(self, sandbox):
|
|
"""Test that syntax errors are reported."""
|
|
result = await sandbox.execute("def foo(")
|
|
assert not result.success
|
|
assert result.stderr != ""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_error(self, sandbox):
|
|
"""Test that runtime errors are reported."""
|
|
result = await sandbox.execute("x = 1/0")
|
|
assert not result.success
|
|
assert "ZeroDivisionError" in result.stderr
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_name_error(self, sandbox):
|
|
"""Test that name errors are reported."""
|
|
result = await sandbox.execute("print(undefined_variable)")
|
|
assert not result.success
|
|
assert "NameError" in result.stderr
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_module_import(self, sandbox):
|
|
"""Test that unsupported multi-module imports are caught gracefully."""
|
|
result = await sandbox.execute("import json, string")
|
|
assert not result.success
|
|
assert result.stderr != ""
|
|
|
|
|
|
class TestSandboxListDocuments:
|
|
"""Test list_documents function in sandbox."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_documents_empty(self, sandbox):
|
|
"""list_documents returns empty list for empty database."""
|
|
result = await sandbox.execute(
|
|
"docs = await list_documents()\nprint(type(docs).__name__, len(docs))"
|
|
)
|
|
assert result.success
|
|
assert "list 0" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_list_documents_with_data(self, temp_db_path):
|
|
"""list_documents returns documents when populated."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="Test content",
|
|
uri="test://doc1",
|
|
title="Test Document",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"docs = await list_documents()\n"
|
|
"print(len(docs))\n"
|
|
"print(docs[0]['title'])\n"
|
|
"print(docs[0]['source'])"
|
|
)
|
|
assert result.success
|
|
assert "1" in result.stdout
|
|
assert "Test Document" in result.stdout
|
|
assert temp_db_path.stem in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_documents_carries_metadata(self, temp_db_path):
|
|
"""Rows carry the document's metadata, so a corpus-wide pass over it is
|
|
one call rather than a file read per document."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Test content")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Test content",
|
|
embedding=[0.1] * config.embeddings.model.vector_dim,
|
|
order=0,
|
|
)
|
|
],
|
|
uri="test://doc1",
|
|
title="Test Document",
|
|
metadata={"author": "Ada"},
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"docs = await list_documents()\nprint(docs[0]['metadata']['author'])"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
assert result.success, result.stderr
|
|
assert "Ada" in result.stdout
|
|
|
|
|
|
class TestSandboxSearch:
|
|
"""Test search function in sandbox."""
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_with_data(self, temp_db_path):
|
|
"""Test search function works."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="The quick brown fox jumps over the lazy dog.",
|
|
uri="test://animals",
|
|
title="Animals",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"results = await search('fox', limit=5)\n"
|
|
"print(len(results))\n"
|
|
"if results:\n"
|
|
" print('fox' in results[0]['content'].lower())"
|
|
)
|
|
assert result.success
|
|
assert "True" in result.stdout or "1" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_returns_doc_item_refs_and_labels(self, temp_db_path):
|
|
"""Search results include doc_item_refs and labels."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="The quick brown fox jumps over the lazy dog.",
|
|
uri="test://animals",
|
|
title="Animals",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"results = await search('fox', limit=1)\n"
|
|
"r = results[0]\n"
|
|
"print('doc_item_refs' in r)\n"
|
|
"print('labels' in r)\n"
|
|
"print(type(r['doc_item_refs']).__name__)\n"
|
|
"print(type(r['labels']).__name__)"
|
|
)
|
|
assert result.success
|
|
assert "True\nTrue" in result.stdout
|
|
assert "list\nlist" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_returns_expanded_content(self, temp_db_path):
|
|
"""search() returns context-expanded results."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="The quick brown fox jumps over the lazy dog.",
|
|
uri="test://animals",
|
|
title="Animals",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"results = await search('fox', limit=1)\n"
|
|
"print(type(results[0]['content']).__name__)\n"
|
|
"print('fox' in results[0]['content'].lower())"
|
|
)
|
|
assert result.success
|
|
assert "str" in result.stdout
|
|
assert "True" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_returns_the_matched_chunks_metadata(
|
|
self, temp_db_path, monkeypatch
|
|
):
|
|
"""Results carry the stored metadata of the chunk that matched, custom
|
|
keys included."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
from haiku.rag.embeddings import EmbedderWrapper
|
|
|
|
config = AppConfig()
|
|
dim = config.embeddings.model.vector_dim
|
|
|
|
async def embed_query(self, text):
|
|
return [0.1] * dim
|
|
|
|
monkeypatch.setattr(EmbedderWrapper, "embed_query", embed_query)
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Paragraph fourteen.")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Paragraph fourteen.",
|
|
embedding=[0.1] * dim,
|
|
order=0,
|
|
metadata={"para_no": "14"},
|
|
)
|
|
],
|
|
uri="test://paras",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"results = await search('fourteen', limit=1)\n"
|
|
"print(results[0]['chunk_meta']['para_no'])"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
assert result.success, result.stderr
|
|
assert "14" in result.stdout
|
|
|
|
|
|
class TestSandboxExternalFunctionEdgeCases:
|
|
"""Test edge cases in external function dispatch."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unknown_external_function(self, sandbox):
|
|
"""Test that calling an unregistered external function resumes with KeyError."""
|
|
original_build = sandbox._build_external_functions
|
|
|
|
def patched_build():
|
|
fns = original_build()
|
|
fns["search"] = None
|
|
return fns
|
|
|
|
sandbox._build_external_functions = patched_build
|
|
|
|
result = await sandbox.execute(
|
|
"try:\n"
|
|
" await search('hello')\n"
|
|
"except:\n"
|
|
" print('caught')\n"
|
|
"print('done')"
|
|
)
|
|
assert result.success
|
|
assert "caught" in result.stdout
|
|
assert "done" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_external_function_raises_exception(self, sandbox):
|
|
"""Test that exceptions from async external functions surface as errors.
|
|
|
|
With run_monty_async, exceptions from async external functions
|
|
propagate as MontyRuntimeError rather than being catchable inside
|
|
Monty's try/except.
|
|
"""
|
|
original_build = sandbox._build_external_functions
|
|
|
|
def patched_build():
|
|
fns = original_build()
|
|
|
|
async def failing_search(*args, **kwargs):
|
|
raise ValueError("external error")
|
|
|
|
fns["search"] = failing_search
|
|
return fns
|
|
|
|
sandbox._build_external_functions = patched_build
|
|
|
|
result = await sandbox.execute("await search('hello')")
|
|
assert not result.success
|
|
assert "external error" in result.stderr
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_failing_search_reaches_the_program_by_type_only(
|
|
self, sandbox, monkeypatch, caplog
|
|
):
|
|
"""A host-side failure inside search() names its exception type to
|
|
the program; the message and traceback go to the log."""
|
|
|
|
async def boom(self, *args, **kwargs):
|
|
raise ValueError("failed at /secret/path")
|
|
|
|
monkeypatch.setattr(HaikuRAG, "search", boom)
|
|
|
|
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"):
|
|
result = await sandbox.execute("await search('hello')")
|
|
|
|
assert not result.success
|
|
assert "search() failed: ValueError" in result.stderr
|
|
assert "/secret/path" not in result.stderr
|
|
assert any(
|
|
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
|
|
for r in caplog.records
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_failing_document_read_reaches_the_program_by_type_only(
|
|
self, temp_db_path, monkeypatch, caplog
|
|
):
|
|
"""A program can catch a failed file read, and what it catches names
|
|
the exception type only."""
|
|
from haiku.rag.store.models.document import Document
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.document_repository.create(
|
|
Document(content="x", uri="test://read", title="Read")
|
|
)
|
|
repository = type(client.document_repository)
|
|
|
|
async def boom(self, *args, **kwargs):
|
|
raise ValueError("failed at /secret/path")
|
|
|
|
monkeypatch.setattr(repository, "get_content", boom)
|
|
sb = Sandbox(
|
|
db_path=temp_db_path, config=AppConfig(), context=AnalysisContext()
|
|
)
|
|
try:
|
|
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"):
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"try:\n"
|
|
f" Path('/documents/{doc.id}/content.txt').read_text()\n"
|
|
"except Exception as e:\n"
|
|
" print('caught:', e)"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
|
|
assert result.success, result.stderr
|
|
assert "caught:" in result.stdout
|
|
assert "ValueError" in result.stdout
|
|
assert "/secret/path" not in result.stdout
|
|
assert any(
|
|
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
|
|
for r in caplog.records
|
|
)
|
|
|
|
|
|
class TestSandboxOutputTruncation:
|
|
"""Test output truncation behavior."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_truncate_stdout_on_runtime_error(self, temp_db_path):
|
|
"""Test stdout is truncated when a runtime error occurs after large output."""
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
config = AppConfig()
|
|
config.analysis.max_output_chars = 20
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute("print('a' * 100)\nx = 1/0")
|
|
assert not result.success
|
|
assert "ZeroDivisionError" in result.stderr
|
|
assert result.stdout.endswith("... (output truncated)")
|
|
assert len(result.stdout) < 100
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_truncate_successful_output(self, temp_db_path):
|
|
"""Test output is truncated on successful execution with large output."""
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
config = AppConfig()
|
|
config.analysis.max_output_chars = 20
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute("print('b' * 100)")
|
|
assert result.success
|
|
assert result.stdout.endswith("... (output truncated)")
|
|
assert len(result.stdout) < 100
|
|
|
|
|
|
class TestSandboxVFS:
|
|
"""Test virtual filesystem for document access."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_database_has_no_documents(self, sandbox):
|
|
"""Empty database has no document directories."""
|
|
result = await sandbox.execute(
|
|
"from pathlib import Path\nprint(Path('/documents').exists())"
|
|
)
|
|
assert result.success
|
|
# /documents dir may or may not exist when empty, both are valid
|
|
# The key is it doesn't error
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_iterdir_discovers_documents(self, temp_db_path):
|
|
"""Path('/documents').iterdir() lists document directories."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="Test content",
|
|
uri="test://doc1",
|
|
title="Test Document",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"dirs = list(Path('/documents').iterdir())\n"
|
|
"print(len(dirs))\n"
|
|
"print(dirs[0].is_dir())"
|
|
)
|
|
assert result.success
|
|
assert "1" in result.stdout
|
|
assert "True" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_metadata_json(self, temp_db_path):
|
|
"""metadata.json contains document title, uri and stored metadata."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="Test content",
|
|
uri="test://doc1",
|
|
title="Test Document",
|
|
metadata={"author": "Ada"},
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"import json\n"
|
|
f"meta = json.loads(Path('/documents/{doc.id}/metadata.json').read_text())\n"
|
|
"print(meta['title'])\n"
|
|
"print(meta['uri'])\n"
|
|
"print(meta['metadata']['author'])"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert "Test Document" in result.stdout
|
|
assert "test://doc1" in result.stdout
|
|
assert "Ada" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_content_txt(self, temp_db_path):
|
|
"""content.txt returns full document text (lazy loaded)."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="Content about foxes and dogs.",
|
|
uri="test://doc",
|
|
title="Fox Document",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
f"content = Path('/documents/{doc.id}/content.txt').read_text()\n"
|
|
"print('foxes' in content.lower())"
|
|
)
|
|
assert result.success
|
|
assert "True" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_items_jsonl(self, temp_db_path):
|
|
"""items.jsonl returns document items as JSONL (lazy loaded)."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="The quick brown fox jumps over the lazy dog.",
|
|
uri="test://animals",
|
|
title="Animals",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"import json\n"
|
|
f"text = Path('/documents/{doc.id}/items.jsonl').read_text()\n"
|
|
"lines = text.strip().split('\\n')\n"
|
|
"print(len(lines) > 0)\n"
|
|
"item = json.loads(lines[0])\n"
|
|
"print('self_ref' in item)\n"
|
|
"print('label' in item)\n"
|
|
"print('text' in item)\n"
|
|
"print('page_numbers' in item)\n"
|
|
"print('chunk_ids' in item)"
|
|
)
|
|
assert result.success
|
|
assert result.stdout.count("True") == 6
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chunks_jsonl(self, temp_db_path):
|
|
"""chunks.jsonl lists a document's chunks in order with their stored
|
|
metadata; a chunk found by its metadata leads to its items through
|
|
their chunk_ids."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
dim = config.embeddings.model.vector_dim
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Paragraph thirteen.")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Paragraph fourteen.")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Paragraph thirteen.",
|
|
embedding=[0.1] * dim,
|
|
order=0,
|
|
metadata={"para_no": "13", "doc_item_refs": ["#/texts/0"]},
|
|
),
|
|
Chunk(
|
|
content="Paragraph fourteen.",
|
|
embedding=[0.1] * dim,
|
|
order=1,
|
|
metadata={"para_no": "14", "doc_item_refs": ["#/texts/1"]},
|
|
),
|
|
],
|
|
uri="test://paras",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"import json\n"
|
|
f"root = Path('/documents/{doc.id}')\n"
|
|
"def rows(name):\n"
|
|
" return [json.loads(l) for l in (root / name).read_text().strip().split('\\n')]\n"
|
|
"chunks = rows('chunks.jsonl')\n"
|
|
"print(len(chunks))\n"
|
|
"hit = [c for c in chunks if c['metadata'].get('para_no') == '14']\n"
|
|
"print(len(hit))\n"
|
|
"items = rows('items.jsonl')\n"
|
|
"print([i['text'] for i in items if hit[0]['chunk_id'] in i['chunk_ids']])"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
assert result.success, result.stderr
|
|
assert result.stdout.splitlines() == ["2", "1", "['Paragraph fourteen.']"]
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_open_read(self, temp_db_path):
|
|
"""open() and a with-block read document files through the VFS."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="Content about foxes and dogs.",
|
|
uri="test://doc",
|
|
title="Fox Document",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
f"with open('/documents/{doc.id}/content.txt') as f:\n"
|
|
" data = f.read()\n"
|
|
"print('foxes' in data.lower())"
|
|
)
|
|
assert result.success
|
|
assert "True" in result.stdout
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_open_readlines(self, temp_db_path):
|
|
"""readlines() splits a newline-delimited VFS file into lines."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.create_document(
|
|
content="The quick brown fox jumps over the lazy dog.",
|
|
uri="test://animals",
|
|
title="Animals",
|
|
)
|
|
|
|
context = AnalysisContext()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
f"lines = open('/documents/{doc.id}/items.jsonl').readlines()\n"
|
|
"print(len(lines) > 0)\n"
|
|
"import json\n"
|
|
"print('self_ref' in json.loads(lines[0]))"
|
|
)
|
|
assert result.success
|
|
assert result.stdout.count("True") == 2
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"filename",
|
|
["content.txt", "items.jsonl", "chunks.jsonl", "toc.json", "metadata.json"],
|
|
)
|
|
async def test_write_denied_for_every_document_file(self, temp_db_path, filename):
|
|
"""Every file in the document VFS is read-only, metadata.json included."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.")
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Foxes and dogs.",
|
|
embedding=[0.1] * config.embeddings.model.vector_dim,
|
|
order=0,
|
|
)
|
|
],
|
|
uri="test://readonly",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"try:\n"
|
|
f" Path('/documents/{doc.id}/{filename}').write_text('nope')\n"
|
|
" print('WROTE')\n"
|
|
"except PermissionError:\n"
|
|
" print('DENIED')"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert "DENIED" in result.stdout
|
|
assert "WROTE" not in result.stdout
|
|
finally:
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_open_for_writing_is_denied(self, temp_db_path):
|
|
"""`open()` in write mode is refused, not only `Path.write_text`."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.")
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Foxes and dogs.",
|
|
embedding=[0.1] * config.embeddings.model.vector_dim,
|
|
order=0,
|
|
)
|
|
],
|
|
uri="test://openwrite",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"try:\n"
|
|
f" with open('/documents/{doc.id}/content.txt', 'w') as f:\n"
|
|
" f.write('nope')\n"
|
|
" print('WROTE')\n"
|
|
"except PermissionError:\n"
|
|
" print('DENIED')"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert "DENIED" in result.stdout
|
|
assert "WROTE" not in result.stdout
|
|
finally:
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_open_file_objects_are_not_iterable(self, temp_db_path):
|
|
"""Pins the limitation the instructions warn about: pydantic/monty#490.
|
|
|
|
A failure here means Monty gained iteration support and the
|
|
`for line in f` prohibition in the analysis instructions is now wrong.
|
|
"""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="one\ntwo")
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="one\ntwo",
|
|
embedding=[0.1] * config.embeddings.model.vector_dim,
|
|
order=0,
|
|
)
|
|
],
|
|
uri="test://lines",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
f"for line in open('/documents/{doc.id}/content.txt'):\n print(line)"
|
|
)
|
|
assert result.success is False
|
|
assert "not iterable" in result.stderr
|
|
|
|
# The documented alternatives do work.
|
|
result = await sb.execute(
|
|
f"print(len(open('/documents/{doc.id}/content.txt').readlines()))"
|
|
)
|
|
assert result.success, result.stderr
|
|
finally:
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_context_filter_limits_vfs(self, temp_db_path):
|
|
"""Context filter restricts which documents appear in VFS."""
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
await client.create_document(
|
|
content="Public content",
|
|
uri="public://doc1",
|
|
title="Public Doc",
|
|
)
|
|
await client.create_document(
|
|
content="Private content",
|
|
uri="private://doc2",
|
|
title="Private Doc",
|
|
)
|
|
|
|
context = AnalysisContext(filter="uri LIKE 'public://%'")
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
"import json\n"
|
|
"dirs = list(Path('/documents').iterdir())\n"
|
|
"print(len(dirs))\n"
|
|
"meta = json.loads((dirs[0] / 'metadata.json').read_text())\n"
|
|
"print(meta['title'])"
|
|
)
|
|
assert result.success
|
|
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-capability 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:
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_vfs_reads_use_injected_connection(self, temp_db_path):
|
|
"""An injected connection services VFS reads on the calling loop.
|
|
|
|
No dedicated background loop/thread is spawned: all DB access for the
|
|
sandbox runs on the loop driving execute(), through the one connection.
|
|
"""
|
|
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) as rag:
|
|
sb = Sandbox(
|
|
db_path=temp_db_path,
|
|
config=config,
|
|
context=AnalysisContext(),
|
|
rag=rag,
|
|
)
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert "Foxes" in result.stdout
|
|
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_vfs_read_concurrent_with_connection_read(self, temp_db_path):
|
|
"""A VFS read and a direct read on the shared connection run together.
|
|
|
|
Both serialize through the shared lock, so concurrent tasks never have
|
|
two operations in flight on the one connection at once.
|
|
"""
|
|
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) as rag:
|
|
lock = asyncio.Lock()
|
|
sb = Sandbox(
|
|
db_path=temp_db_path,
|
|
config=config,
|
|
context=AnalysisContext(),
|
|
rag=rag,
|
|
lock=lock,
|
|
)
|
|
read_code = (
|
|
"from pathlib import Path\n"
|
|
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
|
|
)
|
|
|
|
async def direct_read() -> str | None:
|
|
async with lock:
|
|
return await rag.document_repository.get_content(doc.id)
|
|
|
|
exec_result, content = await asyncio.gather(
|
|
sb.execute(read_code),
|
|
direct_read(),
|
|
)
|
|
assert exec_result.success, exec_result.stderr
|
|
assert "Foxes" in exec_result.stdout
|
|
assert content and "Foxes" in content
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_vfs_reads_repeatable_across_executes(self, temp_db_path):
|
|
"""Repeated VFS reads across execute() calls return consistent content."""
|
|
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:
|
|
read = (
|
|
"from pathlib import Path\n"
|
|
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
|
|
)
|
|
first = await sb.execute(read)
|
|
second = await sb.execute(read)
|
|
assert first.success and second.success
|
|
assert "Foxes" in first.stdout
|
|
assert first.stdout == second.stdout
|
|
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
|
|
finally:
|
|
await 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:
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_is_safe_without_vfs_read(self, temp_db_path):
|
|
"""close() is safe and idempotent before any code has run."""
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
config = AppConfig()
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
await sb.close()
|
|
await sb.close()
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_close_is_idempotent_after_vfs_read(self, temp_db_path):
|
|
"""close() tears down the worker and is idempotent; no thread lingers."""
|
|
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
|
|
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
|
|
await sb.close()
|
|
await sb.close()
|
|
|
|
|
|
class TestSandboxReadDeadline:
|
|
"""The VFS bridge suspends the worker for the length of a read, so Monty
|
|
cannot check its duration budget while one is in flight. The sandbox
|
|
enforces the budget itself, before each read."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_failed_read_reaches_the_program_by_type_only(
|
|
self, sandbox, caplog
|
|
):
|
|
"""The bridged read hands the program the exception type, not the
|
|
message, and logs the traceback."""
|
|
sandbox._loop = asyncio.get_running_loop()
|
|
|
|
async def failing_read():
|
|
raise ValueError("failed at /secret/path")
|
|
|
|
with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"):
|
|
with pytest.raises(RuntimeError, match="document read failed: ValueError"):
|
|
await asyncio.to_thread(sandbox._run_on_loop, failing_read())
|
|
|
|
assert any(
|
|
r.exc_info and "failed at /secret/path" in str(r.exc_info[1])
|
|
for r in caplog.records
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_after_deadline_raises_without_scheduling(self, sandbox):
|
|
"""A read attempted past the deadline fails instead of querying."""
|
|
scheduled = False
|
|
|
|
async def _never_runs():
|
|
nonlocal scheduled
|
|
scheduled = True
|
|
|
|
sandbox._loop = asyncio.get_running_loop()
|
|
sandbox._deadline = sandbox._loop.time() - 1.0
|
|
|
|
coro = _never_runs()
|
|
with pytest.raises(TimeoutError, match="time limit exceeded"):
|
|
sandbox._run_on_loop(coro)
|
|
|
|
coro.close()
|
|
assert scheduled is False
|
|
|
|
def test_session_budget_covers_every_permitted_execution(self, temp_db_path):
|
|
"""Monty spends its duration budget across the session's whole life, so a
|
|
per-call value would let the first call starve the rest."""
|
|
config = AppConfig()
|
|
config.analysis.code_timeout = 5.0
|
|
config.analysis.max_executions = 3
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
|
|
assert sb._session_limits() == {"max_duration_secs": 15.0}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refused_read_fails_the_execution(self, temp_db_path, monkeypatch):
|
|
"""The refusal surfaces as a failed result, not a raised exception."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
docling = DoclingDocument(name="d")
|
|
docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.")
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.import_document(
|
|
docling,
|
|
[
|
|
Chunk(
|
|
content="Foxes and dogs.",
|
|
embedding=[0.1] * config.embeddings.model.vector_dim,
|
|
order=0,
|
|
)
|
|
],
|
|
uri="test://deadline",
|
|
)
|
|
|
|
def _past_deadline(*_args, **_kwargs):
|
|
raise TimeoutError(
|
|
"time limit exceeded: no further document reads after 60.0s"
|
|
)
|
|
|
|
monkeypatch.setattr(Sandbox, "_run_on_loop", _past_deadline)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
|
|
)
|
|
assert result.success is False
|
|
assert "no further document reads" in result.stderr
|
|
finally:
|
|
await sb.close()
|
|
|
|
|
|
class TestSandboxWorkerCrash:
|
|
"""A dead worker must not poison every later call in the run."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_crashed_worker_is_replaced(self, temp_db_path):
|
|
"""The crash fails one call. The next call gets a fresh session."""
|
|
import os
|
|
import signal
|
|
|
|
config = AppConfig()
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
first = await sb.execute("x = 1\nprint(x)")
|
|
assert first.success, first.stderr
|
|
assert sb._session is not None
|
|
pid = sb._session.worker_pid
|
|
assert pid is not None
|
|
|
|
os.kill(pid, signal.SIGKILL)
|
|
|
|
crashed = await sb.execute("print(2)")
|
|
assert crashed.success is False
|
|
assert "restarted" in crashed.stderr
|
|
|
|
# Without the discard every later call fails on the dead session.
|
|
recovered = await sb.execute("print(3)")
|
|
assert recovered.success, recovered.stderr
|
|
assert "3" in recovered.stdout
|
|
|
|
# The replacement worker starts empty, which the failure said.
|
|
lost = await sb.execute("print(x)")
|
|
assert lost.success is False
|
|
finally:
|
|
await sb.close()
|
|
|
|
|
|
class TestSandboxRequestTimeout:
|
|
"""The pool watchdog bounds a call that never reads."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runaway_compute_is_killed_and_the_next_call_recovers(
|
|
self, temp_db_path
|
|
):
|
|
"""Code that never reads escapes the read deadline. The watchdog kills it."""
|
|
config = AppConfig()
|
|
config.analysis.code_timeout = 1.0
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
runaway = await sb.execute(
|
|
"x = 0\nfor i in range(500000000):\n x += i\nprint(x)"
|
|
)
|
|
assert runaway.success is False
|
|
assert "restarted" in runaway.stderr
|
|
|
|
recovered = await sb.execute("print('alive')")
|
|
assert recovered.success, recovered.stderr
|
|
assert "alive" in recovered.stdout
|
|
finally:
|
|
await sb.close()
|
|
|
|
|
|
class TestSandboxClose:
|
|
@pytest.mark.asyncio
|
|
async def test_a_failing_teardown_still_releases_the_rest(self, tmp_path):
|
|
"""Each of the session, the pool and the held connection is released,
|
|
whichever of them fails."""
|
|
from unittest.mock import AsyncMock
|
|
|
|
sb = Sandbox(
|
|
db_path=tmp_path / "x.lancedb",
|
|
config=AppConfig(),
|
|
context=AnalysisContext(),
|
|
)
|
|
session, pool, opened = AsyncMock(), AsyncMock(), AsyncMock()
|
|
session.__aexit__.side_effect = RuntimeError("worker already gone")
|
|
sb._session, sb._pool = session, pool # ty: ignore[invalid-assignment]
|
|
sb._opened = opened
|
|
|
|
await sb.close()
|
|
|
|
pool.__aexit__.assert_awaited_once()
|
|
opened.__aexit__.assert_awaited_once()
|
|
assert sb._session is None
|
|
assert sb._pool is None
|
|
assert sb._opened is None
|