Security fixes for RLM. Fix type() builtin, AST validation for dict key access, sql injection

TOOLS_REFACTORING_PLAN.md

@@ -0,0 +1,200 @@
+# Tools Extraction Refactoring Plan
+
+## Goal
+
+Extract tools from haiku.rag agents into a reusable `tools/` module, enabling users to create pydantic-ai agents outside haiku.rag and compose toolsets as needed.
+
+## Target API
+
+```python
+from pydantic_ai import Agent
+from haiku.rag import HaikuRAG
+from haiku.rag.tools import ToolContext, create_search_toolset, create_document_toolset
+
+async with HaikuRAG(db_path) as client:
+    context = ToolContext()
+    search_tools = create_search_toolset(client, config, context)
+    doc_tools = create_document_toolset(client, config, context)
+
+    agent = Agent(
+        'anthropic:claude-sonnet',
+        toolsets=[search_tools, doc_tools]
+    )
+    result = await agent.run("Find documents about X")
+
+    # Access accumulated state after run
+    search_state = context.get("haiku.rag.search")
+    for result in search_state.results:
+        print(f"{result.document_title}")
+```
+
+## Design Principles
+
+1. **ToolContext is a pure generic container** - No special-cased fields. Toolsets register their own Pydantic model state under namespaces.
+
+2. **Shared state via same namespace** - Multiple toolsets can share state (e.g., citations, filters) by registering under the same namespace.
+
+3. **App manages identity** - ToolContext has no session/user identity. The app layer manages `session_id -> ToolContext` mapping.
+
+4. **Toolsets are stateless factories** - `create_*_toolset()` returns a `FunctionToolset`. State lives in the context they're given.
+
+## ToolContext Design
+
+```python
+class ToolContext(BaseModel):
+    """Generic state container for toolsets.
+
+    Toolsets register Pydantic model state under namespaces.
+    Multiple toolsets can share state via the same namespace.
+    """
+    _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
+
+    def register(self, namespace: str, state: BaseModel) -> None: ...
+    def get(self, namespace: str) -> BaseModel | None: ...
+    def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T: ...
+    def clear_namespace(self, namespace: str) -> None: ...
+    def clear_all(self) -> None: ...
+    def dump_namespaces(self) -> dict[str, dict[str, Any]]: ...
+    def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T: ...
+```
+
+## Toolset State Examples
+
+Each toolset defines its own state model:
+
+```python
+# Search toolset state
+class SearchState(BaseModel):
+    results: list[SearchResult] = []
+    filter: str | None = None
+
+SEARCH_NAMESPACE = "haiku.rag.search"
+
+# QA toolset state
+class QAState(BaseModel):
+    history: list[QAResult] = []
+
+QA_NAMESPACE = "haiku.rag.qa"
+
+# Shared citation state (used by multiple toolsets)
+class CitationState(BaseModel):
+    registry: dict[str, int] = {}
+
+    def get_or_assign_index(self, chunk_id: str) -> int:
+        if chunk_id in self.registry:
+            return self.registry[chunk_id]
+        new_index = len(self.registry) + 1
+        self.registry[chunk_id] = new_index
+        return new_index
+
+CITATION_NAMESPACE = "haiku.rag.citations"
+```
+
+## Multi-User/Session Management
+
+App layer manages context routing:
+
+```python
+# App maintains context per session
+contexts: dict[str, ToolContext] = {}
+
+def get_context(session_id: str) -> ToolContext:
+    if session_id not in contexts:
+        contexts[session_id] = ToolContext()
+    return contexts[session_id]
+
+# When running agent
+context = get_context(user_session_id)
+toolsets = [create_search_toolset(client, config, context)]
+await agent.run(prompt, toolsets=toolsets)
+```
+
+## New Module Structure
+
+```
+haiku_rag_slim/haiku/rag/
+├── tools/                      # NEW
+│   ├── __init__.py             # Public exports
+│   ├── context.py              # ToolContext (generic state container)
+│   ├── models.py               # QAResult, AnalysisResult
+│   ├── filters.py              # build_document_filter, combine_filters
+│   ├── search.py               # create_search_toolset()
+│   ├── document.py             # create_document_toolset()
+│   ├── qa.py                   # create_qa_toolset()
+│   └── analysis.py             # create_analysis_toolset()
+├── agents/                     # REFACTORED to use tools/
+```
+
+## Implementation Chunks
+
+### Chunk 1: Create tools module foundation ✅ DONE
+- Created `tools/__init__.py`, `tools/context.py`, `tools/models.py`, `tools/filters.py`
+- Created `ToolContext` as generic namespace-based Pydantic model
+- Moved filter utilities from `agents/chat/state.py` to `tools/filters.py`
+- Created result models (`QAResult`, `AnalysisResult`)
+- Added tests for ToolContext and filters
+
+### Chunk 2: Create SearchToolset ✅ DONE
+- Created `tools/search.py` with `create_search_toolset()`
+- Defined `SearchState` model for accumulating search results
+- Core search logic: `client.search()` → `client.expand_context()` → `format_for_agent()`
+- Results accumulated in `SearchState` under `SEARCH_NAMESPACE`
+- Added 13 tests for SearchToolset
+
+### Chunk 3: Refactor QA Agent to use SearchToolset ✅ DONE
+- Updated `agents/qa/agent.py` to use `create_search_toolset()`
+- Added `base_filter` and `tool_name` parameters to `create_search_toolset()`
+- QA agent now uses ToolContext + SearchState for result accumulation
+- Public interface (`answer(question, filter)`) unchanged
+- All 5 QA tests pass
+
+### Chunk 4: Create DocumentToolset ✅ DONE
+- Created `tools/document.py` with `create_document_toolset()`
+- Defined `DocumentState`, `DocumentInfo`, `DocumentListResponse` models
+- Extracted `list_documents`, `get_document`, `summarize_document` tools
+- Moved `find_document` helper (now public)
+- Added 13 tests
+
+### Chunk 5: Create QAToolset ✅ DONE
+- Created `tools/qa.py` with `create_qa_toolset()`
+- Defined `QAState` model (tracks QA history)
+- Runs research graph, returns structured `QAResult`
+- Supports `base_filter`, `tool_name`, `session_context`, `prior_answers` params
+- Added 7 tests
+
+### Chunk 6: Create AnalysisToolset ✅ DONE
+- Created `tools/analysis.py` with `create_analysis_toolset()`
+- Defined `AnalysisState` model (tracks CodeExecution history)
+- Extracted `analyze` tool (RLM delegation with filter support)
+- Fixed circular import by using direct submodule imports
+- Added 6 tests
+
+### Chunk 7: Refactor Chat Agent ✅ DONE
+- Removed `analyze` tool from chat agent (kept hardcoded, not composing toolsets)
+- Reverted system prompt to pre-analyze version
+- Removed `test_analyze_tool` test and cassette file
+- All 47 chat agent tests pass
+
+### Chunk 8: Refactor Research Graph
+- Update `_search_one_step_logic` to use search toolset
+- Verify research tests pass
+
+### Chunk 9: Public API and Documentation
+- Export from `haiku.rag.tools` and `haiku.rag`
+- Update CLAUDE.md
+- Add usage examples
+
+## Verification
+
+- Run `pytest` after each chunk
+- Run `ty check` and `ruff check`
+- Test with existing agents (QA, Chat, Research)
+- Test with external agent using new toolsets
+
+## Critical Files
+
+- `haiku_rag_slim/haiku/rag/agents/chat/agent.py` - largest tool collection
+- `haiku_rag_slim/haiku/rag/agents/qa/agent.py` - simplest, good starting point
+- `haiku_rag_slim/haiku/rag/agents/chat/state.py` - filter utilities (now moved)
+- `haiku_rag_slim/haiku/rag/agents/research/graph.py` - search tool inside step
+- `haiku_rag_slim/haiku/rag/store/models/chunk.py` - SearchResult.format_for_agent()

haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
 
 from haiku.rag.agents.rlm.dependencies import RLMContext
 from haiku.rag.config.models import RLMConfig
+from haiku.rag.store.repositories.document import _escape_sql_string
 
 if TYPE_CHECKING:
     from haiku.rag.client import HaikuRAG
@@ -84,7 +85,9 @@ class REPLEnvironment:
         "str": str,
         "sum": sum,
         "tuple": tuple,
-        "type": type,
+        "type": (
+            lambda obj: type(obj)
+        ),  # Single-arg only, blocks type(name, bases, dict)
         "zip": zip,
         "Exception": Exception,
         "ValueError": ValueError,
@@ -206,13 +209,14 @@ class REPLEnvironment:
                 doc = await self.client.get_document_by_id(id_or_title)
                 if doc:
                     return doc.content
+                safe_input = _escape_sql_string(id_or_title)
                 docs = await self.client.list_documents(
-                    filter=f"title = '{id_or_title}'"
+                    filter=f"title = '{safe_input}'"
                 )
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
                     return full_doc.content if full_doc else None
-                docs = await self.client.list_documents(filter=f"uri = '{id_or_title}'")
+                docs = await self.client.list_documents(filter=f"uri = '{safe_input}'")
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
                     return full_doc.content if full_doc else None
