From f0016ebcd2ebbf4189cca28a16697833d2b3c524 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 22 Apr 2026 13:49:24 +0300 Subject: [PATCH] scope citations, searches, and executions to the current invocation --- CHANGELOG.md | 1 + docs/skills/analysis.md | 9 +++--- docs/skills/rag.md | 8 +++--- haiku_rag_slim/haiku/rag/skills/_deps.py | 24 +++++++++++++++- tests/skills/test_analysis.py | 35 ++++++++++++++++++++++++ tests/skills/test_rag.py | 32 ++++++++++++++++++++++ 6 files changed, 100 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a441e78c..acfaccac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - **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. - **Analysis sandbox persists variables across `execute_code` calls within one invocation.** Re-enables the incremental-exploration workflow (search in one call, process results in the next). Each new skill invocation constructs a fresh `Sandbox` via the analysis lifespan, so there is no cross-invocation leak. +- **Skill state is scoped to the current invocation.** Lifespans now clear `citations`, `searches`, and (for analysis) `executions` at the start of each invocation, so state deltas sent to the AG-UI client reflect only the in-progress turn. `citation_index` is preserved across invocations so past-turn citation chunk ids remain resolvable, and `document_filter` is preserved as session-level config. ## [0.41.0] - 2026-04-20 diff --git a/docs/skills/analysis.md b/docs/skills/analysis.md index 00723540..3a90e310 100644 --- a/docs/skills/analysis.md +++ b/docs/skills/analysis.md @@ -37,10 +37,11 @@ class AnalysisState(BaseModel): searches: dict[str, list[SearchResult]] = {} ``` -- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. -- **executions** — Each `execute_code` call appends a `CodeExecutionEntry` with code, stdout, stderr, and success status. -- **citation_index** / **citations** — Same per-turn citation tracking as the RAG skill. -- **searches** — Search results from both the `search` tool and sandbox-internal searches. +- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Persists across invocations as session-level configuration. +- **executions** — Each `execute_code` call appends a `CodeExecutionEntry` with code, stdout, stderr, and success status. Cleared at the start of each invocation; mirrors the sandbox lifecycle (variables persist across calls within one invocation, a fresh sandbox is built per invocation). +- **citation_index** — Citations indexed by chunk ID. Accumulates across invocations (same semantics as the RAG skill). +- **citations** — Cleared at the start of each invocation; holds only the in-progress turn. +- **searches** — Search results from both the `search` tool and sandbox-internal searches. Cleared at the start of each invocation. ## Usage with RAG Skill diff --git a/docs/skills/rag.md b/docs/skills/rag.md index 064ea86c..90917d44 100644 --- a/docs/skills/rag.md +++ b/docs/skills/rag.md @@ -36,7 +36,7 @@ class RAGState(BaseModel): searches: dict[str, list[SearchResult]] = {} ``` -- **citation_index** — All citations indexed by chunk ID (deduplicated across turns). -- **citations** — Per-turn lists of chunk IDs registered via the `cite` tool. -- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Set this to scope queries to specific documents. -- **searches** — Search results keyed by query string. +- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical turns' chunk IDs remain resolvable in the UI scrollback. +- **citations** — Chunk IDs registered via the `cite` tool. Cleared at the start of each invocation; holds only the in-progress turn. +- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Persists across invocations as session-level configuration. +- **searches** — Search results keyed by query string. Cleared at the start of each invocation. diff --git a/haiku_rag_slim/haiku/rag/skills/_deps.py b/haiku_rag_slim/haiku/rag/skills/_deps.py index 8f8795d3..fdad98a8 100644 --- a/haiku_rag_slim/haiku/rag/skills/_deps.py +++ b/haiku_rag_slim/haiku/rag/skills/_deps.py @@ -2,7 +2,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from haiku.rag.config.models import AppConfig from haiku.skills.state import SkillRunDeps @@ -23,6 +23,26 @@ class AnalysisRunDeps(RAGRunDeps): sandbox: "Sandbox | None" = None +def _reset_invocation_state(state: Any) -> None: + """Clear state fields scoped to a single invocation. + + Keeps ``citation_index`` (accumulates resolved citations across the session + for lookup) and ``document_filter`` (session-level). Clears ``citations``, + ``searches``, and (for analysis) ``executions``. + """ + if state is None: + return + citations = getattr(state, "citations", None) + if citations is not None: + citations.clear() + searches = getattr(state, "searches", None) + if searches is not None: + searches.clear() + executions = getattr(state, "executions", None) + if executions is not None: + executions.clear() + + def make_rag_lifespan(db_path: Path, config: AppConfig): @asynccontextmanager async def lifespan(deps: RAGRunDeps) -> AsyncIterator[None]: @@ -31,6 +51,7 @@ def make_rag_lifespan(db_path: Path, config: AppConfig): async with HaikuRAG(db_path, config=config, read_only=True) as rag: deps.rag = rag deps.search_count = 0 + _reset_invocation_state(deps.state) yield return lifespan @@ -52,6 +73,7 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig): config=config, context=AnalysisContext(filter=doc_filter), ) + _reset_invocation_state(deps.state) yield return lifespan diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py index 1119b4d8..273c6d67 100644 --- a/tests/skills/test_analysis.py +++ b/tests/skills/test_analysis.py @@ -265,3 +265,38 @@ class TestAnalysisLifespan: skill = create_skill(config=test_app_config, db_path=temp_db_path) assert skill.deps_type is AnalysisRunDeps assert skill.lifespan is not None + + 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 + from haiku.rag.skills._tools import CodeExecutionEntry + from haiku.rag.skills.analysis import AnalysisState + + config = AppConfig() + lifespan = make_analysis_lifespan(rag_db, config) + + state = AnalysisState( + document_filter="title = 'AI Overview'", + executions=[CodeExecutionEntry(code="prior", stdout="", success=True)], + citation_index={ + "c1": Citation( + index=1, + chunk_id="c1", + document_id="d1", + document_title="t", + document_uri="u", + content="x", + page_numbers=[], + headings=[], + ) + }, + citations=[["c1"]], + searches={"prior": []}, + ) + deps = AnalysisRunDeps(state=state) + async with lifespan(deps): + assert state.executions == [] + assert state.citations == [] + assert state.searches == {} + assert "c1" in state.citation_index + assert state.document_filter == "title = 'AI Overview'" diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index e73cb34d..4890583f 100644 --- a/tests/skills/test_rag.py +++ b/tests/skills/test_rag.py @@ -358,3 +358,35 @@ class TestLifespan: skill = create_skill(config=test_app_config, db_path=temp_db_path) assert skill.deps_type is RAGRunDeps assert skill.lifespan is not None + + 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 + from haiku.rag.skills.rag import RAGState + + config = AppConfig() + lifespan = make_rag_lifespan(rag_db, config) + + state = RAGState( + document_filter="title = 'AI Overview'", + citation_index={ + "c1": Citation( + index=1, + chunk_id="c1", + document_id="d1", + document_title="t", + document_uri="u", + content="x", + page_numbers=[], + headings=[], + ) + }, + citations=[["c1"]], + searches={"prior": []}, + ) + deps = RAGRunDeps(state=state) + async with lifespan(deps): + assert state.citations == [] + assert state.searches == {} + assert "c1" in state.citation_index # preserved for cross-turn lookup + assert state.document_filter == "title = 'AI Overview'"