open one HaikuRAG client per skill invocation via lifespan

CHANGELOG.md

@@ -1,6 +1,11 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- **Skills share a single `HaikuRAG` client per invocation** via the new `haiku.skills>=0.15.0` `lifespan` hook. The skill's sub-agent opens one read-only client on entry, all tool calls reuse it, and it closes on exit — replacing the old pattern of open/close around every `search` / `list_documents` / `get_document` call.
+- **`max_searches` tracked on `RAGRunDeps.search_count`** instead of a module-level `ctx.run_id`-keyed dict. Eliminates a memory leak in long-running processes where old run ids were never evicted.
+
 ## [0.41.0] - 2026-04-20
 
 ### Added

haiku_rag_slim/haiku/rag/skills/_deps.py

@@ -0,0 +1,36 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from haiku.rag.config.models import AppConfig
+from haiku.skills.state import SkillRunDeps
+
+if TYPE_CHECKING:
+    from haiku.rag.agents.analysis.sandbox import Sandbox
+    from haiku.rag.client import HaikuRAG
+
+
+@dataclass
+class RAGRunDeps(SkillRunDeps):
+    rag: "HaikuRAG | None" = None
+    search_count: int = 0
+
+
+@dataclass
+class AnalysisRunDeps(RAGRunDeps):
+    sandbox: "Sandbox | None" = None
+
+
+def make_rag_lifespan(db_path: Path, config: AppConfig):
+    @asynccontextmanager
+    async def lifespan(deps: RAGRunDeps) -> AsyncIterator[None]:
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+            deps.rag = rag
+            deps.search_count = 0
+            yield
+
+    return lifespan

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -5,9 +5,10 @@ from pydantic import BaseModel
 from pydantic_ai import RunContext
 
 from haiku.rag.agents.research.models import Citation
+from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
+from haiku.rag.skills._deps import RAGRunDeps
 from haiku.rag.store.models.chunk import SearchResult
-from haiku.skills.state import SkillRunDeps
 
 
 class CodeExecutionEntry(BaseModel):
@@ -18,22 +19,13 @@ class CodeExecutionEntry(BaseModel):
 
 
 async def skill_search(
-    db_path: Path,
+    rag: HaikuRAG,
-    config: AppConfig,
     query: str,
     limit: int | None = None,
     document_filter: str | None = None,
 ) -> tuple[str, list[SearchResult]]:
-    from haiku.rag.client import HaikuRAG
+    results = await rag.search(query, limit=limit, filter=document_filter)
-
+    results = await rag.expand_context(results)
-    async with HaikuRAG(db_path, config=config, read_only=True) as rag:
-        results = await rag.search(
-            query,
-            limit=limit,
-            filter=document_filter,
-        )
-        results = await rag.expand_context(results)
-
     formatted = "\n\n---\n\n".join(
         r.format_for_agent(rank=i + 1, total=len(results))
         for i, r in enumerate(results)
@@ -42,55 +34,55 @@ async def skill_search(
 
 
 async def skill_list_documents(
-    db_path: Path,
+    rag: HaikuRAG,
-    config: AppConfig,
     filter: str | None = None,
 ) -> list[dict[str, Any]]:
-    from haiku.rag.client import HaikuRAG
+    documents = await rag.list_documents(filter=filter)
-
+    return [
-    async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+        {
-        documents = await rag.list_documents(filter=filter)
+            "id": doc.id,
-        return [
+            "title": doc.title,
-            {
+            "uri": doc.uri,
-                "id": doc.id,
+            "metadata": doc.metadata,
-                "title": doc.title,
+            "created_at": str(doc.created_at),
-                "uri": doc.uri,
+            "updated_at": str(doc.updated_at),
-                "metadata": doc.metadata,
+        }
-                "created_at": str(doc.created_at),
+        for doc in documents
-                "updated_at": str(doc.updated_at),
+    ]
-            }
-            for doc in documents
-        ]
 
 
 async def skill_get_document(
-    db_path: Path,
+    rag: HaikuRAG,
-    config: AppConfig,
     query: str,
 ) -> dict[str, Any] | None:
-    from haiku.rag.client import HaikuRAG
+    document = await rag.resolve_document(query)
-
+    if document is None:
-    async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+        return None
-        document = await rag.resolve_document(query)
+    return {
-        if document is None:
+        "id": document.id,
-            return None
+        "content": document.content,
-        return {
+        "title": document.title,
-            "id": document.id,
+        "uri": document.uri,
-            "content": document.content,
+        "metadata": document.metadata,
-            "title": document.title,
+        "created_at": str(document.created_at),
-            "uri": document.uri,
+        "updated_at": str(document.updated_at),
-            "metadata": document.metadata,
+    }
-            "created_at": str(document.created_at),
-            "updated_at": str(document.updated_at),
-        }
 
 
-def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> Any:
+def _get_state(ctx: RunContext[RAGRunDeps], state_type: type[BaseModel]) -> Any:
     if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type):
         return ctx.deps.state
     return None
 
 
+def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
+    if ctx.deps is None or ctx.deps.rag is None:
+        raise RuntimeError(
+            "RAGRunDeps.rag is not set — skill lifespan must run before tools."
+        )
+    return ctx.deps.rag
+
+
 def _register_citations(state: Any, citations: "list[Citation]") -> None:
     """Add citations to the index and record the turn's chunk IDs."""
     chunk_ids = []
@@ -174,10 +166,9 @@ def create_skill_tools(
 
     if "search" in tool_names:
         max_searches = config.qa.max_searches
-        search_counts: dict[str, int] = {}
 
         async def search(
-            ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
+            ctx: RunContext[RAGRunDeps], query: str, limit: int | None = None
         ) -> str:
             """Search the knowledge base using hybrid search (vector + full-text).
 
@@ -187,9 +178,8 @@ def create_skill_tools(
                 query: The search query.
                 limit: Maximum number of results.
             """
-            rid = ctx.run_id or ""
+            ctx.deps.search_count += 1
-            search_counts[rid] = search_counts.get(rid, 0) + 1
+            if ctx.deps.search_count > max_searches:
-            if search_counts[rid] > max_searches:
                 return (
                     "Search limit reached. Answer the question using "
                     "the results you already have."
@@ -197,8 +187,7 @@ def create_skill_tools(
 
             state = _get_state(ctx, state_type)
             formatted, results = await skill_search(
-                db_path,
+                _require_rag(ctx),
-                config,
                 query,
                 limit=limit,
                 document_filter=state.document_filter if state else None,
@@ -212,36 +201,34 @@ def create_skill_tools(
     if "list_documents" in tool_names:
 
         async def list_documents(
-            ctx: RunContext[SkillRunDeps],
+            ctx: RunContext[RAGRunDeps],
         ) -> list[dict[str, Any]]:
             """List all documents in the knowledge base."""
             state = _get_state(ctx, state_type)
-            result = await skill_list_documents(
+            return await skill_list_documents(
-                db_path,
+                _require_rag(ctx),
-                config,
                 filter=state.document_filter if state else None,
             )
-            return result
 
         tools["list_documents"] = list_documents
 
     if "get_document" in tool_names:
 
         async def get_document(
-            ctx: RunContext[SkillRunDeps], query: str
+            ctx: RunContext[RAGRunDeps], query: str
         ) -> dict[str, Any] | None:
             """Retrieve a document by ID, title, or URI.
 
             Args:
                 query: Document ID, title, or URI to look up.
             """
-            return await skill_get_document(db_path, config, query)
+            return await skill_get_document(_require_rag(ctx), query)
 
         tools["get_document"] = get_document
 
     if "execute_code" in tool_names:
 
-        async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
+        async def execute_code(ctx: RunContext[RAGRunDeps], code: str) -> str:
             """Execute Python code in a sandboxed interpreter.
 
             The code has access to search(), list_documents(), llm() functions
@@ -291,7 +278,7 @@ def create_skill_tools(
 
     if "cite" in tool_names:
 
-        async def cite(ctx: RunContext[SkillRunDeps], chunk_ids: list[str]) -> str:
+        async def cite(ctx: RunContext[RAGRunDeps], chunk_ids: list[str]) -> str:
             """Register chunk IDs as citations for your answer.
 
             Call this after searching, with the chunk_id values from search

haiku_rag_slim/haiku/rag/skills/analysis.py

@@ -60,6 +60,7 @@ def create_skill(
         config: haiku.rag AppConfig instance. If None, uses get_config().
     """
     from haiku.rag.config import get_config
+    from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan
     from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
 
     if config is None:
@@ -93,4 +94,6 @@ def create_skill(
         extras=extras,
         state_type=STATE_TYPE,
         state_namespace=STATE_NAMESPACE,
+        deps_type=AnalysisRunDeps,
+        lifespan=make_rag_lifespan(db_path, config),
     )

haiku_rag_slim/haiku/rag/skills/rag.py

@@ -75,6 +75,7 @@ def create_skill(
         config: haiku.rag AppConfig instance. If None, uses get_config().
     """
     from haiku.rag.config import get_config
+    from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
     from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
 
     if config is None:
@@ -103,4 +104,6 @@ def create_skill(
         extras=extras,
         state_type=STATE_TYPE,
         state_namespace=STATE_NAMESPACE,
+        deps_type=RAGRunDeps,
+        lifespan=make_rag_lifespan(db_path, config),
     )

haiku_rag_slim/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 
 dependencies = [
     "docling-core>=2.71.0,<2.72",
-    "haiku.skills>=0.14.0",
+    "haiku.skills>=0.15.0",
     "httpx>=0.28.1",
     "jinja2>=3.1.0",
     "jsonpatch>=1.33",

tests/skills/conftest.py

@@ -7,15 +7,20 @@ from pydantic_ai import RunContext
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.embeddings import EmbedderWrapper
-from haiku.skills.state import SkillRunDeps
+from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
 
 VECTOR_DIM = 2560
 
 
-def _make_ctx(state=None):
+def _make_ctx(state=None, rag=None):
-    """Create a mock RunContext with SkillRunDeps."""
+    """Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
+    from haiku.rag.skills.analysis import AnalysisState
+
     ctx = MagicMock(spec=RunContext)
-    ctx.deps = SkillRunDeps(state=state)
+    if isinstance(state, AnalysisState):
+        ctx.deps = AnalysisRunDeps(state=state, rag=rag)
+    else:
+        ctx.deps = RAGRunDeps(state=state, rag=rag)
     return ctx
 
 
@@ -68,3 +73,10 @@ async def rag_db(temp_db_path):
             uri="test://ml-basics",
         )
     return temp_db_path
+
+
+@pytest.fixture
+async def rag_client(rag_db):
+    """Yield an open read-only HaikuRAG client on the sample db."""
+    async with HaikuRAG(rag_db, read_only=True) as rag:
+        yield rag

tests/skills/test_analysis.py

@@ -192,3 +192,28 @@ class TestExecuteCodeTool:
         )
         assert "Error" in result
         assert "read-only" in result
+
+
+class TestAnalysisLifespan:
+    async def test_opens_one_client_per_invocation(self, rag_db):
+        from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan
+
+        config = AppConfig()
+        lifespan = make_rag_lifespan(rag_db, config)
+        deps = AnalysisRunDeps()
+        async with lifespan(deps):
+            assert deps.rag is not None
+            assert deps.rag.is_read_only
+            assert deps.search_count == 0
+            docs = await deps.rag.list_documents()
+            assert len(docs) == 2
+
+    async def test_skill_has_lifespan_and_deps_type(
+        self, test_app_config, temp_db_path
+    ):
+        from haiku.rag.skills._deps import AnalysisRunDeps
+        from haiku.rag.skills.analysis import create_skill
+
+        skill = create_skill(config=test_app_config, db_path=temp_db_path)
+        assert skill.deps_type is AnalysisRunDeps
+        assert skill.lifespan is not None

tests/skills/test_rag.py

@@ -157,50 +157,50 @@ class TestSkillExtras:
 
 
 class TestSearchTool:
-    async def test_search_returns_formatted_string(self, rag_db):
+    async def test_search_returns_formatted_string(self, rag_db, rag_client):
         from haiku.rag.skills.rag import create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
-        ctx = _make_ctx()
+        ctx = _make_ctx(rag=rag_client)
         result = await search(ctx, query="artificial intelligence")
         assert isinstance(result, str)
         assert len(result) > 0
 
-    async def test_search_updates_state(self, rag_db):
+    async def test_search_updates_state(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
         state = RAGState()
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
         await search(ctx, query="artificial intelligence")
         assert "artificial intelligence" in state.searches
         results = state.searches["artificial intelligence"]
         assert len(results) > 0
         assert isinstance(results[0], SearchResult)
 
-    async def test_search_applies_document_filter_from_state(self, rag_db):
+    async def test_search_applies_document_filter_from_state(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
         state = RAGState(document_filter="title = 'AI Overview'")
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
         result = await search(ctx, query="artificial intelligence")
         assert "AI Overview" in result
         assert "ML Basics" not in result
 
-    async def test_search_without_state(self, rag_db):
+    async def test_search_without_state(self, rag_db, rag_client):
         from haiku.rag.skills.rag import create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
-        ctx = _make_ctx(state=None)
+        ctx = _make_ctx(state=None, rag=rag_client)
         result = await search(ctx, query="artificial intelligence")
         assert isinstance(result, str)
 
-    async def test_search_rate_limited(self, rag_db):
+    async def test_search_rate_limited(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         config = AppConfig()
@@ -208,69 +208,71 @@ class TestSearchTool:
         skill = create_skill(db_path=rag_db, config=config)
         search = _get_tool(skill, "search")
         state = RAGState()
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
-        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 ctx.deps.search_count == 3
         assert len(state.searches) == 2
 
 
 class TestListDocumentsTool:
-    async def test_list_documents_returns_results(self, rag_db):
+    async def test_list_documents_returns_results(self, rag_db, rag_client):
         from haiku.rag.skills.rag import create_skill
 
         skill = create_skill(db_path=rag_db)
         list_docs = _get_tool(skill, "list_documents")
-        ctx = _make_ctx()
+        ctx = _make_ctx(rag=rag_client)
         results = await list_docs(ctx)
         assert isinstance(results, list)
         assert len(results) == 2
 
-    async def test_list_documents_applies_document_filter_from_state(self, rag_db):
+    async def test_list_documents_applies_document_filter_from_state(
+        self, rag_db, rag_client
+    ):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         skill = create_skill(db_path=rag_db)
         list_docs = _get_tool(skill, "list_documents")
         state = RAGState(document_filter="title = 'AI Overview'")
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
         results = await list_docs(ctx)
         assert len(results) == 1
         assert results[0]["title"] == "AI Overview"
 
 
 class TestGetDocumentTool:
-    async def test_get_document_by_title(self, rag_db):
+    async def test_get_document_by_title(self, rag_db, rag_client):
         from haiku.rag.skills.rag import create_skill
 
         skill = create_skill(db_path=rag_db)
         get_doc = _get_tool(skill, "get_document")
-        ctx = _make_ctx()
+        ctx = _make_ctx(rag=rag_client)
         result = await get_doc(ctx, query="AI Overview")
         assert result is not None
         assert result["title"] == "AI Overview"
 
-    async def test_get_document_not_found(self, rag_db):
+    async def test_get_document_not_found(self, rag_db, rag_client):
         from haiku.rag.skills.rag import create_skill
 
         skill = create_skill(db_path=rag_db)
         get_doc = _get_tool(skill, "get_document")
-        ctx = _make_ctx()
+        ctx = _make_ctx(rag=rag_client)
         result = await get_doc(ctx, query="nonexistent document xyz")
         assert result is None
 
 
 class TestCiteTool:
-    async def test_cite_registers_citations(self, rag_db):
+    async def test_cite_registers_citations(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
         cite = _get_tool(skill, "cite")
         state = RAGState()
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
 
         await search(ctx, query="artificial intelligence")
         chunk_ids = [
@@ -286,14 +288,14 @@ class TestCiteTool:
         assert len(state.citations[0]) == 2
         assert all(cid in state.citation_index for cid in chunk_ids)
 
-    async def test_cite_deduplicates_in_index(self, rag_db):
+    async def test_cite_deduplicates_in_index(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill
 
         skill = create_skill(db_path=rag_db)
         search = _get_tool(skill, "search")
         cite = _get_tool(skill, "cite")
         state = RAGState()
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
 
         await search(ctx, query="artificial intelligence")
         chunk_ids = [
@@ -316,3 +318,43 @@ class TestCiteTool:
         ctx = _make_ctx(state=None)
         result = await cite(ctx, chunk_ids=["nonexistent"])
         assert "No state" in result
+
+
+class TestLifespan:
+    async def test_opens_one_client_per_invocation(self, rag_db):
+        """Lifespan opens one HaikuRAG client, available on ctx.deps.rag throughout."""
+        from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
+
+        config = AppConfig()
+        lifespan = make_rag_lifespan(rag_db, config)
+        deps = RAGRunDeps()
+        async with lifespan(deps):
+            assert deps.rag is not None
+            assert deps.rag.is_read_only
+            assert deps.search_count == 0
+            docs = await deps.rag.list_documents()
+            assert len(docs) == 2
+        # after exit the client has been closed; field still references it
+        assert deps.rag is not None
+
+    async def test_search_count_resets_per_invocation(self, rag_db):
+        from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
+
+        config = AppConfig()
+        lifespan = make_rag_lifespan(rag_db, config)
+
+        deps = RAGRunDeps(search_count=42)
+        async with lifespan(deps):
+            assert deps.search_count == 0
+
+        deps2 = RAGRunDeps(search_count=5)
+        async with lifespan(deps2):
+            assert deps2.search_count == 0
+
+    def test_skill_has_lifespan_and_deps_type(self, test_app_config, temp_db_path):
+        from haiku.rag.skills._deps import RAGRunDeps
+        from haiku.rag.skills.rag import create_skill
+
+        skill = create_skill(config=test_app_config, db_path=temp_db_path)
+        assert skill.deps_type is RAGRunDeps
+        assert skill.lifespan is not None

uv.lock

@@ -1570,7 +1570,7 @@ requires-dist = [
     { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.21.1" },
     { name = "docling", marker = "extra == 'docling'", specifier = ">=2.84.0" },
     { name = "docling-core", specifier = ">=2.71.0,<2.72" },
-    { name = "haiku-skills", specifier = ">=0.14.0" },
+    { name = "haiku-skills", specifier = ">=0.15.0" },
     { name = "httpx", specifier = ">=0.28.1" },
     { name = "jinja2", specifier = ">=3.1.0" },
     { name = "jsonpatch", specifier = ">=1.33" },
@@ -1604,7 +1604,7 @@ provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jin
 
 [[package]]
 name = "haiku-skills"
-version = "0.14.0"
+version = "0.15.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "ag-ui-protocol" },
@@ -1614,9 +1614,9 @@ dependencies = [
     { name = "pyyaml" },
     { name = "skills-ref" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/89/c4/82a6b82f70726a2e759aad4c6c553309f2cc2ca7f3c157b8a15fd14da709/haiku_skills-0.14.0.tar.gz", hash = "sha256:27074a171060a0ecae6b89c6b7756b3b8ed0dfb957a3f5076ab1d158e68feb82", size = 250637, upload-time = "2026-04-16T08:47:33.152Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/86/a1/e2bd00a72d002f9db1c53c068167ed436a457dae0f8996399f116c087f6a/haiku_skills-0.15.0.tar.gz", hash = "sha256:ce93e6846e05397f5d96c144f956edd395b9e7308cb5cef6213c49c183a08bbc", size = 252030, upload-time = "2026-04-22T09:00:13.775Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e5/7a/bc53abcbae8bf1aa013379f49f6690fbaff55f80a7b997824103a5f44384/haiku_skills-0.14.0-py3-none-any.whl", hash = "sha256:698d0012bcf06f43499c30aaf6a3b7b772c66737696bfa3f93a6b113fb0f6b17", size = 31613, upload-time = "2026-04-16T08:47:31.68Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/68/3df2c9761fc0b0592b60f4723c87adeda5f335ea0835b4cf77ca07784379/haiku_skills-0.15.0-py3-none-any.whl", hash = "sha256:a1771b16e0ffe7da775f28d38c704e029f8e8757791616d7f27e73feb2c16fd0", size = 32041, upload-time = "2026-04-22T09:00:12.824Z" },
 ]
 
 [[package]]