fix(vector-io): persist vector store metadata to kvstore in Milvus, Chroma, and Weaviate (#6371)

# What does this PR do?

Ran into #5954 while digging into why `file_search` silently returns
empty results after a server restart. The Milvus, Chroma, and Weaviate
providers only write vector store metadata to the in-memory `self.cache`
in `register_vector_store`, but both `initialize()` and
`_get_and_cache_vector_store_index()` try to read that metadata back
from the kvstore. So any store created at runtime via
`client.vector_stores.create(...)` is just gone after a restart, and in
a multi-instance setup it's only visible to whichever instance happened
to handle the create call.

I saw #5997 took a swing at this a while ago but stalled out without
tests. This adds the kvstore write on register and the matching delete
on unregister for all three providers, same as what pgvector,
sqlite_vec, and qdrant already do.

One wrinkle: Weaviate's kvstore is optional (persistence can simply be
unconfigured), so there the write and delete are skipped when no kvstore
is set instead of raising — matching its existing "registry will not
persist across restarts" behavior. Milvus and Chroma always have a
kvstore by the time register can be called, so they get the same
RuntimeError guard the other providers use.

## Test Plan

New tests in
`tests/unit/providers/vector_io/test_vector_store_kvstore_persistence.py`,
parametrized over all three providers with a real sqlite kvstore and
mocked backend clients:

- register persists the metadata to the kvstore
- a fresh adapter on the same kvstore (simulated restart) can resolve
the store
- unregister removes the kvstore entry

The register and restart tests fail on main without the fix and pass
with it. Full `tests/unit/providers/vector_io/` suite: 286 passed. Ruff
check and format clean.

---------

Signed-off-by: goingforstudying-ctrl <goingforstudying-ctrl@users.noreply.github.com>
Co-authored-by: goingforstudying-ctrl <goingforstudying-ctrl@users.noreply.github.com>

Author: goingforstudying-ctrl

src/ogx/providers/remote/vector_io/chroma/chroma.py

@@ -290,6 +290,11 @@ async def shutdown(self) -> None:
         await super().shutdown()
 
     async def register_vector_store(self, vector_store: VectorStore) -> None:
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before registering vector stores.")
+        key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
+        await self.kvstore.set(key=key, value=vector_store.model_dump_json())
+
         collection = await maybe_await(
             self.client.get_or_create_collection(
                 name=vector_store.identifier, metadata={"vector_store": vector_store.model_dump_json()}
@@ -300,12 +305,15 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
         )
 
     async def unregister_vector_store(self, vector_store_id: str) -> None:
-        if vector_store_id not in self.cache:
-            log.warning(f"Vector DB {vector_store_id} not found")
-            return
+        if vector_store_id in self.cache:
+            await self.cache[vector_store_id].index.delete()
+            del self.cache[vector_store_id]
+        else:
+            log.warning("Vector DB not found", vector_store_id=vector_store_id)
 
-        await self.cache[vector_store_id].index.delete()
-        del self.cache[vector_store_id]
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before unregistering vector stores.")
+        await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")
 
     async def insert_chunks(self, request: InsertChunksRequest) -> None:
         index = await self._get_and_cache_vector_store_index(request.vector_store_id)

src/ogx/providers/remote/vector_io/milvus/milvus.py

@@ -524,6 +524,11 @@ async def shutdown(self) -> None:
         await super().shutdown()
 
     async def register_vector_store(self, vector_store: VectorStore) -> None:
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before registering vector stores.")
+        key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
+        await self.kvstore.set(key=key, value=vector_store.model_dump_json())
+
         use_native_hybrid = isinstance(self.config, RemoteMilvusVectorIOConfig)
         if isinstance(self.config, RemoteMilvusVectorIOConfig):
             consistency_level = self.config.consistency_level
@@ -579,6 +584,10 @@ async def unregister_vector_store(self, vector_store_id: str) -> None:
             await self.cache[vector_store_id].index.delete()
             del self.cache[vector_store_id]
 
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before unregistering vector stores.")
+        await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")
+
     async def insert_chunks(self, request: InsertChunksRequest) -> None:
         index = await self._get_and_cache_vector_store_index(request.vector_store_id)
         if not index:

src/ogx/providers/remote/vector_io/weaviate/weaviate.py

@@ -379,18 +379,24 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
                 ],
             )
 
+        if self.kvstore is not None:
+            key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
+            await self.kvstore.set(key=key, value=vector_store.model_dump_json())
+
         self.cache[vector_store.identifier] = VectorStoreWithIndex(
             vector_store, WeaviateIndex(client=client, collection_name=sanitized_collection_name), self.inference_api
         )
 
     async def unregister_vector_store(self, vector_store_id: str) -> None:
         client = self._get_client()
         sanitized_collection_name = sanitize_collection_name(vector_store_id, weaviate_format=True)
-        if vector_store_id not in self.cache or client.collections.exists(sanitized_collection_name) is False:
-            return
-        client.collections.delete(sanitized_collection_name)
-        await self.cache[vector_store_id].index.delete()
-        del self.cache[vector_store_id]
+        if vector_store_id in self.cache and client.collections.exists(sanitized_collection_name):
+            client.collections.delete(sanitized_collection_name)
+            await self.cache[vector_store_id].index.delete()
+            del self.cache[vector_store_id]
+
+        if self.kvstore is not None:
+            await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")
 
     async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None:
         if vector_store_id in self.cache:

tests/unit/providers/vector_io/test_vector_store_kvstore_persistence.py

@@ -0,0 +1,173 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+# Stub optional provider dependencies so these tests run without the heavy
+# backend packages installed (same pattern as tests/unit/providers/test_milvus_weights.py).
+if "pymilvus" not in sys.modules:
+    pymilvus = ModuleType("pymilvus")
+    pymilvus.AnnSearchRequest = object
+    pymilvus.DataType = SimpleNamespace(
+        VARCHAR="VARCHAR",
+        FLOAT_VECTOR="FLOAT_VECTOR",
+        JSON="JSON",
+        SPARSE_FLOAT_VECTOR="SPARSE_FLOAT_VECTOR",
+    )
+    pymilvus.Function = object
+    pymilvus.FunctionType = SimpleNamespace(BM25="BM25")
+    pymilvus.MilvusClient = object
+    pymilvus.RRFRanker = object
+    pymilvus.WeightedRanker = object
+    sys.modules["pymilvus"] = pymilvus
+
+if "chromadb" not in sys.modules:
+    chromadb = MagicMock(name="chromadb")
+    chromadb.AsyncHttpClient = AsyncMock()
+    sys.modules["chromadb"] = chromadb
+
+if "weaviate" not in sys.modules:
+    weaviate = MagicMock(name="weaviate")
+    sys.modules["weaviate"] = weaviate
+    sys.modules["weaviate.classes"] = weaviate.classes
+    sys.modules["weaviate.classes.init"] = weaviate.classes.init
+    sys.modules["weaviate.classes.query"] = weaviate.classes.query
+
+from ogx.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
+from ogx.providers.remote.vector_io.chroma import chroma as chroma_module
+from ogx.providers.remote.vector_io.chroma.chroma import ChromaVectorIOAdapter
+from ogx.providers.remote.vector_io.chroma.config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
+from ogx.providers.remote.vector_io.milvus import milvus as milvus_module
+from ogx.providers.remote.vector_io.milvus.milvus import MilvusVectorIOAdapter
+from ogx.providers.remote.vector_io.weaviate import weaviate as weaviate_module
+from ogx.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig
+from ogx.providers.remote.vector_io.weaviate.weaviate import WeaviateVectorIOAdapter
+from ogx_api.vector_stores import VectorStore
+
+
+def _vector_store(store_id: str) -> VectorStore:
+    return VectorStore(
+        identifier=store_id,
+        provider_id="test-provider",
+        embedding_model="test-embedding-model",
+        embedding_dimension=8,
+    )
+
+
+async def _make_milvus_adapter(kvstore_config, tmp_path, monkeypatch):
+    fake_client = MagicMock()
+    fake_client.has_collection.return_value = True
+    monkeypatch.setattr(milvus_module, "MilvusClient", lambda **kwargs: fake_client)
+    config = InlineMilvusVectorIOConfig(
+        db_path=str(tmp_path / "milvus.db"),
+        persistence=kvstore_config,
+    )
+    adapter = MilvusVectorIOAdapter(config, inference_api=MagicMock(), files_api=None)
+    await adapter.initialize()
+    return adapter
+
+
+async def _make_chroma_adapter(kvstore_config, tmp_path, monkeypatch):
+    # Always fake the Chroma HTTP client: with chromadb installed locally the
+    # module stub above is skipped, and a real client would try to reach
+    # localhost:8000 during initialize(), making the test environment-dependent.
+    fake_client = MagicMock()
+    fake_client.get_or_create_collection = AsyncMock(
+        side_effect=lambda name, metadata=None: SimpleNamespace(name=name, metadata=metadata)
+    )
+    fake_client.get_collection = AsyncMock(side_effect=lambda name: SimpleNamespace(name=name))
+    fake_client.delete_collection = AsyncMock()
+    monkeypatch.setattr(chroma_module.chromadb, "AsyncHttpClient", AsyncMock(return_value=fake_client))
+    config = RemoteChromaVectorIOConfig(
+        url="http://localhost:8000",
+        persistence=kvstore_config,
+    )
+    adapter = ChromaVectorIOAdapter(config, inference_api=MagicMock(), files_api=None)
+    await adapter.initialize()
+    return adapter
+
+
+async def _make_weaviate_adapter(kvstore_config, tmp_path, monkeypatch):
+    config = WeaviateVectorIOConfig(
+        weaviate_cluster_url="localhost:8080",
+        persistence=kvstore_config,
+    )
+    adapter = WeaviateVectorIOAdapter(config, inference_api=MagicMock(), files_api=None)
+    fake_client = MagicMock()
+    monkeypatch.setattr(adapter, "_get_client", lambda: fake_client)
+    await adapter.initialize()
+    return adapter
+
+
+@pytest.mark.parametrize(
+    "make_adapter,prefix_attr",
+    [
+        (_make_milvus_adapter, (milvus_module, "VECTOR_DBS_PREFIX")),
+        (_make_chroma_adapter, (chroma_module, "VECTOR_DBS_PREFIX")),
+        (_make_weaviate_adapter, (weaviate_module, "VECTOR_DBS_PREFIX")),
+    ],
+    ids=["milvus", "chroma", "weaviate"],
+)
+async def test_register_vector_store_persists_metadata(
+    make_adapter, prefix_attr, unique_kvstore_config, tmp_path, monkeypatch
+):
+    module, attr = prefix_attr
+    adapter = await make_adapter(unique_kvstore_config, tmp_path, monkeypatch)
+    vector_store = _vector_store("vs-persist")
+
+    await adapter.register_vector_store(vector_store)
+
+    raw = await adapter.kvstore.get(f"{getattr(module, attr)}vs-persist")
+    assert raw is not None
+    persisted = VectorStore.model_validate_json(raw)
+    assert persisted.identifier == vector_store.identifier
+    assert persisted.embedding_model == vector_store.embedding_model
+    assert persisted.embedding_dimension == vector_store.embedding_dimension
+
+
+@pytest.mark.parametrize(
+    "make_adapter",
+    [_make_milvus_adapter, _make_chroma_adapter, _make_weaviate_adapter],
+    ids=["milvus", "chroma", "weaviate"],
+)
+async def test_vector_store_survives_restart(make_adapter, unique_kvstore_config, tmp_path, monkeypatch):
+    adapter = await make_adapter(unique_kvstore_config, tmp_path, monkeypatch)
+    vector_store = _vector_store("vs-restart")
+    await adapter.register_vector_store(vector_store)
+
+    # Simulate a server restart: a brand new adapter backed by the same kvstore
+    restarted = await make_adapter(unique_kvstore_config, tmp_path, monkeypatch)
+
+    index = await restarted._get_and_cache_vector_store_index("vs-restart")
+    assert index is not None
+    assert index.vector_store.identifier == "vs-restart"
+
+
+@pytest.mark.parametrize(
+    "make_adapter,prefix_attr",
+    [
+        (_make_milvus_adapter, (milvus_module, "VECTOR_DBS_PREFIX")),
+        (_make_chroma_adapter, (chroma_module, "VECTOR_DBS_PREFIX")),
+        (_make_weaviate_adapter, (weaviate_module, "VECTOR_DBS_PREFIX")),
+    ],
+    ids=["milvus", "chroma", "weaviate"],
+)
+async def test_unregister_vector_store_removes_metadata(
+    make_adapter, prefix_attr, unique_kvstore_config, tmp_path, monkeypatch
+):
+    module, attr = prefix_attr
+    adapter = await make_adapter(unique_kvstore_config, tmp_path, monkeypatch)
+    vector_store = _vector_store("vs-unregister")
+    await adapter.register_vector_store(vector_store)
+
+    await adapter.unregister_vector_store("vs-unregister")
+
+    assert await adapter.kvstore.get(f"{getattr(module, attr)}vs-unregister") is None
+    assert "vs-unregister" not in adapter.cache