Merge pull request #348 from ggozad/feat/skills-lifespan

Use haiku.skills lifespan for per-invocation client, sandbox, state

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

@@ -277,11 +277,8 @@ class Sandbox:
                     external_functions=self._build_external_functions(),
                     os=self._vfs,
                 )
-        repl = self._repl
+        assert self._repl is not None and self._vfs is not None
-        vfs = self._vfs
+        return self._repl, self._vfs
-        if repl is None or vfs is None:
-            raise RuntimeError("Sandbox initialization failed")
-        return repl, vfs
 
     async def execute(self, code: str) -> SandboxResult:
         """Execute Python code in the Monty REPL.

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -7,7 +7,7 @@ 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.skills._deps import AnalysisRunDeps, RAGRunDeps
 from haiku.rag.store.models.chunk import SearchResult
 
 
@@ -76,10 +76,9 @@ def _get_state(ctx: RunContext[RAGRunDeps], state_type: type[BaseModel]) -> Any:
 
 
 def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
-    if ctx.deps is None or ctx.deps.rag is None:
+    assert ctx.deps is not None and ctx.deps.rag is not None, (
-        raise RuntimeError(
+        "RAGRunDeps.rag is not set — skill lifespan must run before tools."
-            "RAGRunDeps.rag is not set — skill lifespan must run before tools."
+    )
-        )
     return ctx.deps.rag
 
 
@@ -227,7 +226,6 @@ def create_skill_tools(
         tools["get_document"] = get_document
 
     if "execute_code" in tool_names:
-        from haiku.rag.skills._deps import AnalysisRunDeps
 
         async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
             """Execute Python code in a sandboxed interpreter.
@@ -242,10 +240,9 @@ def create_skill_tools(
             Args:
                 code: Python code to execute.
             """
-            if ctx.deps is None or ctx.deps.sandbox is None:
+            assert ctx.deps is not None and ctx.deps.sandbox is not None, (
-                raise RuntimeError(
+                "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
-                    "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
+            )
-                )
             sandbox = ctx.deps.sandbox
             result = await sandbox.execute(code)
 

tests/skills/test_analysis.py

@@ -266,6 +266,19 @@ class TestAnalysisLifespan:
         assert skill.deps_type is AnalysisRunDeps
         assert skill.lifespan is not None
 
+    async def test_run_skill_end_to_end_opens_client_and_sandbox(
+        self, allow_model_requests, rag_db
+    ):
+        """Full sub-agent path: lifespan builds client + sandbox, tools see them."""
+        from pydantic_ai.models.test import TestModel
+
+        from haiku.rag.skills.analysis import create_skill
+        from haiku.skills.agent import _run_skill
+
+        skill = create_skill(db_path=rag_db)
+        result, *_ = await _run_skill(TestModel(), skill, "Print the document count.")
+        assert result
+
     async def test_lifespan_clears_executions_citations_searches(self, rag_db):
         from haiku.rag.agents.research.models import Citation
         from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan

tests/skills/test_rag.py

@@ -359,6 +359,19 @@ class TestLifespan:
         assert skill.deps_type is RAGRunDeps
         assert skill.lifespan is not None
 
+    async def test_run_skill_end_to_end_opens_and_closes_client(
+        self, allow_model_requests, rag_db
+    ):
+        """Full sub-agent path: lifespan opens the client, tools see it, lifespan closes it."""
+        from pydantic_ai.models.test import TestModel
+
+        from haiku.rag.skills.rag import create_skill
+        from haiku.skills.agent import _run_skill
+
+        skill = create_skill(db_path=rag_db)
+        result, *_ = await _run_skill(TestModel(), skill, "List the documents.")
+        assert result
+
     async def test_lifespan_clears_citations_and_searches_but_keeps_index(self, rag_db):
         from haiku.rag.agents.research.models import Citation
         from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan