diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py
index 739a29be..eb0128c7 100644
--- a/haiku_rag_slim/haiku/rag/client/__init__.py
+++ b/haiku_rag_slim/haiku/rag/client/__init__.py
@@ -139,7 +139,6 @@ class HaikuRAG:
         self._read_only = read_only
         self._requested_sources = sources
         self._clients: dict[str, HaikuRAG] = {}
-        self._clients_lock = asyncio.Lock()
         self._source: str | None = None
         self._session: SingleDatabaseSession | FederatedSession | None = None
         self._owns_session = True
@@ -418,14 +417,6 @@ class HaikuRAG:
         if cached is not None:
             await aclose_quietly(cached, name)
 
-    async def _await_vacuum_tasks(self) -> None:
-        if isinstance(self._session, SingleDatabaseSession):
-            await self._session.drain_vacuum()
-
-    def _schedule_vacuum(self) -> None:
-        if isinstance(self._session, SingleDatabaseSession):
-            self._session.schedule_vacuum()
-
     # =========================================================================
     # Processing Primitives
     # =========================================================================
@@ -487,9 +478,9 @@ class HaikuRAG:
     ) -> Document:
         from haiku.rag.client.documents import create_document
 
-        self._single_session("create_document")
+        session = self._single_session("create_document")
 
-        return await create_document(self, content, uri, title, metadata, format)
+        return await create_document(session, content, uri, title, metadata, format)
 
     async def import_document(
         self,
@@ -501,10 +492,10 @@ class HaikuRAG:
     ) -> Document:
         from haiku.rag.client.documents import import_document
 
-        self._single_session("import_document")
+        session = self._single_session("import_document")
 
         return await import_document(
-            self, docling_document, chunks, uri, title, metadata
+            session, docling_document, chunks, uri, title, metadata
         )
 
     async def import_documents(
@@ -513,9 +504,9 @@ class HaikuRAG:
     ) -> list[Document]:
         from haiku.rag.client.documents import import_documents
 
-        self._single_session("import_documents")
+        session = self._single_session("import_documents")
 
-        return await import_documents(self, imports)
+        return await import_documents(session, imports)
 
     async def create_document_from_source(
         self,
@@ -530,10 +521,10 @@ class HaikuRAG:
     ) -> Document | list[Document]:
         from haiku.rag.client.documents import create_document_from_source
 
-        self._single_session("create_document_from_source")
+        session = self._single_session("create_document_from_source")
 
         return await create_document_from_source(
-            self,
+            session,
             source,
             title,
             metadata,
@@ -556,10 +547,10 @@ class HaikuRAG:
     ) -> Document:
         from haiku.rag.client.documents import update_document
 
-        self._single_session("update_document")
+        session = self._single_session("update_document")
 
         return await update_document(
-            self,
+            session,
             document_id,
             content,
             metadata,
@@ -841,9 +832,9 @@ class HaikuRAG:
     ) -> AsyncGenerator[str, None]:
         from haiku.rag.client.rebuild import rebuild_database
 
-        self._single_session("rebuild_database")
+        session = self._single_session("rebuild_database")
 
-        async for doc_id in rebuild_database(self, mode):
+        async for doc_id in rebuild_database(session, mode):
             yield doc_id
 
     async def vacuum(self) -> None:
diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py
index d049a723..31f5498b 100644
--- a/haiku_rag_slim/haiku/rag/client/documents.py
+++ b/haiku_rag_slim/haiku/rag/client/documents.py
@@ -11,9 +11,12 @@ from urllib.parse import quote, urlparse
 from haiku.rag.client.exceptions import UnsupportedSourceError
 from haiku.rag.client.processing import (
     _write_fetch_body,
+    chunk,
+    convert,
     ensure_chunks_embedded,
     get_extension_from_content_type_or_url,
 )
+from haiku.rag.client.session import SingleDatabaseSession
 from haiku.rag.client.titles import resolve_title
 from haiku.rag.converters import get_converter
 from haiku.rag.store.models.chunk import Chunk
@@ -25,7 +28,6 @@ from haiku.rag.uri import is_local_uri, uri_to_path
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
 
-    from haiku.rag.client import HaikuRAG
     from haiku.rag.ingester.metadata import MetadataProvider
     from haiku.rag.sources.base import FetchResult, Source
 
