diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py index 040ad023..e2b66799 100644 --- a/tests/skills/test_analysis.py +++ b/tests/skills/test_analysis.py @@ -159,3 +159,36 @@ class TestExecuteCodeTool: ctx, code="docs = await list_documents()\nprint(len(docs))" ) assert "1" in result + + async def test_execute_code_accumulates_search_results(self, rag_db): + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(db_path=rag_db) + execute_code = _get_tool(skill, "execute_code") + state = AnalysisState() + ctx = _make_ctx(state) + await execute_code( + ctx, code="results = await search('intelligence')\nprint(len(results))" + ) + assert "_sandbox" in state.searches + assert len(state.searches["_sandbox"]) > 0 + + async def test_execute_code_vfs_write_denied(self, rag_db): + from haiku.rag.skills.analysis import create_skill + + skill = create_skill(db_path=rag_db) + execute_code = _get_tool(skill, "execute_code") + state = AnalysisState() + ctx = _make_ctx(state) + result = await execute_code( + ctx, + code=( + "from pathlib import Path\n" + "import json\n" + "dirs = list(Path('/documents').iterdir())\n" + "p = dirs[0] / 'content.txt'\n" + "p.write_text('hacked')" + ), + ) + assert "Error" in result + assert "read-only" in result diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index d39f080b..6fee0262 100644 --- a/tests/skills/test_rag.py +++ b/tests/skills/test_rag.py @@ -200,6 +200,23 @@ class TestSearchTool: result = await search(ctx, query="artificial intelligence") assert isinstance(result, str) + async def test_search_rate_limited(self, rag_db): + from haiku.rag.skills.rag import RAGState, create_skill + + config = AppConfig() + config.qa.max_searches = 2 + skill = create_skill(db_path=rag_db, config=config) + search = _get_tool(skill, "search") + state = RAGState() + ctx = _make_ctx(state) + ctx.run_id = "test-run" + + await search(ctx, query="first") + await search(ctx, query="second") + result = await search(ctx, query="third") + assert "Search limit reached" in result + assert len(state.searches) == 2 + class TestListDocumentsTool: async def test_list_documents_returns_results(self, rag_db):