Merge pull request #553 from ggozad/feat/session-reuse

Share the LanceDB session and open one client per server
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 15:11:24 +03:00 committed by GitHub
commit adcf27f650
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 705 additions and 130 deletions

View file

@ -5,12 +5,14 @@
- `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata.
- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes.
- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`.
- `lancedb.read_consistency_interval_seconds` (default 30), `lancedb.index_cache_size_bytes` and `lancedb.metadata_cache_size_bytes`. The LanceDB session is shared across connections in a process, so its index and metadata caches survive a connection being closed.
### Changed
- Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
- `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.
- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift.
- `import_documents` embeds chunks across the whole batch in one pass instead of per document.
- `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`.
- `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema.

View file

@ -122,6 +122,18 @@ The `storage_options` keys are case-insensitive and passed directly to the under
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization and vector indexing are still performed normally.
### Caching and Read Consistency
```yaml
lancedb:
read_consistency_interval_seconds: 30 # null to never re-check
index_cache_size_bytes: 536870912 # null for the LanceDB default
metadata_cache_size_bytes: 268435456
```
- **read_consistency_interval_seconds**: how often a connection checks for writes from another process. `null` never checks, so a long-lived reader never sees the ingester's writes. `0` checks on every read.
- **index_cache_size_bytes** / **metadata_cache_size_bytes**: sizes for the caches held by the LanceDB session, which is shared across every connection in the process. The first vector query loads the index into it, so on object storage the cache is what stops the next connection refetching it. Size it for the total set of indexes a process keeps warm, against the memory available to it.
### Deployment Pattern: One Writer, Many Readers
LanceDB on S3 supports **exactly one writer + N readers** per database URI. Multiple writers against the same URI can race on the manifest commit and corrupt state. This is a LanceDB property, not something `haiku.rag` enforces.

View file

@ -129,6 +129,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
state: StateT | None = field(default=None, repr=False)
outer_state: dict[str, Any] | None = field(default=None, repr=False)
rag: HaikuRAG | None = field(default=None, repr=False)
"""A connection this capability opened, and must close."""
borrowed_rag: HaikuRAG | None = field(default=None, repr=False)
"""A caller's connection, reused and never closed here."""
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
search_count: int = field(default=0, repr=False)
@ -347,6 +350,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
raise error
async def _ensure_rag(self) -> HaikuRAG:
if self.borrowed_rag is not None:
return self.borrowed_rag
if self.rag is None:
async with self.resource_lock:
if self.rag is None:

View file