@@ -230,13 +234,14 @@ class REPLEnvironment:
                 doc = await self.client.get_document_by_id(id_or_title)
                 if doc:
                     return doc.get_docling_document()
+                safe_input = _escape_sql_string(id_or_title)
                 docs = await self.client.list_documents(
-                    filter=f"title = '{id_or_title}'"
+                    filter=f"title = '{safe_input}'"
                 )
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
                     return full_doc.get_docling_document() if full_doc else None
-                docs = await self.client.list_documents(filter=f"uri = '{id_or_title}'")
+                docs = await self.client.list_documents(filter=f"uri = '{safe_input}'")
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
                     return full_doc.get_docling_document() if full_doc else None
@@ -305,6 +310,16 @@ class REPLEnvironment:
                     raise SecurityError(
                         f"Access to private/dunder attribute '{node.attr}' is not allowed"
                     )
+            # Block dictionary key access to dunder/private strings
+            # This prevents type.__dict__['__subclasses__'] attacks
+            if isinstance(node, ast.Subscript):
+                if isinstance(node.slice, ast.Constant):
+                    if isinstance(
+                        node.slice.value, str
+                    ) and node.slice.value.startswith("_"):
+                        raise SecurityError(
+                            f"Dictionary access to '{node.slice.value}' is not allowed"
+                        )
 
     def _execute_sync(self, code: str) -> REPLResult:
         """Internal synchronous execution - must be called from executor thread."""

haiku_rag_slim/haiku/rag/client.py

@@ -22,7 +22,10 @@ from haiku.rag.store.engine import Store
 from haiku.rag.store.models.chunk import Chunk, SearchResult
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
-from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.store.repositories.document import (
+    DocumentRepository,
+    _escape_sql_string,
+)
 from haiku.rag.store.repositories.settings import SettingsRepository
 
 if TYPE_CHECKING:
@@ -1322,7 +1325,8 @@ class HaikuRAG:
             for doc_ref in documents:
                 doc = await self.get_document_by_id(doc_ref)
                 if not doc:
-                    docs = await self.list_documents(filter=f"title = '{doc_ref}'")
+                    safe_ref = _escape_sql_string(doc_ref)
+                    docs = await self.list_documents(filter=f"title = '{safe_ref}'")
                     if docs and docs[0].id:
                         doc = await self.get_document_by_id(docs[0].id)
                 if doc:

haiku_rag_slim/haiku/rag/store/repositories/document.py

@@ -77,9 +77,10 @@ class DocumentRepository:
 
     async def get_by_id(self, entity_id: str) -> Document | None:
         """Get a document by its ID."""
+        safe_id = _escape_sql_string(entity_id)
         results = list(
             self.store.documents_table.search()
-            .where(f"id = '{entity_id}'")
+            .where(f"id = '{safe_id}'")
             .limit(1)
             .to_pydantic(DocumentRecord)
         )
@@ -104,8 +105,9 @@ class DocumentRepository:
         entity.updated_at = datetime.fromisoformat(now)
 
         # Update the record
+        safe_id = _escape_sql_string(entity.id)
         self.store.documents_table.update(
-            where=f"id = '{entity.id}'",
+            where=f"id = '{safe_id}'",
             values={
                 "content": entity.content,
                 "uri": entity.uri,
@@ -136,7 +138,8 @@ class DocumentRepository:
         await self.chunk_repository.delete_by_document_id(entity_id)
 
         # Delete the document
-        self.store.documents_table.delete(f"id = '{entity_id}'")
+        safe_id = _escape_sql_string(entity_id)
+        self.store.documents_table.delete(f"id = '{safe_id}'")
         return True
 
     async def list_all(

tests/agents/rlm/test_sandbox.py

@@ -1,6 +1,13 @@
+from pathlib import Path
+
 import pytest
 
 
+@pytest.fixture(scope="module")
+def vcr_cassette_dir():
+    return str(Path(__file__).parent.parent.parent / "cassettes" / "test_sandbox")
+
+
 class TestSafeBuiltins:
     """Test that safe builtins are available."""
 
@@ -572,6 +579,130 @@ class TestPreloadedDocuments:
             assert "test://doc" in result.stdout
 
 
+class TestSandboxEscapeVectors:
+    """Test that known sandbox escape techniques are blocked.
+
+    Each test contains actual exploit code that would work without the fix.
+    """
+
+    @pytest.mark.asyncio
+    async def test_type_dict_subclasses_escape_blocked(self, repl_env_empty):
+        """Cannot escape via type.__dict__['__subclasses__'].
+
+        Without fix: This would enumerate all loaded classes and find
+        subprocess.Popen to execute arbitrary shell commands.
+        """
+        result = await repl_env_empty.execute_async("""
+# EXPLOIT: Access __subclasses__ via dict to bypass AST check
+subclasses_method = type.__dict__['__subclasses__']
+all_classes = subclasses_method(object)
+print(f"Found {len(all_classes)} classes")
+""")
+        assert not result.success
+        assert "not allowed" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    async def test_popen_shell_execution_blocked(self, repl_env_empty):
+        """Cannot execute shell commands via Popen.
+
+        Without fix: This would execute 'whoami' and return the username.
+        """
+        result = await repl_env_empty.execute_async("""
+# EXPLOIT: Find subprocess.Popen and execute shell commands
+subclasses_method = type.__dict__['__subclasses__']
+all_classes = subclasses_method(object)
+popen = [c for c in all_classes if c.__name__ == 'Popen'][0]
+proc = popen('whoami', shell=True, stdout=-1)
+print(proc.stdout.read())
+""")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_socket_creation_blocked(self, repl_env_empty):
+        """Cannot create network sockets for data exfiltration.
+
+        Without fix: This would create a socket that could connect to external servers.
+        """
+        result = await repl_env_empty.execute_async("""
+# EXPLOIT: Find socket class and create network connection
+subclasses_method = type.__dict__['__subclasses__']
+all_classes = subclasses_method(object)
+socket_cls = [c for c in all_classes if c.__name__ == 'socket'][0]
+s = socket_cls(2, 1)  # AF_INET, SOCK_STREAM
+print(f"Created socket: {s}")
+""")
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_type_three_arg_class_creation_blocked(self, repl_env_empty):
+        """Cannot use type() with 3 arguments to create classes dynamically."""
+        result = await repl_env_empty.execute_async(
+            "EvilClass = type('EvilClass', (object,), {'x': 1})"
+        )
+        assert not result.success
+
+    @pytest.mark.asyncio
+    async def test_dict_key_dunder_access_blocked(self, repl_env_empty):
+        """Cannot access dunder methods via dictionary key access."""
+        result = await repl_env_empty.execute_async(
+            "method = str.__dict__['__add__']\nprint(method)"
+        )
+        assert not result.success
+        assert "not allowed" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    async def test_dict_key_private_access_blocked(self, repl_env_empty):
+        """Cannot access private attributes via dictionary key access."""
+        result = await repl_env_empty.execute_async(
+            "method = object.__dict__['_private']\nprint(method)"
+        )
+        assert not result.success
+        assert "not allowed" in result.stderr.lower()
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_sql_injection_in_get_document_blocked(self, temp_db_path):
+        """SQL injection in get_document cannot bypass context filter.
+
+        Without fix: Injecting quotes would leak documents that should be
+        protected by the context filter.
+        """
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            # Create documents: one secret, one public
+            await client.create_document(
+                content="TOP SECRET: Launch codes 1234",
+                uri="secret://classified",
+                title="Classified Intel",
+            )
+            await client.create_document(
+                content="Public weather report",
+                uri="public://weather",
+                title="Weather",
+            )
+
+            # Sandbox restricted to public:// only
+            context = RLMContext(filter="uri LIKE 'public://%'")
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+
+            # EXPLOIT: SQL injection to access secret document
+            result = await repl.execute_async("""
+# Injection payload breaks out of quotes and adds OR clause
+content = get_document("x' OR uri LIKE 'secret://%")
+if content:
+    print(f"LEAKED: {content}")
+else:
+    print("NO LEAK")
+""")
+            assert result.success
+            assert "TOP SECRET" not in result.stdout
+            assert "Launch codes" not in result.stdout
+
+
 class TestSecurityEscapes:
     """Test that common security escape attempts are blocked."""
 

tests/test_client.py

@@ -1384,3 +1384,50 @@ async def test_client_convert_with_html_format(temp_db_path):
         labels = [str(getattr(item, "label", "")) for item, _ in items]
 
         assert "title" in labels
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_sql_injection_is_blocked_with_escaping(temp_db_path):
+    """SQL injection is blocked when using _escape_sql_string.
+
+    This test verifies that _escape_sql_string properly prevents SQL injection
+    by escaping single quotes in user input.
+    """
+    from haiku.rag.store.repositories.document import _escape_sql_string
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Create documents
+        await client.create_document(
+            content="Secret classified data XYZ",
+            uri="secret://doc",
+            title="Secret",
+        )
+        await client.create_document(
+            content="Public report about weather",
+            uri="public://report",
+            title="Weather Report",
+        )
+
+        # Without escaping, this injection would match all documents
+        # by breaking out of the string literal: title = 'x' OR title LIKE '%'
+        injection_payload = "x' OR title LIKE '%"
+
+        # With proper escaping, single quotes become double quotes
+        # so the filter becomes: title = 'x'' OR title LIKE ''%'
+        # which searches for a literal title containing the injection string
+        safe_payload = _escape_sql_string(injection_payload)
+        docs = await client.list_documents(filter=f"title = '{safe_payload}'")
+
+        # Should find 0 documents (injection is escaped, searching for literal string)
+        assert len(docs) == 0
+
+        # Verify the escaping works correctly
+        assert safe_payload == "x'' OR title LIKE ''%"
+
+        # Verify unescaped injection would have matched documents (for test validity)
+        # This demonstrates that the injection works without escaping
+        docs_unescaped = await client.list_documents(
+            filter=f"title = '{injection_payload}'"
+        )
+        assert len(docs_unescaped) == 2  # SQL injection succeeds without escaping