@@ -86,7 +88,9 @@ async def _prepare_document_from_docling(
 
 
 async def _prepare_and_title(
-    client: "HaikuRAG", document: Document, docling_document: "DoclingDocument"
+    session: SingleDatabaseSession,
+    document: Document,
+    docling_document: "DoclingDocument",
 ) -> None:
     """Fill the document from its converted form and title it if it has none.
 
@@ -97,10 +101,27 @@ async def _prepare_and_title(
     stored_content = await _prepare_document_from_docling(document, docling_document)
     if document.title is None:
         document.title = await resolve_title(
-            client._config, docling_document, stored_content
+            session.config, docling_document, stored_content
         )
 
 
+async def chunk_document(
+    session: SingleDatabaseSession,
+    docling_document: "DoclingDocument",
+    *,
+    existing_picture_data: dict[str, bytes] | None = None,
+    document_id: str | None = None,
+) -> list[Chunk]:
+    """Chunk and embed through the database this write belongs to."""
+    return await chunk(
+        session.config,
+        docling_document,
+        embedder=session.store.embedder,
+        existing_picture_data=existing_picture_data,
+        document_id=document_id,
+    )
+
+
 def parent_uri_filter(parent_uri: str) -> str:
     """SQL `WHERE` clause matching documents whose ``metadata.parent_uri``
     equals ``parent_uri``. ``metadata`` is stored as a JSON string produced by
@@ -112,7 +133,7 @@ def parent_uri_filter(parent_uri: str) -> str:
 
 
 async def _store_document_with_chunks(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     document: Document,
     chunks: list[Chunk],
     docling_document: "DoclingDocument",
@@ -121,16 +142,18 @@ async def _store_document_with_chunks(
 
     Handles versioning/rollback on failure.
     """
-    chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
+    chunks = await ensure_chunks_embedded(
+        session.config, chunks, session.store.embedder
+    )
     items = await asyncio.to_thread(extract_items, "", docling_document)
 
-    async with client.store.write_transaction():
+    async with session.store.write_transaction():
         # A concurrent ingestion of the same URI may have created the document
         # while this one was converting/embedding outside the lock. LanceDB has
         # no unique constraint on `uri`, so re-check under the lock and update in
         # place rather than inserting a duplicate.
         existing = (
-            await client.get_document_by_uri(document.uri)
+            await session.get_document_by_uri(document.uri)
             if document.uri is not None
             else None
         )
@@ -138,9 +161,9 @@ async def _store_document_with_chunks(
         if existing is not None:
             document.id = existing.id
             document.created_at = existing.created_at
-            stored_doc = await client.document_repository.update(document)
+            stored_doc = await session.document_repository.update(document)
         else:
-            stored_doc = await client.document_repository.create(document)
+            stored_doc = await session.document_repository.create(document)
 
         assert stored_doc.id is not None, "Document ID should not be None after storing"
         for order, chunk in enumerate(chunks):
@@ -150,22 +173,22 @@ async def _store_document_with_chunks(
             item.document_id = stored_doc.id
 
         if existing is not None:
-            await client.chunk_repository.replace_for_document(stored_doc.id, chunks)
-            await client.document_item_repository.replace_for_document(
+            await session.chunk_repository.replace_for_document(stored_doc.id, chunks)
+            await session.document_item_repository.replace_for_document(
                 stored_doc.id, items
             )
         else:
-            await client.chunk_repository.create(chunks)
-            await client.document_item_repository.create_items(stored_doc.id, items)
+            await session.chunk_repository.create(chunks)
+            await session.document_item_repository.create_items(stored_doc.id, items)
 
-    if client._config.storage.auto_vacuum:
-        client._schedule_vacuum()
+    if session.config.storage.auto_vacuum:
+        session.schedule_vacuum()
 
     return stored_doc
 
 
 async def _update_document_with_chunks(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     document: Document,
     chunks: list[Chunk],
     docling_document: "DoclingDocument | None" = None,
@@ -183,10 +206,12 @@ async def _update_document_with_chunks(
     existing_picture_data: dict[str, bytes] | None = None
     if docling_document is not None:
         existing_picture_data = (
-            await client.document_item_repository.get_all_picture_data(document.id)
+            await session.document_item_repository.get_all_picture_data(document.id)
         )
 
-    chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
+    chunks = await ensure_chunks_embedded(
+        session.config, chunks, session.store.embedder
+    )
 
     items: list[DocumentItem] | None = None
     if docling_document is not None:
@@ -194,29 +219,29 @@ async def _update_document_with_chunks(
             extract_items, document.id, docling_document, existing_picture_data
         )
 
-    async with client.store.write_transaction():
-        updated_doc = await client.document_repository.update(document)
+    async with session.store.write_transaction():
+        updated_doc = await session.document_repository.update(document)
 
         assert updated_doc.id is not None
         for order, chunk in enumerate(chunks):
             chunk.document_id = updated_doc.id
             chunk.order = order
 
-        await client.chunk_repository.replace_for_document(updated_doc.id, chunks)
+        await session.chunk_repository.replace_for_document(updated_doc.id, chunks)
 
         if items is not None:
-            await client.document_item_repository.replace_for_document(
+            await session.document_item_repository.replace_for_document(
                 updated_doc.id, items
             )
 
-    if client._config.storage.auto_vacuum:
-        client._schedule_vacuum()
+    if session.config.storage.auto_vacuum:
+        session.schedule_vacuum()
 
     return updated_doc
 
 
 async def create_document(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     content: str,
     uri: str | None = None,
     title: str | None = None,
@@ -227,9 +252,9 @@ async def create_document(
 
     Converts the content, chunks it, and generates embeddings.
     """
-    converter = get_converter(client._config)
+    converter = get_converter(session.config)
     docling_document = await converter.convert_text(content, format=format)
-    chunks = await client.chunk(docling_document)
+    chunks = await chunk_document(session, docling_document)
 
     document = Document(
         content="",
@@ -237,13 +262,15 @@ async def create_document(
         title=title,
         metadata=metadata or {},
     )
-    await _prepare_and_title(client, document, docling_document)
+    await _prepare_and_title(session, document, docling_document)
 
-    return await _store_document_with_chunks(client, document, chunks, docling_document)
+    return await _store_document_with_chunks(
+        session, document, chunks, docling_document
+    )
 
 
 async def import_document(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     docling_document: "DoclingDocument",
     chunks: list[Chunk],
     uri: str | None = None,
@@ -261,13 +288,15 @@ async def import_document(
         title=title,
         metadata=metadata or {},
     )
-    await _prepare_and_title(client, document, docling_document)
+    await _prepare_and_title(session, document, docling_document)
 
-    return await _store_document_with_chunks(client, document, chunks, docling_document)
+    return await _store_document_with_chunks(
+        session, document, chunks, docling_document
+    )
 
 
 async def _store_documents_with_chunks(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     prepared: list[tuple[Document, list[Chunk], "DoclingDocument"]],
 ) -> list[Document]:
     """Store many documents with their chunks in a single table version each.
@@ -276,9 +305,9 @@ async def _store_documents_with_chunks(
     and document_items tables once apiece. Restores all tables on any failure.
     """
     flat = await ensure_chunks_embedded(
-        client._config,
+        session.config,
         [chunk for _, chunks, _ in prepared for chunk in chunks],
-        client.embedder,
+        session.store.embedder,
     )
     embedded: list[list[Chunk]] = []
     position = 0
@@ -291,8 +320,8 @@ async def _store_documents_with_chunks(
 
     all_item_lists = await asyncio.to_thread(_extract_all_items)
 
-    async with client.store.write_transaction():
-        created = await client.document_repository.create(
+    async with session.store.write_transaction():
+        created = await session.document_repository.create(
             [doc for doc, _, _ in prepared]
         )
 
@@ -308,17 +337,17 @@ async def _store_documents_with_chunks(
                 item.document_id = doc.id
             all_items.extend(item_list)
 
-        await client.chunk_repository.create(all_chunks)
-        await client.document_item_repository.create_all(all_items)
+        await session.chunk_repository.create(all_chunks)
+        await session.document_item_repository.create_all(all_items)
 
-    if client._config.storage.auto_vacuum:
-        client._schedule_vacuum()
+    if session.config.storage.auto_vacuum:
+        session.schedule_vacuum()
 
     return created
 
 
 async def import_documents(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     imports: list[DocumentImport],
 ) -> list[Document]:
     """Batch-import pre-processed documents with their chunks.
@@ -338,14 +367,14 @@ async def import_documents(
             title=item.title,
             metadata=item.metadata or {},
         )
-        await _prepare_and_title(client, document, item.docling_document)
+        await _prepare_and_title(session, document, item.docling_document)
         prepared.append((document, item.chunks, item.docling_document))
 
-    return await _store_documents_with_chunks(client, prepared)
+    return await _store_documents_with_chunks(session, prepared)
 
 
 async def _refresh_doc_metadata(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     doc: Document,
     *,
     title: str | None,
@@ -367,12 +396,12 @@ async def _refresh_doc_metadata(
         updated = True
 
     if updated:
-        async with client.store._write_lock:
-            result = await client.document_repository.update_meta(doc)
+        async with session.store._write_lock:
+            result = await session.document_repository.update_meta(doc)
         # Reclaim the document_meta churn from rolling source_revision sweeps.
         # The vacuum is debounced, and document_meta is tiny, so this is cheap.
-        if client._config.storage.auto_vacuum:
-            client._schedule_vacuum()
+        if session.config.storage.auto_vacuum:
+            session.schedule_vacuum()
         return result
     return doc
 
@@ -398,7 +427,7 @@ async def _provider_metadata(
 
 
 async def _ingest_fetch_result(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     result: "FetchResult",
     *,
     title: str | None,
@@ -417,7 +446,7 @@ async def _ingest_fetch_result(
     fallback. Callers pass it when ``result.uri`` cannot yield the right
     extension, e.g. embedded attachments whose name lives in a URI fragment."""
 
-    converter = get_converter(client._config)
+    converter = get_converter(session.config)
     if filename is not None:
         file_extension = Path(filename).suffix.lower()
     else:
@@ -446,9 +475,11 @@ async def _ingest_fetch_result(
 
     try:
         with logfire.span("document.convert", uri=result.uri):
-            docling_document = await client.convert(target_path, source_uri=result.uri)
+            docling_document = await convert(
+                session.config, target_path, source_uri=result.uri
+            )
         with logfire.span("document.chunk", uri=result.uri) as chunk_span:
-            chunks = await client.chunk(docling_document)
+            chunks = await chunk_document(session, docling_document)
             chunk_span.set_attribute("chunks_created", len(chunks))
     finally:
         if cleanup_path is not None:
@@ -460,13 +491,13 @@ async def _ingest_fetch_result(
         existing_doc.metadata = final_metadata
         if title is not None:
             existing_doc.title = title
-        await _prepare_and_title(client, existing_doc, docling_document)
+        await _prepare_and_title(session, existing_doc, docling_document)
         with logfire.span("document.store", uri=result.uri, op="update") as store_span:
             updated = await _update_document_with_chunks(
-                client, existing_doc, chunks, docling_document
+                session, existing_doc, chunks, docling_document
             )
             store_span.set_attribute("document_id", updated.id)
-        await _reconcile_pdf_attachments(client, updated, result.body, depth=depth)
+        await _reconcile_pdf_attachments(session, updated, result.body, depth=depth)
         return updated
 
     document = Document(
@@ -475,13 +506,13 @@ async def _ingest_fetch_result(
         title=title,
         metadata=final_metadata,
     )
-    await _prepare_and_title(client, document, docling_document)
+    await _prepare_and_title(session, document, docling_document)
     with logfire.span("document.store", uri=result.uri, op="create") as store_span:
         created = await _store_document_with_chunks(
-            client, document, chunks, docling_document
+            session, document, chunks, docling_document
         )
         store_span.set_attribute("document_id", created.id)
-    await _reconcile_pdf_attachments(client, created, result.body, depth=depth)
+    await _reconcile_pdf_attachments(session, created, result.body, depth=depth)
     return created
 
 
@@ -545,7 +576,7 @@ def _extract_pdf_attachments(
 
 
 async def _reconcile_pdf_attachments(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     parent_doc: Document,
     parent_body: bytes,
     *,
@@ -559,7 +590,7 @@ async def _reconcile_pdf_attachments(
     path runs uniformly — child PDFs recurse into this helper one level deeper,
     bounded by ``MAX_ATTACHMENT_DEPTH``.
     """
-    if not client._config.processing.extract_pdf_attachments:
+    if not session.config.processing.extract_pdf_attachments:
         return
     if not parent_doc.uri:
         return
@@ -572,7 +603,7 @@ async def _reconcile_pdf_attachments(
     if new_attachments is None:
         return
 
-    existing = await client.list_documents(filter=parent_uri_filter(parent_doc.uri))
+    existing = await session.list_documents(filter=parent_uri_filter(parent_doc.uri))
     existing_by_uri: dict[str, Document] = {d.uri: d for d in existing if d.uri}
 
     for child_uri, (name, data, content_type, content_hash) in new_attachments.items():
@@ -594,7 +625,7 @@ async def _reconcile_pdf_attachments(
         )
         try:
             await _ingest_fetch_result(
-                client,
+                session,
                 child_fr,
                 title=None,
                 user_metadata={},
@@ -615,11 +646,11 @@ async def _reconcile_pdf_attachments(
 
     for child_uri, child in existing_by_uri.items():
         if child_uri not in new_attachments and child.id:
-            await client.delete_document(child.id)
+            await session.delete_document(child.id)
 
 
 async def create_document_from_source(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     source: str | Path,
     title: str | None = None,
     metadata: dict | None = None,
@@ -674,7 +705,7 @@ async def create_document_from_source(
             for child in walk_files(local_path):
                 if child.is_file() and filter.include_file(str(child)):
                     doc = await create_document_from_source(
-                        client,
+                        session,
                         child,
                         title=None,
                         metadata=metadata,
@@ -692,7 +723,7 @@ async def create_document_from_source(
 
         # Match the old _create_document_from_file behaviour: fail fast on
         # unsupported extension before reading any bytes.
-        converter = get_converter(client._config)
+        converter = get_converter(session.config)
         if local_path.suffix.lower() not in converter.supported_extensions:
             raise UnsupportedSourceError(
                 f"Unsupported file extension: {local_path.suffix}"
@@ -733,7 +764,7 @@ async def create_document_from_source(
         else:
             stored_uri = source_str
 
-        existing_doc = await client.get_document_by_uri(stored_uri)
+        existing_doc = await session.get_document_by_uri(stored_uri)
 
         # Cheap revision-based short-circuit: only worth a HEAD when we have a
         # stored revision to compare against. All sources persist their native
@@ -748,7 +779,7 @@ async def create_document_from_source(
             current_revision = await fetcher.head(source_str)
             if current_revision == stored_revision:
                 return await _refresh_doc_metadata(
-                    client,
+                    session,
                     existing_doc,
                     title=title,
                     user_metadata=metadata,
@@ -781,7 +812,7 @@ async def create_document_from_source(
             if result.revision is not None:
                 source_meta["source_revision"] = result.revision
             return await _refresh_doc_metadata(
-                client,
+                session,
                 existing_doc,
                 title=title,
                 user_metadata=user_metadata,
@@ -789,7 +820,7 @@ async def create_document_from_source(
             )
 
         return await _ingest_fetch_result(
-            client,
+            session,
             result,
             title=title,
             user_metadata=user_metadata,
@@ -802,7 +833,7 @@ async def create_document_from_source(
 
 
 async def update_document(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     document_id: str,
     content: str | None = None,
     metadata: dict | None = None,
@@ -830,7 +861,7 @@ async def update_document(
 
     # Caller-supplied chunks without a docling document replace neither blob,
     # and the row is written back whole, so they have to make the round trip.
-    existing_doc = await client.document_repository.get_by_id(
+    existing_doc = await session.document_repository.get_by_id(
         document_id, include_blobs=chunks is not None and docling_document is None
     )
     if existing_doc is None:
@@ -844,10 +875,10 @@ async def update_document(
         existing_doc.uri = uri
 
     if content is None and chunks is None and docling_document is None:
-        async with client.store._write_lock:
-            updated = await client.document_repository.update_meta(existing_doc)
-        if client._config.storage.auto_vacuum:
-            client._schedule_vacuum()
+        async with session.store._write_lock:
+            updated = await session.document_repository.update_meta(existing_doc)
+        if session.config.storage.auto_vacuum:
+            session.schedule_vacuum()
         return updated
 
     if chunks is not None:
@@ -857,26 +888,26 @@ async def update_document(
             existing_doc.content = content
 
         return await _update_document_with_chunks(
-            client, existing_doc, chunks, docling_document
+            session, existing_doc, chunks, docling_document
         )
 
     if docling_document is not None:
         await _prepare_document_from_docling(existing_doc, docling_document)
 
-        new_chunks = await client.chunk(docling_document)
+        new_chunks = await chunk_document(session, docling_document)
         return await _update_document_with_chunks(
-            client, existing_doc, new_chunks, docling_document
+            session, existing_doc, new_chunks, docling_document
         )
 
     assert content is not None
     existing_doc.content = content
-    converter = get_converter(client._config)
+    converter = get_converter(session.config)
     converted_docling = await converter.convert_text(existing_doc.content, format="md")
     await _prepare_document_from_docling(existing_doc, converted_docling)
 
-    new_chunks = await client.chunk(converted_docling)
+    new_chunks = await chunk_document(session, converted_docling)
     return await _update_document_with_chunks(
-        client, existing_doc, new_chunks, converted_docling
+        session, existing_doc, new_chunks, converted_docling
     )
 
 
diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py
index 3f5f6b3f..be8dffe5 100644
--- a/haiku_rag_slim/haiku/rag/client/rebuild.py
+++ b/haiku_rag_slim/haiku/rag/client/rebuild.py
@@ -10,8 +10,11 @@ from lancedb.pydantic import LanceModel
 
 from haiku.rag.client.documents import (
     check_source_accessible,
+    chunk_document,
     create_document_from_source,
 )
+from haiku.rag.client.session import SingleDatabaseSession
+from haiku.rag.client.titles import generate_title
 from haiku.rag.converters import get_converter
 from haiku.rag.store.compression import compress_docling_split
 from haiku.rag.store.models.chunk import Chunk
@@ -23,7 +26,7 @@ from haiku.rag.store.schema import ChunkRecordBase
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
 
-    from haiku.rag.client import HaikuRAG, RebuildMode
+    from haiku.rag.client import RebuildMode
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +72,7 @@ class _StagingMarkerRecord(LanceModel):
 
 
 async def rebuild_database(
-    client: "HaikuRAG", mode: "RebuildMode"
+    session: SingleDatabaseSession, mode: "RebuildMode"
 ) -> AsyncGenerator[str, None]:
     """Rebuild the database with the specified mode.
 
@@ -82,17 +85,17 @@ async def rebuild_database(
     """
     from haiku.rag.client import RebuildMode
 
-    async with client.store._rebuild_lock:
+    async with session.store._rebuild_lock:
         if mode == RebuildMode.SET_EMBEDDER:
-            await _set_embedder(client)
+            await _set_embedder(session)
             return
 
-        async for doc_id in _rebuild_locked(client, mode):
+        async for doc_id in _rebuild_locked(session, mode):
             yield doc_id
 
 
 async def _rebuild_locked(
-    client: "HaikuRAG", mode: "RebuildMode"
+    session: SingleDatabaseSession, mode: "RebuildMode"
 ) -> AsyncGenerator[str, None]:
     from haiku.rag.client import RebuildMode
 
@@ -100,7 +103,7 @@ async def _rebuild_locked(
     # interrupted rebuild. Returns True only when phase 1 was already
     # complete and the current mode is EMBED_ONLY, in which case we resume
     # phase 2 from the existing staging table instead of recopying.
-    resume_from_staging = await _resolve_rebuild_recovery(client, mode)
+    resume_from_staging = await _resolve_rebuild_recovery(session, mode)
 
     # Wait for any already-scheduled background vacuum before the destructive
     # table operations at the top of RECHUNK / FULL. Rebuild drops and
@@ -109,52 +112,52 @@ async def _rebuild_locked(
     # lance. Note: FULL calls create_document_from_source inside its loop,
     # which may schedule *new* background vacuums — those run after the
     # destructive phase and are fine.
-    await client._await_vacuum_tasks()
+    await session.drain_vacuum()
 
-    settings_repo = SettingsRepository(client.store)
+    settings_repo = SettingsRepository(session.store)
     await settings_repo.save_current_settings()
 
     # Light listing — id/uri/title/metadata only. Each rebuild function
     # fetches full content (including the multi-MB docling_pages blob) one
     # document at a time so a 1000-doc database doesn't pull ~15 GB of
     # blobs into memory before the loop starts.
-    documents = await client.list_documents(include_content=False)
+    documents = await session.list_documents(include_content=False)
 
     if mode == RebuildMode.TITLE_ONLY:
-        async for doc_id in _rebuild_title_only(client, documents):
+        async for doc_id in _rebuild_title_only(session, documents):
             yield doc_id
     elif mode == RebuildMode.EMBED_ONLY:
         async for doc_id in _rebuild_embed_only(
-            client, documents, resume_from_staging=resume_from_staging
+            session, documents, resume_from_staging=resume_from_staging
         ):
             yield doc_id
     elif mode == RebuildMode.RECHUNK:
-        await client.chunk_repository.delete_all()
-        await client.store.recreate_embeddings_table()
-        async for doc_id in _rebuild_rechunk(client, documents):
+        await session.chunk_repository.delete_all()
+        await session.store.recreate_embeddings_table()
+        async for doc_id in _rebuild_rechunk(session, documents):
             yield doc_id
     elif mode == RebuildMode.DESCRIPTIONS:
-        await client.chunk_repository.delete_all()
-        await client.store.recreate_embeddings_table()
-        async for doc_id in _rebuild_descriptions(client, documents):
+        await session.chunk_repository.delete_all()
+        await session.store.recreate_embeddings_table()
+        async for doc_id in _rebuild_descriptions(session, documents):
             yield doc_id
     else:  # FULL
-        await client.chunk_repository.delete_all()
-        await client.store.recreate_embeddings_table()
-        async for doc_id in _rebuild_full(client, documents):
+        await session.chunk_repository.delete_all()
+        await session.store.recreate_embeddings_table()
+        async for doc_id in _rebuild_full(session, documents):
             yield doc_id
 
     # Final maintenance if auto_vacuum enabled. Swallowing only so that a
     # failed post-rebuild optimize doesn't mask a successful rebuild — but
     # log it so the failure is visible in the output.
-    if client._config.storage.auto_vacuum:
+    if session.config.storage.auto_vacuum:
         try:
-            await client.store.vacuum()
+            await session.store.vacuum()
         except Exception:
             logger.warning("Post-rebuild vacuum failed", exc_info=True)
 
 
-async def _set_embedder(client: "HaikuRAG") -> None:
+async def _set_embedder(session: SingleDatabaseSession) -> None:
     """Adopt the current embedder identity without re-embedding.
 
     Only valid when the vector dimension is unchanged — the stored vectors stay
@@ -163,10 +166,10 @@ async def _set_embedder(client: "HaikuRAG") -> None:
     """
     from haiku.rag.store.exceptions import ConfigMismatchError
 
-    settings_repo = SettingsRepository(client.store)
+    settings_repo = SettingsRepository(session.store)
     stored = await settings_repo.get_current_settings()
     stored_dim = stored.get("embeddings", {}).get("model", {}).get("vector_dim")
-    current_dim = client._config.embeddings.model.vector_dim
+    current_dim = session.config.embeddings.model.vector_dim
 
     if stored_dim is not None and current_dim != stored_dim:
         raise ConfigMismatchError(
@@ -178,7 +181,9 @@ async def _set_embedder(client: "HaikuRAG") -> None:
 
 
 async def _hydrate(
-    client: "HaikuRAG", light_docs: list[Document], include_blobs: bool = True
+    session: SingleDatabaseSession,
+    light_docs: list[Document],
+    include_blobs: bool = True,
 ) -> AsyncGenerator[Document, None]:
     """Yield fully-loaded documents one at a time from a light listing.
 
@@ -189,7 +194,7 @@ async def _hydrate(
     """
     for light_doc in light_docs:
         assert light_doc.id is not None
-        doc = await client.document_repository.get_by_id(
+        doc = await session.document_repository.get_by_id(
             light_doc.id, include_blobs=include_blobs
         )
         if doc is None:
@@ -199,22 +204,22 @@ async def _hydrate(
 
 
 async def _rebuild_title_only(
-    client: "HaikuRAG", documents: list[Document]
+    session: SingleDatabaseSession, documents: list[Document]
 ) -> AsyncGenerator[str, None]:
     """Generate titles for documents that don't have one.
 
     A title comes from the content or the docling structure, never the page
     rasters, so those are left out of the per-document load.
     """
-    repo = client.document_repository
+    repo = session.document_repository
     untitled = [d for d in documents if d.title is None]
-    async for doc in _hydrate(client, untitled, include_blobs=False):
+    async for doc in _hydrate(session, untitled, include_blobs=False):
         assert doc.id is not None
         structure = await repo.get_docling_data(doc.id)
         if structure is not None:
             doc.docling_document = structure.docling_document
         try:
-            title = await client.generate_title(doc)
+            title = await generate_title(session.config, doc)
         except Exception:
             logger.warning(
                 "Failed to generate title for document %s", doc.id, exc_info=True
@@ -226,7 +231,9 @@ async def _rebuild_title_only(
             yield doc.id
 
 
-async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool:
+async def _resolve_rebuild_recovery(
+    session: SingleDatabaseSession, mode: "RebuildMode"
+) -> bool:
     """Resolve any partially-completed rebuild state from a previous crash.
 
     Returns ``True`` if ``_rebuild_embed_only`` should resume from the
@@ -243,7 +250,7 @@ async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") ->
     """
     from haiku.rag.client import RebuildMode
 
-    db = client.store.db
+    db = session.store.db
     tables = (await db.list_tables()).tables
     has_staging = _STAGING_TABLE_NAME in tables
     has_marker = _STAGING_MARKER_TABLE_NAME in tables
@@ -286,7 +293,7 @@ async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") ->
     return False
 
 
-async def _populate_staging_table(client: "HaikuRAG") -> None:
+async def _populate_staging_table(session: SingleDatabaseSession) -> None:
     """Stream the non-vector columns of the chunks table into staging.
 
     Uses ``to_batches`` for a single streaming read (no offset/limit
@@ -297,7 +304,7 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
     Requires ``_resolve_rebuild_recovery`` to have cleared any leftover
     staging table first: ``create_table`` raises if the name is already taken.
     """
-    db = client.store.db
+    db = session.store.db
     tables = (await db.list_tables()).tables
 
     staging = await db.create_table(_STAGING_TABLE_NAME, schema=_StagingChunkRecord)
@@ -305,7 +312,7 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
         return
 
     stream = (
-        await client.store.chunks_table.query()
+        await session.store.chunks_table.query()
         .select(["id", "document_id", "content", "metadata", "order"])
         .to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE)
     )
@@ -324,13 +331,13 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
         await staging.add(records)
 
 
-async def _mark_phase1_complete(client: "HaikuRAG") -> None:
+async def _mark_phase1_complete(session: SingleDatabaseSession) -> None:
     """Create the marker table that designates staging as authoritative.
 
     Called after ``_populate_staging_table`` finishes. On crash recovery the
     marker's presence flips ``_rebuild_embed_only`` into resume mode.
     """
-    db = client.store.db
+    db = session.store.db
     if _STAGING_MARKER_TABLE_NAME in (await db.list_tables()).tables:
         return
     marker = await db.create_table(
@@ -339,7 +346,7 @@ async def _mark_phase1_complete(client: "HaikuRAG") -> None:
     await marker.add([_StagingMarkerRecord(id="phase1_complete")])
 
 
-async def _drop_staging_tables(client: "HaikuRAG") -> None:
+async def _drop_staging_tables(session: SingleDatabaseSession) -> None:
     """Drop the marker first, then the staging table.
 
     Ordering matters: if a crash interrupts cleanup between the two drops,
@@ -347,7 +354,7 @@ async def _drop_staging_tables(client: "HaikuRAG") -> None:
     partial phase 1 → drops staging harmlessly. The reverse order would
     leak a marker pointing at nothing.
     """
-    db = client.store.db
+    db = session.store.db
     tables = (await db.list_tables()).tables
     if _STAGING_MARKER_TABLE_NAME in tables:
         await db.drop_table(_STAGING_MARKER_TABLE_NAME)
@@ -384,7 +391,7 @@ async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chu
 
 
 async def _rebuild_embed_only(
-    client: "HaikuRAG",
+    session: SingleDatabaseSession,
     documents: list[Document],
     *,
     resume_from_staging: bool = False,
@@ -414,19 +421,19 @@ async def _rebuild_embed_only(
     """
     from haiku.rag.embeddings import contextualize, embed_chunks
 
-    db = client.store.db
-    embedder = client.chunk_repository.embedder
+    db = session.store.db
+    embedder = session.chunk_repository.embedder
 
     if not resume_from_staging:
         # Phase 1: copy chunks into staging, then mark it complete. After the
         # marker exists, a crash will resume phase 2 from staging.
-        await _populate_staging_table(client)
-        await _mark_phase1_complete(client)
+        await _populate_staging_table(session)
+        await _mark_phase1_complete(session)
 
     # Recreate the chunks table fresh (idempotent; handles vector-dim
     # changes and discards any partial new chunks from a prior crashed
     # phase 2).
-    await client.store.recreate_embeddings_table()
+    await session.store.recreate_embeddings_table()
 
     staging_table = await db.open_table(_STAGING_TABLE_NAME)
 
@@ -443,7 +450,7 @@ async def _rebuild_embed_only(
         # chunks route through embed_image rather than being text-embedded.
         # Bytes live in document_items (embed-only never touches that table).
         if embedder.supports_images:
-            picture_data = await client.document_item_repository.get_all_picture_data(
+            picture_data = await session.document_item_repository.get_all_picture_data(
                 doc.id
             )
             for chunk in chunks:
@@ -462,7 +469,7 @@ async def _rebuild_embed_only(
                     )
 
         content_fts_list = contextualize(chunks)
-        embedded_chunks = await embed_chunks(chunks, embedder, client._config)
+        embedded_chunks = await embed_chunks(chunks, embedder, session.config)
 
         for chunk, content_fts, embedded in zip(
             chunks, content_fts_list, embedded_chunks
@@ -471,7 +478,7 @@ async def _rebuild_embed_only(
             assert chunk.document_id is not None
             assert embedded.embedding is not None
             pending_records.append(
-                client.store.ChunkRecord(
+                session.store.ChunkRecord(
                     id=chunk.id,
                     document_id=chunk.document_id,
                     content=chunk.content,
@@ -491,16 +498,16 @@ async def _rebuild_embed_only(
         yield doc.id
 
         if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
-            await client.store.chunks_table.add(pending_records)
+            await session.store.chunks_table.add(pending_records)
             pending_records = []
 
     if pending_records:
-        await client.store.chunks_table.add(pending_records)
+        await session.store.chunks_table.add(pending_records)
 
     # Phase 2 finished. Drop the recovery state — marker first so a crash
     # between the two drops leaves only staging behind, which the next
     # rebuild discards harmlessly.
-    await _drop_staging_tables(client)
+    await _drop_staging_tables(session)
 
     for doc in documents:
         if doc.id and doc.id not in yielded_docs:
@@ -508,7 +515,7 @@ async def _rebuild_embed_only(
 
 
 async def _flush_rebuild_batch(
-    client: "HaikuRAG", documents: list[Document], chunks: list[Chunk]
+    session: SingleDatabaseSession, documents: list[Document], chunks: list[Chunk]
 ) -> None:
     """Batch write documents and chunks during rebuild.
 
@@ -552,12 +559,12 @@ async def _flush_rebuild_batch(
         )
 
     await (
-        client.store.documents_table.merge_insert("id")
+        session.store.documents_table.merge_insert("id")
         .when_matched_update_all()
         .execute(doc_records)
     )
     await (
-        client.store.document_meta_table.merge_insert("id")
+        session.store.document_meta_table.merge_insert("id")
         .when_matched_update_all()
         .when_not_matched_insert_all()
         .execute(meta_records)
@@ -565,7 +572,7 @@ async def _flush_rebuild_batch(
 
     # Batch create all chunks (single LanceDB version)
     if chunks:
-        await client.chunk_repository.create(chunks)
+        await session.chunk_repository.create(chunks)
 
     # Repopulate document items from stored docling data. The stored docling
     # blob has had its picture URIs stripped (compress_docling_split), so
@@ -576,28 +583,28 @@ async def _flush_rebuild_batch(
         docling_doc = doc.get_docling_document()
         if docling_doc is not None:
             existing_picture_data = (
-                await client.document_item_repository.get_all_picture_data(doc.id)
+                await session.document_item_repository.get_all_picture_data(doc.id)
             )
-            await client.document_item_repository.delete_by_document_id(doc.id)
+            await session.document_item_repository.delete_by_document_id(doc.id)
             items = extract_items(
                 doc.id,
                 docling_doc,
                 existing_picture_data=existing_picture_data,
             )
-            await client.document_item_repository.create_items(doc.id, items)
+            await session.document_item_repository.create_items(doc.id, items)
 
 
 async def _rebuild_rechunk(
-    client: "HaikuRAG", documents: list[Document]
+    session: SingleDatabaseSession, documents: list[Document]
 ) -> AsyncGenerator[str, None]:
     """Re-chunk and re-embed each document from its stored docling blob."""
     from haiku.rag.embeddings import embed_chunks
 
     pending_chunks: list[Chunk] = []
     pending_docs: list[Document] = []
-    embedder = client.embedder
+    embedder = session.store.embedder
 
-    async for doc in _hydrate(client, documents):
+    async for doc in _hydrate(session, documents):
         assert doc.id is not None
         docling_document = doc.get_docling_document()
         if docling_document is None:
@@ -609,16 +616,17 @@ async def _rebuild_rechunk(
         # Stored blob has stripped picture URIs; pass the snapshot so
         # build_picture_chunks (inside chunk()) can recover the bytes.
         existing_picture_data = (
-            await client.document_item_repository.get_all_picture_data(doc.id)
+            await session.document_item_repository.get_all_picture_data(doc.id)
             if embedder.supports_images
             else None
         )
-        chunks = await client.chunk(
+        chunks = await chunk_document(
+            session,
             docling_document,
             existing_picture_data=existing_picture_data,
             document_id=doc.id,
         )
-        embedded_chunks = await embed_chunks(chunks, embedder, client._config)
+        embedded_chunks = await embed_chunks(chunks, embedder, session.config)
 
         for order, chunk in enumerate(embedded_chunks):
             chunk.document_id = doc.id
@@ -634,12 +642,12 @@ async def _rebuild_rechunk(
         yield doc.id
 
         if len(pending_docs) >= _REBUILD_BATCH_SIZE:
-            await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+            await _flush_rebuild_batch(session, pending_docs, pending_chunks)
             pending_chunks = []
             pending_docs = []
 
     if pending_docs:
-        await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+        await _flush_rebuild_batch(session, pending_docs, pending_chunks)
 
 
 def _apply_descriptions_sync(
@@ -668,7 +676,9 @@ def _apply_descriptions_sync(
     return len(descriptions)
 
 
-async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
+async def _patch_picture_descriptions(
+    session: SingleDatabaseSession, doc: Document
+) -> int:
     """Run the VLM against pictures lacking a description, patch the docling
     blob in-place. Returns the number of newly described pictures.
     Pictures that already carry ``meta.description.text`` are skipped, so the
@@ -692,7 +702,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
     if not needs_description:
         return 0
 
-    bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk(
+    bytes_by_ref = await session.document_item_repository.get_pictures_for_chunk(
         doc.id, needs_description
     )
     if not bytes_by_ref:
@@ -705,7 +715,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
         )
         return 0
 
-    descriptions = await describe_pictures(bytes_by_ref, config=client._config)
+    descriptions = await describe_pictures(bytes_by_ref, config=session.config)
 
     if not descriptions:
         return 0
@@ -716,7 +726,7 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
 
 
 async def _rebuild_descriptions(
-    client: "HaikuRAG", documents: list[Document]
+    session: SingleDatabaseSession, documents: list[Document]
 ) -> AsyncGenerator[str, None]:
     """Run the VLM over already-stored picture bytes, patch descriptions into
     the docling blob, then re-chunk + re-embed.
@@ -727,7 +737,7 @@ async def _rebuild_descriptions(
     """
     from haiku.rag.embeddings import embed_chunks
 
-    if client._config.processing.pictures != "description":
+    if session.config.processing.pictures != "description":
         raise ValueError(
             "rebuild --descriptions requires processing.pictures = 'description' "
             "in your config."
@@ -735,10 +745,10 @@ async def _rebuild_descriptions(
 
     pending_chunks: list[Chunk] = []
     pending_docs: list[Document] = []
-    embedder = client.embedder
+    embedder = session.store.embedder
 
     described_total = 0
-    async for doc in _hydrate(client, documents):
+    async for doc in _hydrate(session, documents):
         assert doc.id is not None
         docling_document = doc.get_docling_document()
         if docling_document is None:
@@ -747,23 +757,24 @@ async def _rebuild_descriptions(
                 "rebuild --descriptions requires it. Run a full rebuild instead."
             )
 
-        n = await _patch_picture_descriptions(client, doc)
+        n = await _patch_picture_descriptions(session, doc)
         described_total += n
         # Use the (possibly patched) docling document for chunking.
         docling_document = doc.get_docling_document()
         assert docling_document is not None
 
         existing_picture_data = (
-            await client.document_item_repository.get_all_picture_data(doc.id)
+            await session.document_item_repository.get_all_picture_data(doc.id)
             if embedder.supports_images
             else None
         )
-        chunks = await client.chunk(
+        chunks = await chunk_document(
+            session,
             docling_document,
             existing_picture_data=existing_picture_data,
             document_id=doc.id,
         )
-        embedded_chunks = await embed_chunks(chunks, embedder, client._config)
+        embedded_chunks = await embed_chunks(chunks, embedder, session.config)
 
         for order, chunk in enumerate(embedded_chunks):
             chunk.document_id = doc.id
@@ -774,12 +785,12 @@ async def _rebuild_descriptions(
         yield doc.id
 
         if len(pending_docs) >= _REBUILD_BATCH_SIZE:
-            await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+            await _flush_rebuild_batch(session, pending_docs, pending_chunks)
             pending_chunks = []
             pending_docs = []
 
     if pending_docs:
-        await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+        await _flush_rebuild_batch(session, pending_docs, pending_chunks)
 
     logger.info(
         "rebuild --descriptions: %d new picture descriptions added across %d documents",
@@ -789,15 +800,15 @@ async def _rebuild_descriptions(
 
 
 async def _rebuild_full(
-    client: "HaikuRAG", documents: list[Document]
+    session: SingleDatabaseSession, documents: list[Document]
 ) -> AsyncGenerator[str, None]:
     """Full rebuild: re-convert from source, re-chunk, re-embed."""
     from haiku.rag.embeddings import embed_chunks
 
     pending_chunks: list[Chunk] = []
     pending_docs: list[Document] = []
-    converter = get_converter(client._config)
-    embedder = client.embedder
+    converter = get_converter(session.config)
+    embedder = session.store.embedder
 
     for light_doc in documents:
         assert light_doc.id is not None
@@ -805,10 +816,10 @@ async def _rebuild_full(
         # Try to rebuild from source if available — uses the light listing
         # directly, no need to load the stored content/blobs first.
         if light_doc.uri and check_source_accessible(light_doc.uri):
-            # The refresh writes through the client, not the batch buffer, so
+            # The refresh writes through the database, not the batch buffer, so
             # anything pending has to land first.
             if pending_docs:
-                await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+                await _flush_rebuild_batch(session, pending_docs, pending_chunks)
                 pending_chunks = []
                 pending_docs = []
 
@@ -817,7 +828,7 @@ async def _rebuild_full(
                 # point of a FULL rebuild is to re-convert them anyway. Updates
                 # in place, so a failure here cannot cost the document.
                 refreshed = await create_document_from_source(
-                    client,
+                    session,
                     source=light_doc.uri,
                     metadata=light_doc.metadata or {},
                     force=True,
@@ -840,7 +851,7 @@ async def _rebuild_full(
 
         # Fallback: rebuild from stored content. Now we need the full
         # record (content + docling_pages for the round-trip write).
-        doc = await client.document_repository.get_by_id(
+        doc = await session.document_repository.get_by_id(
             light_doc.id, include_blobs=True
         )
         if doc is None:
@@ -848,8 +859,8 @@ async def _rebuild_full(
         assert doc.id is not None
 
         docling_document = await converter.convert_text(doc.content, format="md")
-        chunks = await client.chunk(docling_document)
-        embedded_chunks = await embed_chunks(chunks, embedder, client._config)
+        chunks = await chunk_document(session, docling_document)
+        embedded_chunks = await embed_chunks(chunks, embedder, session.config)
 
         doc.set_docling(docling_document)
 
@@ -862,9 +873,9 @@ async def _rebuild_full(
         yield doc.id
 
         if len(pending_docs) >= _REBUILD_BATCH_SIZE:
-            await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+            await _flush_rebuild_batch(session, pending_docs, pending_chunks)
             pending_chunks = []
             pending_docs = []
 
     if pending_docs:
-        await _flush_rebuild_batch(client, pending_docs, pending_chunks)
+        await _flush_rebuild_batch(session, pending_docs, pending_chunks)
diff --git a/tests/conftest.py b/tests/conftest.py
index fb5229c1..15f4654c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -30,6 +30,9 @@ from .services import reachable  # noqa: E402
 if TYPE_CHECKING:
     from vcr import VCR
 
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.client.session import SingleDatabaseSession
+
 setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
 logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
 
@@ -216,3 +219,14 @@ def docling_serve_url() -> str:
     if not reachable("localhost", 5001):
         pytest.skip(f"docling-serve not reachable on localhost:5001 — {_COMPOSE_HINT}")
     return "http://localhost:5001"
+
+
+def writing(client: "HaikuRAG") -> "SingleDatabaseSession":
+    """The database a write implementation works on, from a client holding one.
+
+    Write implementations take a session rather than a client, so a set can
+    never reach them. Tests that call one directly go through here."""
+    from haiku.rag.client.session import SingleDatabaseSession
+
+    assert isinstance(client._session, SingleDatabaseSession)
+    return client._session
diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py
index 0f1412d4..87a00534 100644
--- a/tests/store/test_document_items.py
+++ b/tests/store/test_document_items.py
@@ -12,6 +12,7 @@ from haiku.rag.store.models.document_item import (
     extract_items,
 )
 from haiku.rag.store.repositories.document_item import DocumentItemRepository
+from tests.conftest import writing
 
 
 def _make_docling_doc():
@@ -416,7 +417,9 @@ class TestDocumentItemPopulation:
 
             # Use _store_document_with_chunks directly with empty chunks
             # to avoid needing embeddings
-            created = await _store_document_with_chunks(rag, document, [], docling_doc)
+            created = await _store_document_with_chunks(
+                writing(rag), document, [], docling_doc
+            )
             assert created.id is not None
 
             count = await rag.document_item_repository.get_item_count(created.id)
@@ -446,7 +449,9 @@ class TestDocumentItemPopulation:
                 uri="test://doc",
             )
             document.set_docling(docling_doc)
-            created = await _store_document_with_chunks(rag, document, [], docling_doc)
+            created = await _store_document_with_chunks(
+                writing(rag), document, [], docling_doc
+            )
             assert created.id is not None
             assert await rag.document_item_repository.get_item_count(created.id) == 6
 
@@ -455,7 +460,7 @@ class TestDocumentItemPopulation:
             new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
             created.set_docling(new_doc)
 
-            await _update_document_with_chunks(rag, created, [], new_doc)
+            await _update_document_with_chunks(writing(rag), created, [], new_doc)
             assert await rag.document_item_repository.get_item_count(created.id) == 1
 
     async def test_delete_document_cascades_items(self, temp_db_path):
@@ -470,7 +475,9 @@ class TestDocumentItemPopulation:
                 uri="test://doc",
             )
             document.set_docling(docling_doc)
-            created = await _store_document_with_chunks(rag, document, [], docling_doc)
+            created = await _store_document_with_chunks(
+                writing(rag), document, [], docling_doc
+            )
             assert created.id is not None
             assert await rag.document_item_repository.get_item_count(created.id) == 6
 
@@ -846,7 +853,9 @@ class TestPictureDataPreservedThroughRoundTrip:
         async with HaikuRAG(temp_db_path, create=True) as rag:
             document = Document(content="Hello world", uri="test://doc")
             document.set_docling(docling_doc)
-            created = await _store_document_with_chunks(rag, document, [], docling_doc)
+            created = await _store_document_with_chunks(
+                writing(rag), document, [], docling_doc
+            )
             assert created.id is not None
 
             original = await rag.document_item_repository.get_all_picture_data(
@@ -859,7 +868,7 @@ class TestPictureDataPreservedThroughRoundTrip:
             assert from_blob is not None
             assert all(p.image is None for p in from_blob.pictures)
 
-            await _update_document_with_chunks(rag, created, [], from_blob)
+            await _update_document_with_chunks(writing(rag), created, [], from_blob)
 
             after = await rag.document_item_repository.get_all_picture_data(created.id)
             assert after.get("#/pictures/0") == original.get("#/pictures/0")
diff --git a/tests/test_context.py b/tests/test_context.py
index 48130fdd..bd907d8b 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -12,6 +12,7 @@ from haiku.rag.context import (
 )
 from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.document_item import DocumentItem
+from tests.conftest import writing
 
 
 def _item(
@@ -352,7 +353,7 @@ class TestExpandWithItems:
 
         async with HaikuRAG(temp_db_path, create=True) as rag:
             doc = await _store_document_with_chunks(
-                rag,
+                writing(rag),
                 Document(content="test"),
                 [],
                 __import__(
diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py
index 2d68a481..314f207c 100644
--- a/tests/test_context_enhancement.py
+++ b/tests/test_context_enhancement.py
@@ -7,6 +7,7 @@ from haiku.rag.client.documents import _store_document_with_chunks
 from haiku.rag.client.processing import ensure_chunks_embedded
 from haiku.rag.config.models import AppConfig
 from haiku.rag.store.models import SearchResult
+from tests.conftest import writing
 
 
 async def create_document_with_docling(
@@ -318,7 +319,9 @@ async def test_expand_context_single_item_document(temp_db_path):
     async with HaikuRAG(temp_db_path, create=True) as client:
         document = Document(content="Simple test content")
         document.set_docling(docling_doc)
-        doc = await _store_document_with_chunks(client, document, [], docling_doc)
+        doc = await _store_document_with_chunks(
+            writing(client), document, [], docling_doc
+        )
         assert doc.id is not None
 
         # Create a search result with a doc_item_ref pointing to the item
diff --git a/tests/test_pdf_attachments.py b/tests/test_pdf_attachments.py
index f283aba3..f5f0145c 100644
--- a/tests/test_pdf_attachments.py
+++ b/tests/test_pdf_attachments.py
@@ -12,6 +12,7 @@ from haiku.rag.client.documents import (
 )
 from haiku.rag.sources import FetchResult
 from haiku.rag.store.models.document import Document
+from tests.conftest import writing
 
 
 def build_pdf(attachments: list[tuple[str, bytes]]) -> bytes:
@@ -27,7 +28,7 @@ def build_pdf(attachments: list[tuple[str, bytes]]) -> bytes:
 
 
 async def fake_ingest_fetch_result(
-    client,
+    session,
     result: FetchResult,
     *,
     title,
@@ -55,9 +56,9 @@ async def fake_ingest_fetch_result(
         existing_doc.metadata = final_metadata
         if title is not None:
             existing_doc.title = title
-        doc = await client.document_repository.update(existing_doc)
+        doc = await session.document_repository.update(existing_doc)
     else:
-        doc = await client.document_repository.create(
+        doc = await session.document_repository.create(
             Document(
                 content="",
                 uri=stored_uri,
@@ -65,7 +66,7 @@ async def fake_ingest_fetch_result(
                 metadata=final_metadata,
             )
         )
-    await _reconcile_pdf_attachments(client, doc, result.body, depth=depth)
+    await _reconcile_pdf_attachments(session, doc, result.body, depth=depth)
     return doc
 
 
@@ -101,7 +102,7 @@ async def test_first_ingest_creates_one_doc_per_attachment(temp_db_path, monkeyp
         parent_uri = "file:///fixtures/parent.pdf"
         parent = await _make_parent(client, parent_uri, pdf_bytes)
 
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
 
         children = await client.list_documents(filter=parent_uri_filter(parent_uri))
         assert len(children) == 2
@@ -129,7 +130,7 @@ async def test_attachment_with_spaces_in_name_is_percent_encoded(
     async with HaikuRAG(temp_db_path, create=True) as client:
         parent_uri = "file:///fixtures/parent.pdf"
         parent = await _make_parent(client, parent_uri, pdf_bytes)
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
 
         children = await client.list_documents(filter=parent_uri_filter(parent_uri))
         assert len(children) == 1
@@ -145,13 +146,13 @@ async def test_reingest_removes_dropped_attachment(temp_db_path, monkeypatch):
         parent_uri = "file:///fixtures/parent.pdf"
         first = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
         parent = await _make_parent(client, parent_uri, first)
-        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
         assert (
             len(await client.list_documents(filter=parent_uri_filter(parent_uri))) == 2
         )
 
         second = build_pdf([("a.txt", b"A")])
-        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
         remaining = await client.list_documents(filter=parent_uri_filter(parent_uri))
         assert len(remaining) == 1
         assert remaining[0].uri == f"{parent_uri}#attachment=a.txt"
@@ -166,11 +167,11 @@ async def test_reingest_updates_changed_attachment_in_place(temp_db_path, monkey
         parent_uri = "file:///fixtures/parent.pdf"
         first = build_pdf([("a.txt", b"original")])
         parent = await _make_parent(client, parent_uri, first)
-        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
         before = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
 
         second = build_pdf([("a.txt", b"different")])
-        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
         after = (await client.list_documents(filter=parent_uri_filter(parent_uri)))[0]
 
         assert after.id == before.id
@@ -186,10 +187,10 @@ async def test_reingest_adds_new_attachment(temp_db_path, monkeypatch):
         parent_uri = "file:///fixtures/parent.pdf"
         first = build_pdf([("a.txt", b"A")])
         parent = await _make_parent(client, parent_uri, first)
-        await _reconcile_pdf_attachments(client, parent, first, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, first, depth=0)
 
         second = build_pdf([("a.txt", b"A"), ("c.txt", b"C")])
-        await _reconcile_pdf_attachments(client, parent, second, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, second, depth=0)
         children = await client.list_documents(filter=parent_uri_filter(parent_uri))
         assert len(children) == 2
         names = {c.uri for c in children}
@@ -213,7 +214,7 @@ async def test_nested_pdf_attachments_recurse_up_to_cap(temp_db_path, monkeypatc
     async with HaikuRAG(temp_db_path, create=True) as client:
         root_uri = "file:///fixtures/root.pdf"
         parent = await _make_parent(client, root_uri, root)
-        await _reconcile_pdf_attachments(client, parent, root, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, root, depth=0)
 
         l1_uri = f"{root_uri}#attachment=l1.pdf"
         l2_uri = f"{l1_uri}#attachment=l2.pdf"
@@ -234,7 +235,7 @@ async def test_config_off_skips_extraction(temp_db_path, monkeypatch):
         parent_uri = "file:///fixtures/parent.pdf"
         pdf_bytes = build_pdf([("a.txt", b"A")])
         parent = await _make_parent(client, parent_uri, pdf_bytes)
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
 
         assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
 
@@ -256,7 +257,7 @@ async def test_non_pdf_parent_is_ignored(temp_db_path, monkeypatch):
         )
         # Even with PDF bytes, content_type=text/plain blocks extraction.
         pdf_bytes = build_pdf([("a.txt", b"A")])
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
         assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
 
 
@@ -272,7 +273,7 @@ async def test_parent_without_uri_is_skipped(temp_db_path, monkeypatch):
             metadata={"content_type": "application/pdf", "md5": "abc"},
         )
         pdf_bytes = build_pdf([("a.txt", b"A")])
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
         assert await client.list_documents() == []
 
 
@@ -290,7 +291,7 @@ async def test_malformed_pdf_logs_warning_and_skips(temp_db_path, monkeypatch, c
         garbage = b"this is not a pdf at all"
         parent = await _make_parent(client, parent_uri, garbage)
         with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"):
-            await _reconcile_pdf_attachments(client, parent, garbage, depth=0)
+            await _reconcile_pdf_attachments(writing(client), parent, garbage, depth=0)
         assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
 
 
@@ -328,7 +329,7 @@ async def test_unsupported_attachment_continues_loop(temp_db_path, monkeypatch):
         parent_uri = "file:///fixtures/parent.pdf"
         pdf_bytes = build_pdf([("ok.txt", b"keep me"), ("unsupported.xyz", b"data")])
         parent = await _make_parent(client, parent_uri, pdf_bytes)
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
         children = await client.list_documents(filter=parent_uri_filter(parent_uri))
         assert {c.uri for c in children} == {f"{parent_uri}#attachment=ok.txt"}
 
@@ -346,7 +347,9 @@ async def test_joboptions_attachment_skipped_not_routed_as_pdf(temp_db_path, cap
         parent = await _make_parent(client, parent_uri, pdf_bytes)
 
         with caplog.at_level(logging.WARNING, logger="haiku.rag.client.documents"):
-            await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+            await _reconcile_pdf_attachments(
+                writing(client), parent, pdf_bytes, depth=0
+            )
 
         assert await client.list_documents(filter=parent_uri_filter(parent_uri)) == []
 
@@ -360,7 +363,7 @@ async def test_cascade_delete_removes_reconciled_children(temp_db_path, monkeypa
         parent_uri = "file:///fixtures/parent.pdf"
         pdf_bytes = build_pdf([("a.txt", b"A"), ("b.txt", b"B")])
         parent = await _make_parent(client, parent_uri, pdf_bytes)
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
         assert len(await client.list_documents()) == 3
 
         await client.delete_document(parent.id)
@@ -512,7 +515,7 @@ async def test_extract_pdf_attachments_called_off_event_loop_thread(
     async with HaikuRAG(temp_db_path, create=True) as client:
         parent_uri = "file:///fixtures/parent.pdf"
         parent = await _make_parent(client, parent_uri, pdf_bytes)
-        await _reconcile_pdf_attachments(client, parent, pdf_bytes, depth=0)
+        await _reconcile_pdf_attachments(writing(client), parent, pdf_bytes, depth=0)
 
     assert called_from, "_extract_pdf_attachments was never called"
     assert called_from[0] is not event_loop_thread, (
diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py
index 3ef84811..3f464797 100644
--- a/tests/test_picture_in_context.py
+++ b/tests/test_picture_in_context.py
@@ -19,6 +19,7 @@ from haiku.rag.config import AppConfig, get_config
 from haiku.rag.store.models.chunk import Chunk, SearchResult
 from haiku.rag.store.models.document_item import DocumentItem
 from haiku.rag.tools.search import create_search_toolset
+from tests.conftest import writing
 from tests.test_context import _fetch_and_expand
 
 
@@ -244,7 +245,9 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
     async with HaikuRAG(temp_db_path, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
         before = await rag.document_item_repository.get_all_picture_data(created.id)
         assert before.get("#/pictures/0") is not None
@@ -309,7 +312,7 @@ async def test_embed_only_preserves_picture_vectors(temp_db_path, monkeypatch):
         embedded = await embed_chunks(chunks, rag.embedder, rag._config)
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        await _store_document_with_chunks(rag, document, embedded, docling_doc)
+        await _store_document_with_chunks(writing(rag), document, embedded, docling_doc)
 
         before = await _picture_chunk_row(rag)
         assert list(before["vector"]) == pytest.approx(IMAGE_VEC)
@@ -819,7 +822,7 @@ async def test_ingest_emits_picture_chunks_with_multimodal_embedder(
 
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        await _store_document_with_chunks(rag, document, embedded, docling_doc)
+        await _store_document_with_chunks(writing(rag), document, embedded, docling_doc)
 
         all_db_chunks = await rag.chunk_repository.store.chunks_table.query().to_list()
         picture_db_chunks = [
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index f9b0e951..96c860da 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -7,7 +7,7 @@ import pytest
 
 from haiku.rag.client import HaikuRAG, RebuildMode
 from haiku.rag.config import get_config
-from tests.conftest import capture_logs
+from tests.conftest import capture_logs, writing
 
 
 class ChunkData(TypedDict):
@@ -603,12 +603,14 @@ async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch)
         )
         assert doc1.id is not None and doc2.id is not None
 
-        async def fake_generate_title(doc):
+        async def fake_generate_title(config, doc):
             if doc.id == doc1.id:
                 raise RuntimeError("simulated LLM failure")
             return "Second Title"
 
-        monkeypatch.setattr(client, "generate_title", fake_generate_title)
+        monkeypatch.setattr(
+            "haiku.rag.client.rebuild.generate_title", fake_generate_title
+        )
 
         processed_ids = [
             doc_id
@@ -743,7 +745,9 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
 
         # _docling_doc_with_picture has no PageItems, so set_docling leaves
@@ -827,7 +831,9 @@ async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkey
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
 
         called_with: list[dict[str, bytes]] = []
@@ -873,7 +879,7 @@ async def test_patch_picture_descriptions_returns_zero_for_doc_without_pictures(
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         doc = await rag.create_document(content="Just text, no pictures.")
         assert doc.id is not None
-        n = await _patch_picture_descriptions(rag, doc)
+        n = await _patch_picture_descriptions(writing(rag), doc)
         assert n == 0
 
 
@@ -897,7 +903,9 @@ async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path):
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
 
         # Wipe the stored picture bytes to simulate a doc that knows about
@@ -910,7 +918,7 @@ async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path):
         from haiku.rag.client import rebuild as rebuild_module
 
         with capture_logs(rebuild_module.logger, logging.WARNING) as records:
-            n = await _patch_picture_descriptions(rag, created)
+            n = await _patch_picture_descriptions(writing(rag), created)
 
         assert n == 0
         assert any("no stored picture bytes" in r.getMessage() for r in records)
@@ -941,7 +949,9 @@ async def test_patch_picture_descriptions_skips_when_all_already_described(
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
 
         called = False
@@ -956,7 +966,7 @@ async def test_patch_picture_descriptions_skips_when_all_already_described(
             fake_describe,
         )
 
-        n = await _patch_picture_descriptions(rag, created)
+        n = await _patch_picture_descriptions(writing(rag), created)
         assert n == 0
         assert called is False
 
@@ -979,7 +989,9 @@ async def test_rebuild_descriptions_raises_when_blob_is_missing(
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
         assert created.id is not None
 
         # Force the stored doc to come back without a docling blob.
@@ -1251,9 +1263,11 @@ async def test_patch_picture_descriptions_returns_zero_without_descriptions(
     async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
         document = Document(content="x", uri="test://doc")
         document.set_docling(docling_doc)
-        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        created = await _store_document_with_chunks(
+            writing(rag), document, [], docling_doc
+        )
 
-        assert await _patch_picture_descriptions(rag, created) == 0
+        assert await _patch_picture_descriptions(writing(rag), created) == 0
 
 
 @pytest.mark.vcr()
@@ -1462,7 +1476,9 @@ async def test_rebuild_descriptions_flushes_in_batches(temp_db_path, monkeypatch
             docling_doc = _docling_doc_with_picture()
             document = Document(content=f"picture doc {i}", uri=f"test://doc-{i}")
             document.set_docling(docling_doc)
-            created = await _store_document_with_chunks(rag, document, [], docling_doc)
+            created = await _store_document_with_chunks(
+                writing(rag), document, [], docling_doc
+            )
             assert created.id is not None
             ids.append(created.id)
 
@@ -1543,7 +1559,7 @@ async def test_rebuild_embed_only_recovers_picture_bytes(
             embedding=[0.1] * rag.embedder.vector_dim,
         )
         created = await _store_document_with_chunks(
-            rag, document, [picture_chunk, text_chunk], docling_doc
+            writing(rag), document, [picture_chunk, text_chunk], docling_doc
         )
         assert created.id is not None
 
diff --git a/tests/test_title_generation.py b/tests/test_title_generation.py
index a88802c5..0c317268 100644
--- a/tests/test_title_generation.py
+++ b/tests/test_title_generation.py
@@ -353,17 +353,19 @@ class TestRebuildTitleOnly:
             await client.document_repository.update(doc2)
 
             # Make generate_title fail for the first doc, succeed for the second
-            original = HaikuRAG.generate_title
+            import haiku.rag.client.rebuild as rebuild_mod
+
+            original = rebuild_mod.generate_title
             call_count = 0
 
-            async def flaky_generate(self, document):
+            async def flaky_generate(config, document):
                 nonlocal call_count
                 call_count += 1
                 if call_count == 1:
                     raise RuntimeError("LLM failed")
-                return await original(self, document)
+                return await original(config, document)
 
-            monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate)
+            monkeypatch.setattr(rebuild_mod, "generate_title", flaky_generate)
 
             processed_ids = []
             async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py
index 57c2f3b4..a1942950 100644
--- a/tests/test_vacuum_debounce.py
+++ b/tests/test_vacuum_debounce.py
@@ -7,6 +7,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.client.documents import _refresh_doc_metadata
 from haiku.rag.config import get_config
 from haiku.rag.store.models.chunk import Chunk
+from tests.conftest import writing
 
 
 def _docling_doc(name: str, text: str):
@@ -85,7 +86,7 @@ async def test_metadata_refresh_sweep_schedules_vacuum(temp_db_path):
         client._session._vacuum_dirty = False
 
         await _refresh_doc_metadata(
-            client,
+            writing(client),
             doc,
             title=None,
             user_metadata={},
@@ -111,7 +112,7 @@ async def test_metadata_refresh_waits_for_write_lock(temp_db_path):
         async with client.store._write_lock:
             task = asyncio.create_task(
                 _refresh_doc_metadata(
-                    client,
+                    writing(client),
                     doc,
                     title=None,
                     user_metadata={},