@ -1,13 +1,16 @@
from dataclasses import dataclass, field
from functools import cache
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext, ToolFailed
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.capabilities._base import (
CodeExecutionEntry,
RAGCapabilityBase,
@ -162,6 +165,7 @@ def create_capability(
config: AppConfig | None = None,
*,
defer_loading: bool = True,
rag: "HaikuRAG | None" = None,
request_limit: int | None = 30,
vision: bool | None = None,
) -> AnalysisCapability:
@ -180,6 +184,7 @@ def create_capability(
return AnalysisCapability(
db_path=resolve_db_path(db_path, config),
config=config,
borrowed_rag=rag,
state_type=AnalysisState,
state_namespace=STATE_NAMESPACE,
instruction_text=instructions(),

View file

@ -1,13 +1,16 @@
from dataclasses import dataclass
from functools import cache
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.capabilities._base import (
RAGCapabilityBase,
resolve_db_path,
@ -72,6 +75,7 @@ def create_capability(
config: AppConfig | None = None,
*,
defer_loading: bool = True,
rag: "HaikuRAG | None" = None,
request_limit: int | None = 20,
vision: bool | None = None,
) -> RAGCapability:
@ -88,6 +92,7 @@ def create_capability(
return RAGCapability(
db_path=resolve_db_path(db_path, config),
config=config,
borrowed_rag=rag,
state_type=RAGState,
state_namespace=STATE_NAMESPACE,
instruction_text=instructions(),

View file

@ -63,6 +63,7 @@ async def ask(
capability = create_capability(
db_path=client.store.db_path,
config=client._config,
rag=client,
defer_loading=False,
)
deps = _AgentDeps(
@ -115,6 +116,7 @@ async def analyze(
capability = create_capability(
db_path=client.store.db_path,
config=client._config,
rag=client,
defer_loading=False,
)
deps = _AgentDeps(

View file

@ -62,10 +62,21 @@ class StorageConfig(BaseModel):
class LanceDBConfig(BaseModel):
"""LanceDB connection settings.
read_consistency_interval_seconds bounds how stale a reader may be. None
never re-checks, so a long-lived reader never sees another process's writes.
The cache sizes are per process, since the session is shared across
connections.
"""
uri: str = ""
api_key: str = ""
region: str = ""
storage_options: dict[str, str] = Field(default_factory=dict)
read_consistency_interval_seconds: float | None = Field(default=30, ge=0)
index_cache_size_bytes: int | None = Field(default=None, ge=0)
metadata_cache_size_bytes: int | None = Field(default=None, ge=0)
class EmbeddingsConfig(BaseModel):

View file

@ -1,3 +1,6 @@
import asyncio
from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager
from pathlib import Path
from typing import Any
@ -28,7 +31,42 @@ def create_mcp_server(
config: Configuration to use.
read_only: If True, write tools (add_document_*, delete_document) are not registered.
"""
mcp = FastMCP("haiku-rag")
client: HaikuRAG | None = None
stack = AsyncExitStack()
client_lock = asyncio.Lock()
async def _client() -> HaikuRAG:
"""The server's client, opened once.
Opening cost is per connection, and on object storage the first vector
query loads the index into the session cache, so a client per tool call
pays that repeatedly.
"""
nonlocal client
async with client_lock:
if client is None:
client = await stack.enter_async_context(
HaikuRAG(db_path, config=config, read_only=read_only)
)
return client
@asynccontextmanager
async def lifespan(_server: FastMCP) -> AsyncIterator[None]:
# Open eagerly so an unopenable database fails startup rather than
# every tool call.
nonlocal client
await _client()
try:
yield
finally:
# The lifespan can be re-entered; without the reset the next cycle
# hands out the closed client, including when aclose itself fails.
try:
await stack.aclose()
finally:
client = None
mcp = FastMCP("haiku-rag", lifespan=lifespan)
# Write tools - only registered when not in read-only mode
if not read_only:
@ -41,14 +79,14 @@ def create_mcp_server(
) -> str | None:
"""Add a document to the RAG system from a file path."""
try:
async with HaikuRAG(db_path, config=config) as rag:
result = await rag.create_document_from_source(
Path(file_path), title=title, metadata=metadata or {}
)
# Handle both single document and list of documents (directories)
if isinstance(result, list):
return result[0].id if result else None
return result.id
rag = await _client()
result = await rag.create_document_from_source(
Path(file_path), title=title, metadata=metadata or {}
)
# Handle both single document and list of documents (directories)
if isinstance(result, list):
return result[0].id if result else None
return result.id
except Exception:
return None
@ -58,14 +96,14 @@ def create_mcp_server(
) -> str | None:
"""Add a document to the RAG system from a URL."""
try:
async with HaikuRAG(db_path, config=config) as rag:
result = await rag.create_document_from_source(
url, title=title, metadata=metadata or {}
)
# Handle both single document and list of documents
if isinstance(result, list):
return result[0].id if result else None
return result.id
rag = await _client()
result = await rag.create_document_from_source(
url, title=title, metadata=metadata or {}
)
# Handle both single document and list of documents
if isinstance(result, list):
return result[0].id if result else None
return result.id
except Exception:
return None
@ -78,11 +116,11 @@ def create_mcp_server(
) -> str | None:
"""Add a document to the RAG system from text content."""
try:
async with HaikuRAG(db_path, config=config) as rag:
document = await rag.create_document(
content, uri, title=title, metadata=metadata or {}
)
return document.id
rag = await _client()
document = await rag.create_document(
content, uri, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@ -90,10 +128,8 @@ def create_mcp_server(
async def delete_document(document_id: str) -> bool:
"""Delete a document by its ID."""
try:
async with HaikuRAG(
db_path, config=config, skip_validation=True
) as rag:
return await rag.delete_document(document_id)
rag = await _client()
return await rag.delete_document(document_id)
except Exception:
return False
@ -110,10 +146,8 @@ def create_mcp_server(
response (smaller JSON payload for plain-text consumers).
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
return await rag.search(
query, limit=limit, include_images=include_images
)
rag = await _client()
return await rag.search(query, limit=limit, include_images=include_images)
except Exception:
return []
@ -145,10 +179,8 @@ def create_mcp_server(
except Exception:
return []
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
return await rag.search(
raw, limit=limit, include_images=include_images
)
rag = await _client()
return await rag.search(raw, limit=limit, include_images=include_images)
except Exception:
return []
@ -156,8 +188,8 @@ def create_mcp_server(
async def get_document(document_id: str) -> Document | None:
"""Get a document by its ID."""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
return await rag.get_document_by_id(document_id)
rag = await _client()
return await rag.get_document_by_id(document_id)
except Exception:
return None
@ -175,18 +207,18 @@ def create_mcp_server(
filter: Optional SQL WHERE clause to filter documents.
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
documents = await rag.list_documents(limit, offset, filter)
rag = await _client()
documents = await rag.list_documents(limit, offset, filter)
return [
DocumentInfo(
id=doc.id,
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in documents
]
return [
DocumentInfo(
id=doc.id,
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in documents
]
except Exception:
return []
@ -209,11 +241,11 @@ def create_mcp_server(
"""
try:
images = _decode_images(images_base64)
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
answer, citations = await rag.ask(question, images=images)
if cite and citations:
answer += "\n\n" + format_citations(citations)
return answer
rag = await _client()
answer, citations = await rag.ask(question, images=images)
if cite and citations:
answer += "\n\n" + format_citations(citations)
return answer
except Exception as e:
return f"Error answering question: {e!s}"
@ -240,9 +272,9 @@ def create_mcp_server(
"""
try:
images = _decode_images(images_base64)
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
result = await rag.analyze(question, filter=filter, images=images)
return result.answer
rag = await _client()
result = await rag.analyze(question, filter=filter, images=images)
return result.answer
except Exception as e:
return f"Error running analysis capability: {e!s}"

View file

@ -55,25 +55,56 @@ class ConnectionMode(Enum):
return ConnectionMode.OBJECT_STORAGE
_sessions: dict[tuple[int | None, int | None], lancedb.Session] = {}
def _session(config: AppConfig) -> lancedb.Session:
"""The process's session for these cache sizes.
Sessions hold the index and metadata caches. Sharing one across connections
is what keeps a cached index from being refetched per connection, which on
object storage is the dominant cost of the first query.
"""
key = (
config.lancedb.index_cache_size_bytes,
config.lancedb.metadata_cache_size_bytes,
)
if key not in _sessions:
kwargs = {}
if key[0] is not None:
kwargs["index_cache_size_bytes"] = key[0]
if key[1] is not None:
kwargs["metadata_cache_size_bytes"] = key[1]
_sessions[key] = lancedb.Session(**kwargs)
return _sessions[key]
async def connect_lancedb(
config: AppConfig, db_path: Path | None = None
) -> lancedb.AsyncConnection:
interval = config.lancedb.read_consistency_interval_seconds
kwargs: dict[str, Any] = {
"session": _session(config),
"read_consistency_interval": (
timedelta(seconds=interval) if interval is not None else None
),
}
mode = ConnectionMode.from_config(config)
if mode == ConnectionMode.CLOUD:
return await lancedb.connect_async(
uri=config.lancedb.uri,
api_key=config.lancedb.api_key,
region=config.lancedb.region,
**kwargs,
)
elif mode == ConnectionMode.OBJECT_STORAGE:
kwargs: dict[str, Any] = {"uri": config.lancedb.uri}
if config.lancedb.storage_options:
kwargs["storage_options"] = config.lancedb.storage_options
return await lancedb.connect_async(**kwargs)
return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs)
else:
if db_path is None:
raise ValueError("No lancedb.uri configured and no db_path provided")
return await lancedb.connect_async(db_path.absolute())
return await lancedb.connect_async(db_path.absolute(), **kwargs)
class DocumentRecord(LanceModel):
@ -176,6 +207,11 @@ def get_document_items_arrow_schema() -> pa.Schema:
return pa.schema(fields)
def _stored_vector_dim(settings: dict) -> int | None:
"""The vector dimension a database's chunks were written at."""
return settings.get("embeddings", {}).get("model", {}).get("vector_dim")
def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]:
"""The index set each table carries."""
match table_name:
@ -502,29 +538,27 @@ class Store:
self._config, self.db_path
)
# For remote stores (and as a safety net for local paths that exist but
# have no tables — e.g. a previously failed init), detect new DB by
# checking whether any tables exist.
is_new_db = self._is_new_db
if not is_new_db:
existing_tables = (await self.db.list_tables()).tables
if not existing_tables:
is_new_db = True
# Read once and thread onward: on object storage each of these is a
# round trip. A local path that exists with no tables is a failed init,
# so treat it as new.
existing_tables = (await self.db.list_tables()).tables
is_new_db = self._is_new_db or not existing_tables
# For existing databases, read stored vector dimension to create ChunkRecord
# that can read existing chunks. For new databases, use config's dimension.
stored_vector_dim = None
if not is_new_db:
stored_vector_dim = await self._get_stored_vector_dim()
stored_settings: dict = {}
if not is_new_db and "settings" in existing_tables:
self.settings_table = await self.db.open_table("settings")
stored_settings = await self._read_stored_settings()
# Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
# An existing database's chunks can only be read with the dimension they
# were written at.
stored_vector_dim = _stored_vector_dim(stored_settings)
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
# Initialize tables (creates them if they don't exist). For an existing
# DB this raises MigrationRequiredError up front when migrations are
# pending, before creating any newly-introduced table.
await self._init_tables(is_new_db)
await self._init_tables(is_new_db, existing_tables, stored_settings)
# Set version for new databases.
if is_new_db and not self._read_only:
@ -532,7 +566,7 @@ class Store:
# Validate config compatibility after connection is established
if not self._skip_validation:
await self._validate_configuration()
await self._validate_configuration(stored_settings)
async def __aenter__(self):
# If _initialize connects to LanceDB but then fails (e.g. migration
@ -554,33 +588,26 @@ class Store:
"""Whether the store is in read-only mode."""
return self._read_only
async def _get_stored_vector_dim(self) -> int | None:
"""Read the stored vector dimension from the settings table.
async def _read_stored_settings(self) -> dict:
"""The stored settings blob, or {} if it is absent or not a JSON object.
Returns:
The stored vector dimension, or None if not found.
Only decoding failures are tolerated. A storage failure must propagate:
read as empty settings it would look like version 0.0.0, and the
migration check would declare every migration pending.
"""
rows = (
await self.settings_table.query()
.where("id = 'settings'")
.limit(1)
.to_arrow()
).to_pylist()
if not rows or not rows[0].get("settings"):
return {}
try:
existing_tables = (await self.db.list_tables()).tables
if "settings" not in existing_tables:
return None
settings_table = await self.db.open_table("settings")
rows = (
await settings_table.query()
.where("id = 'settings'")
.limit(1)
.to_arrow()
).to_pylist()
if not rows or not rows[0].get("settings"):
return None
settings = json.loads(rows[0]["settings"])
embeddings = settings.get("embeddings", {})
model = embeddings.get("model", {})
return model.get("vector_dim")
except Exception:
return None
decoded = json.loads(rows[0]["settings"])
except (json.JSONDecodeError, TypeError):
return {}
return decoded if isinstance(decoded, dict) else {}
def _assert_writable(self) -> None:
"""Raise ReadOnlyError if the store is in read-only mode."""
@ -711,16 +738,19 @@ class Store:
except Exception as e:
logger.warning(f"Could not create vector index: {e}")
async def _validate_configuration(self) -> None:
async def _validate_configuration(
self, stored_settings: dict | None = None
) -> None:
"""Validate that the configuration is compatible with the database."""
from haiku.rag.store.repositories.settings import SettingsRepository
settings_repo = SettingsRepository(self)
await settings_repo.validate_config_compatibility()
await settings_repo.validate_config_compatibility(stored_settings)
async def _init_tables(self, is_new_db: bool):
async def _init_tables(
self, is_new_db: bool, existing_tables: list[str], stored_settings: dict
):
"""Initialize database tables (create if they don't exist)."""
existing_tables = (await self.db.list_tables()).tables
# Surface pending migrations BEFORE creating any newly-introduced table.
# Otherwise opening a legacy DB would either mutate it (creating an empty
@ -732,8 +762,7 @@ class Store:
and not self._skip_migration_check
and "settings" in existing_tables
):
self.settings_table = await self.db.open_table("settings")
await self._check_migrations()
await self._check_migrations(stored_settings.get("version", "0.0.0"))
missing_tables = set(REQUIRED_TABLES) - set(existing_tables)
@ -780,10 +809,8 @@ class Store:
)
await ensure_indexes(self.document_items_table, "document_items")
# Create or open settings table
if "settings" in existing_tables:
self.settings_table = await self.db.open_table("settings")
else:
# _initialize opened the settings table when the database had one.
if "settings" not in existing_tables:
self.settings_table = await self.db.create_table(
"settings", schema=SettingsRecord
)
@ -797,7 +824,7 @@ class Store:
"""Set the initial version for a new database."""
await self.set_haiku_version(metadata.version("haiku.rag-slim"))
async def _check_migrations(self) -> None:
async def _check_migrations(self, db_version: str) -> None:
"""Raise if migrations are pending. Opening never writes the version.
Raises:
@ -806,7 +833,6 @@ class Store:
from haiku.rag.store.upgrades import get_pending_upgrades
current_version = metadata.version("haiku.rag-slim")
db_version = await self.get_haiku_version()
pending = get_pending_upgrades(db_version)

View file

@ -60,7 +60,9 @@ class SettingsRepository:
)
await self.store.settings_table.add([settings_record])
async def validate_config_compatibility(self) -> None:
async def validate_config_compatibility(
self, stored_settings: dict | None = None
) -> None:
"""Validate the current configuration against stored settings without writing.
Opening a database never modifies it. ``vector_dim`` mismatches raise
@ -72,7 +74,8 @@ class SettingsRepository:
while a read-only open continues. Stored settings are reconciled
explicitly via ``haiku-rag rebuild --set-embedder``, never on open.
"""
stored_settings = await self.get_current_settings()
if stored_settings is None:
stored_settings = await self.get_current_settings()
# Nothing stored to validate against — never write on open.
if not stored_settings:

View file

@ -0,0 +1,132 @@
import pytest
from haiku.rag.capabilities.rag import create_capability
from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio
async def test_a_borrowed_client_is_reused_not_reopened(temp_db_path, monkeypatch):
"""A capability handed a client must not open a second connection to the
same database."""
from haiku.rag.store.engine import Store
async with HaikuRAG(temp_db_path, create=True) as client:
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
assert await capability._ensure_rag() is client
assert opens == 0
@pytest.mark.asyncio
async def test_closing_never_closes_a_borrowed_client(temp_db_path):
"""`_close` owns only what it opened. Closing the caller's client would be a
use-after-close for the caller."""
async with HaikuRAG(temp_db_path, create=True) as client:
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
await capability._ensure_rag()
await capability._close()
# Still usable by its owner.
assert await client.list_documents() == []
@pytest.mark.asyncio
async def test_a_borrowed_client_survives_for_run(temp_db_path):
"""for_run clears the owned connection per run; a borrowed one is the
caller's and carries into the run copy."""
from tests.capabilities.test_capabilities import Deps, make_context
async with HaikuRAG(temp_db_path, create=True) as client:
capability = create_capability(
db_path=client.store.db_path, config=client._config, rag=client
)
run_capability = await capability.for_run(make_context(Deps()))
assert run_capability is not capability
assert run_capability.rag is None
assert run_capability.borrowed_rag is client
assert await run_capability._ensure_rag() is client
@pytest.mark.asyncio
async def test_ask_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
"""`ask` built the capability from a db_path alone, so the capability opened
its own connection to a database the client already had open."""
from haiku.rag.capabilities import rag as rag_capability
from haiku.rag.store.engine import Store
real = rag_capability.create_capability
built = {}
def spy(**kwargs):
built["capability"] = real(**kwargs)
raise RuntimeError("stop before running the agent")
async with HaikuRAG(temp_db_path, create=True) as client:
monkeypatch.setattr(rag_capability, "create_capability", spy)
with pytest.raises(RuntimeError, match="stop before running the agent"):
await client.ask("anything")
capability = built["capability"]
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
assert await capability._ensure_rag() is client
assert opens == 0
@pytest.mark.asyncio
async def test_analyze_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
from haiku.rag.capabilities import analysis as analysis_capability
from haiku.rag.store.engine import Store
real = analysis_capability.create_capability
built = {}
def spy(**kwargs):
built["capability"] = real(**kwargs)
raise RuntimeError("stop before running the agent")
async with HaikuRAG(temp_db_path, create=True) as client:
monkeypatch.setattr(analysis_capability, "create_capability", spy)
with pytest.raises(RuntimeError, match="stop before running the agent"):
await client.analyze("anything")
capability = built["capability"]
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
assert await capability._ensure_rag() is client
assert opens == 0

View file

@ -0,0 +1,71 @@
import lancedb
import pytest
from haiku.rag.store.engine import Store
@pytest.fixture
def counts(monkeypatch):
"""Count the connection-level calls an open makes."""
tally: dict[str, int] = {"list_tables": 0, "open_settings": 0, "settings_query": 0}
list_tables = lancedb.AsyncConnection.list_tables
open_table = lancedb.AsyncConnection.open_table
query = lancedb.AsyncTable.query
async def counted_list_tables(self, *args, **kwargs):
tally["list_tables"] += 1
return await list_tables(self, *args, **kwargs)
async def counted_open_table(self, name, *args, **kwargs):
if name == "settings":
tally["open_settings"] += 1
return await open_table(self, name, *args, **kwargs)
def counted_query(self):
if self.name == "settings":
tally["settings_query"] += 1
return query(self)
monkeypatch.setattr(lancedb.AsyncConnection, "list_tables", counted_list_tables)
monkeypatch.setattr(lancedb.AsyncConnection, "open_table", counted_open_table)
monkeypatch.setattr(lancedb.AsyncTable, "query", counted_query)
return tally
@pytest.mark.asyncio
async def test_reopening_reads_the_table_list_and_settings_once(temp_db_path, counts):
async with Store(temp_db_path, create=True):
pass
for key in counts:
counts[key] = 0
async with Store(temp_db_path):
pass
assert counts["list_tables"] == 1
assert counts["open_settings"] == 1
assert counts["settings_query"] == 1
@pytest.mark.asyncio
async def test_storage_failures_propagate(temp_db_path):
"""A read failure must not read as empty settings: the migration check would
then see version 0.0.0 and declare every migration pending."""
async with Store(temp_db_path, create=True) as store:
def boom():
raise RuntimeError("s3 is having a day")
store.settings_table.query = boom
with pytest.raises(RuntimeError, match="s3 is having a day"):
await store._read_stored_settings()
@pytest.mark.asyncio
async def test_non_dict_settings_read_as_empty(temp_db_path):
async with Store(temp_db_path, create=True) as store:
await store.settings_table.update({"settings": "[]"}, where="id = 'settings'")
assert await store._read_stored_settings() == {}

View file

@ -1,6 +1,8 @@
from datetime import timedelta
from unittest.mock import AsyncMock, patch
import pytest
from pydantic import ValidationError
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, LanceDBConfig
@ -49,7 +51,8 @@ class TestConnectLancedb:
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config, db_path=temp_db_path)
mock_connect.assert_called_once_with(temp_db_path.absolute())
mock_connect.assert_awaited_once()
assert mock_connect.call_args.args == (temp_db_path.absolute(),)
@pytest.mark.asyncio
async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch):
@ -62,7 +65,8 @@ class TestConnectLancedb:
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config, db_path=relative)
mock_connect.assert_called_once_with(relative.absolute())
mock_connect.assert_awaited_once()
assert mock_connect.call_args.args == (relative.absolute(),)
@pytest.mark.asyncio
async def test_cloud_passes_uri_api_key_region(self):
@ -75,9 +79,11 @@ class TestConnectLancedb:
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="db://my-database", api_key="test-key", region="us-west-2"
)
mock_connect.assert_awaited_once()
kwargs = mock_connect.call_args.kwargs
assert kwargs["uri"] == "db://my-database"
assert kwargs["api_key"] == "test-key"
assert kwargs["region"] == "us-west-2"
@pytest.mark.asyncio
async def test_object_storage_passes_uri_and_storage_options(self):
@ -94,13 +100,13 @@ class TestConnectLancedb:
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="s3://bucket/path",
storage_options={
"endpoint": "http://minio:9000",
"region": "us-east-1",
},
)
mock_connect.assert_awaited_once()
kwargs = mock_connect.call_args.kwargs
assert kwargs["uri"] == "s3://bucket/path"
assert kwargs["storage_options"] == {
"endpoint": "http://minio:9000",
"region": "us-east-1",
}
@pytest.mark.asyncio
async def test_object_storage_without_storage_options(self):
@ -109,7 +115,10 @@ class TestConnectLancedb:
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(uri="s3://bucket/path")
mock_connect.assert_awaited_once()
kwargs = mock_connect.call_args.kwargs
assert kwargs["uri"] == "s3://bucket/path"
assert "storage_options" not in kwargs
@pytest.mark.asyncio
async def test_local_without_db_path_raises(self):
@ -267,7 +276,7 @@ class TestInitFailureCleanup:
async def fake_connect(*args, **kwargs):
return mock_conn
async def failing_init_tables(self, is_new_db):
async def failing_init_tables(self, *args):
raise RuntimeError("simulated table init failure")
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
@ -381,7 +390,7 @@ class TestStoreMiscellany:
{"settings": "not json at all"}, where="id = 'settings'"
)
assert await store._get_stored_vector_dim() is None
assert await store._read_stored_settings() == {}
@pytest.mark.asyncio
async def test_vacuum_skips_when_already_running(self, temp_db_path):
@ -398,3 +407,119 @@ class TestStoreMiscellany:
async with Store(temp_db_path, create=True) as store:
with pytest.raises(ValueError, match="Unknown table"):
await store.list_table_versions("not_a_table")
class TestSessionAndConsistency:
@pytest.mark.asyncio
async def test_session_is_shared_across_connections(self):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
await connect_lancedb(config)
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
assert sessions[0] is sessions[1]
@pytest.mark.asyncio
async def test_cache_sizes_select_distinct_sessions(self):
small = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path", index_cache_size_bytes=1 << 20
)
)
large = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path", index_cache_size_bytes=1 << 30
)
)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(small)
await connect_lancedb(large)
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
assert sessions[0] is not sessions[1]
@pytest.mark.asyncio
async def test_both_cache_sizes_are_applied(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
index_cache_size_bytes=2 << 20,
metadata_cache_size_bytes=4 << 20,
)
)
with (
patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
),
patch("haiku.rag.store.engine.lancedb.Session") as mock_session,
):
await connect_lancedb(config)
mock_session.assert_called_once_with(
index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20
)
@pytest.mark.asyncio
async def test_read_consistency_interval_is_forwarded(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path", read_consistency_interval_seconds=5
)
)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
seconds=5
)
@pytest.mark.asyncio
async def test_read_consistency_interval_omitted_when_disabled(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path", read_consistency_interval_seconds=None
)
)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
assert mock_connect.call_args.kwargs["read_consistency_interval"] is None
@pytest.mark.asyncio
async def test_local_connection_also_gets_session_and_consistency(self, tmp_path):
config = AppConfig()
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config, tmp_path / "db.lancedb")
assert mock_connect.call_args.kwargs["session"] is not None
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
seconds=30
)
class TestLanceDBConfigValidation:
def test_negative_values_are_rejected(self):
"""Negatives overflow or panic inside Lance, so reject them here."""
with pytest.raises(ValidationError):
LanceDBConfig(read_consistency_interval_seconds=-1)
with pytest.raises(ValidationError):
LanceDBConfig(index_cache_size_bytes=-1)
with pytest.raises(ValidationError):
LanceDBConfig(metadata_cache_size_bytes=-1)
def test_zero_is_allowed(self):
config = LanceDBConfig(
read_consistency_interval_seconds=0, index_cache_size_bytes=0
)
assert config.read_consistency_interval_seconds == 0

