95fdb46c3a
analysis.code_timeout was enforced only in _run_on_loop, the bridge for database-bound reads; metadata.json, the cached JSONL files and in-code search() and list_documents() never looked at the clock, and Monty's watchdog counts compute only. Every host call now checks the deadline before it starts. Host errors keep their message for every caller: the masking added for the MCP server goes, and the sandbox is one path for the capability and the server alike.
1227 lines
46 KiB
Python
1227 lines
46 KiB
Python
import asyncio
|
|
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_the_documented_modules_import(self, sandbox):
|
|
"""The modules the instructions and the MCP description promise."""
|
|
result = await sandbox.execute(
|
|
"import json, re, math, pathlib, datetime\n"
|
|
"import collections, itertools, functools, dataclasses\n"
|
|
"print(collections.Counter('aab').most_common(1),"
|
|
" list(itertools.islice(itertools.count(), 2)))"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert "[('a', 2)] [0, 1]" in result.stdout
|
|
|
|
@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_keeps_its_message_for_the_program(
|
|
self, sandbox, monkeypatch
|
|
):
|
|
"""A host-side failure inside search() reaches the program with its
|
|
message, which the agent reads to repair its code."""
|
|
|
|
async def boom(self, *args, **kwargs):
|
|
raise ValueError("failed at /secret/path")
|
|
|
|
monkeypatch.setattr(HaikuRAG, "search", boom)
|
|
|
|
result = await sandbox.execute("await search('hello')")
|
|
|
|
assert not result.success
|
|
assert "ValueError: failed at /secret/path" in result.stderr
|
|
|
|
|
|
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_the_deadline_covers_reads_from_memory_and_in_code_calls(
|
|
self, temp_db_path, monkeypatch
|
|
):
|
|
"""Once a call's time is up, a file served from memory and an in-code
|
|
listing are refused like a database read. A slow first read spends the
|
|
budget; the watchdog does not count time spent waiting on the host."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
config = AppConfig()
|
|
config.analysis.code_timeout = 1.0
|
|
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-paths",
|
|
)
|
|
repository = type(client.document_repository)
|
|
|
|
async def slow_content(self, *args, **kwargs):
|
|
await asyncio.sleep(1.3)
|
|
return "body"
|
|
|
|
monkeypatch.setattr(repository, "get_content", slow_content)
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
f"root = Path('/documents/{doc.id}')\n"
|
|
"print(len((root / 'content.txt').read_text()))\n"
|
|
"try:\n"
|
|
" (root / 'metadata.json').read_text()\n"
|
|
" print('static: read')\n"
|
|
"except Exception as e:\n"
|
|
" print('static:', type(e).__name__)\n"
|
|
"await list_documents()\n"
|
|
"print('listed')"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
|
|
assert "static: TimeoutError" in result.stdout
|
|
assert "listed" not in result.stdout
|
|
assert not result.success
|
|
assert "time limit exceeded" in result.stderr
|
|
|
|
@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())
|
|
|
|
limits = sb._session_limits()
|
|
|
|
assert limits["max_duration_secs"] == 15.0
|
|
cap = limits["max_suspensions"]
|
|
assert cap is not None
|
|
assert cap >= 1_000_000
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_program_may_read_more_than_a_thousand_times(self, temp_db_path):
|
|
"""Monty caps host callbacks per checkout at 1000 unless told otherwise;
|
|
a corpus-wide pass over documents reads far more than that."""
|
|
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://many-reads",
|
|
)
|
|
|
|
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
|
|
try:
|
|
result = await sb.execute(
|
|
"from pathlib import Path\n"
|
|
f"p = Path('/documents/{doc.id}/content.txt')\n"
|
|
"n = 0\n"
|
|
"for i in range(1100):\n"
|
|
" n += len(p.read_text())\n"
|
|
"print(n)"
|
|
)
|
|
finally:
|
|
await sb.close()
|
|
|
|
assert result.success, result.stderr
|
|
assert result.stdout.strip() == str(1100 * len("Foxes and dogs."))
|
|
|
|
@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
|