scope citations, searches, and executions to the current invocation

This commit is contained in:
Yiorgis Gozadinos 2026-04-22 13:49:24 +03:00
parent 4a9dd9b49a
commit f0016ebcd2
No known key found for this signature in database
6 changed files with 100 additions and 9 deletions

View file

@ -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. - **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. - **`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. - **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 ## [0.41.0] - 2026-04-20

View file

@ -37,10 +37,11 @@ class AnalysisState(BaseModel):
searches: dict[str, list[SearchResult]] = {} searches: dict[str, list[SearchResult]] = {}
``` ```
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. - **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. - **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** — Same per-turn citation tracking as the RAG skill. - **citation_index** — Citations indexed by chunk ID. Accumulates across invocations (same semantics as the RAG skill).
- **searches** — Search results from both the `search` tool and sandbox-internal searches. - **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 ## Usage with RAG Skill

View file

@ -36,7 +36,7 @@ class RAGState(BaseModel):
searches: dict[str, list[SearchResult]] = {} searches: dict[str, list[SearchResult]] = {}
``` ```
- **citation_index** — All citations indexed by chunk ID (deduplicated across turns). - **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical turns' chunk IDs remain resolvable in the UI scrollback.
- **citations**Per-turn lists of chunk IDs registered via the `cite` tool. - **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. Set this to scope queries to specific documents. - **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. - **searches** — Search results keyed by query string. Cleared at the start of each invocation.

View file

@ -2,7 +2,7 @@ from collections.abc import AsyncIterator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Any
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.skills.state import SkillRunDeps from haiku.skills.state import SkillRunDeps
@ -23,6 +23,26 @@ class AnalysisRunDeps(RAGRunDeps):
sandbox: "Sandbox | None" = None 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): def make_rag_lifespan(db_path: Path, config: AppConfig):
@asynccontextmanager @asynccontextmanager
async def lifespan(deps: RAGRunDeps) -> AsyncIterator[None]: 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: async with HaikuRAG(db_path, config=config, read_only=True) as rag:
deps.rag = rag deps.rag = rag
deps.search_count = 0 deps.search_count = 0
_reset_invocation_state(deps.state)
yield yield
return lifespan return lifespan
@ -52,6 +73,7 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
config=config, config=config,
context=AnalysisContext(filter=doc_filter), context=AnalysisContext(filter=doc_filter),
) )
_reset_invocation_state(deps.state)
yield yield
return lifespan return lifespan

View file

@ -265,3 +265,38 @@ class TestAnalysisLifespan:
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.deps_type is AnalysisRunDeps assert skill.deps_type is AnalysisRunDeps
assert skill.lifespan is not None 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'"

View file

@ -358,3 +358,35 @@ class TestLifespan:
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.deps_type is RAGRunDeps assert skill.deps_type is RAGRunDeps
assert skill.lifespan is not None 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'"