View file

@ -481,3 +481,147 @@ class TestMCPToolsDegradeOnError:
assert "AI Overview" in with_cite
assert await ask(question="q", cite=False) == "the answer"
class TestMCPClientLifetime:
@pytest.mark.asyncio
async def test_tool_calls_share_one_database_open(self, mcp_db, monkeypatch):
from haiku.rag.store.engine import Store
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
mcp = create_mcp_server(mcp_db, read_only=True)
search = await _get_tool(mcp, "search_documents")
list_docs = await _get_tool(mcp, "list_documents")
await search(query="artificial intelligence")
await list_docs()
await search(query="machine learning")
assert opens == 1
@pytest.mark.asyncio
async def test_concurrent_reads_share_one_open(self, mcp_db, monkeypatch):
import asyncio
from haiku.rag.store.engine import Store
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
mcp = create_mcp_server(mcp_db, read_only=True)
list_docs = await _get_tool(mcp, "list_documents")
results = await asyncio.gather(*(list_docs() for _ in range(5)))
assert opens == 1
assert all(len(r) == 2 for r in results)
@pytest.mark.asyncio
async def test_a_write_is_visible_to_the_next_read(self, mcp_db):
"""One connection sees its own writes, whatever the consistency interval."""
mcp = create_mcp_server(mcp_db, read_only=False)
list_docs = await _get_tool(mcp, "list_documents")
delete_doc = await _get_tool(mcp, "delete_document")
docs = await list_docs()
assert await delete_doc(document_id=docs[0].id) is True
assert len(await list_docs()) == len(docs) - 1
@pytest.mark.asyncio
async def test_lifespan_opens_and_closes_once(self, mcp_db, monkeypatch):
from haiku.rag.store.engine import Store
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
mcp = create_mcp_server(mcp_db, read_only=True)
# _lifespan_manager is what every transport enters; the public
# lifespan() combines provider lifespans only.
async with mcp._lifespan_manager():
assert opens == 1, "startup should open the database, not the first call"
search = await _get_tool(mcp, "search_documents")
await search(query="artificial intelligence")
assert opens == 1
assert opens == 1
@pytest.mark.asyncio
async def test_startup_fails_when_the_database_cannot_open(self, tmp_path):
mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb", read_only=True)
with pytest.raises(FileNotFoundError):
async with mcp._lifespan_manager():
pass
@pytest.mark.asyncio
async def test_a_second_lifespan_cycle_opens_a_fresh_client(
self, mcp_db, monkeypatch
):
from haiku.rag.store.engine import Store
opens = 0
initialize = Store._initialize
async def counted(self):
nonlocal opens
opens += 1
return await initialize(self)
monkeypatch.setattr(Store, "_initialize", counted)
mcp = create_mcp_server(mcp_db, read_only=True)
search = await _get_tool(mcp, "search_documents")
async with mcp._lifespan_manager():
await search(query="artificial intelligence")
assert opens == 1
async with mcp._lifespan_manager():
results = await search(query="artificial intelligence")
assert opens == 2
assert len(results) > 0
@pytest.mark.asyncio
async def test_same_dim_drift_starts_read_only_but_not_writable(self, mcp_db):
"""Validation is unchanged: same-dimension identity drift warns in
read-only mode and raises in writable mode. The MCP server no longer
opts out of it for deletion."""
from haiku.rag.config import Config
from haiku.rag.store.repositories.settings import ConfigMismatchError
drifted = Config.model_copy(deep=True)
drifted.embeddings.model.name = "a-different-model"
async with create_mcp_server(
mcp_db, config=drifted, read_only=True
)._lifespan_manager():
pass
with pytest.raises(ConfigMismatchError):
async with create_mcp_server(
mcp_db, config=drifted, read_only=False
)._lifespan_manager():
pass