close connections

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -314,7 +314,7 @@ class MetaReferenceAgentsImpl(Agents):
         return paginate_records(session_dicts, start_index, limit)
 
     async def shutdown(self) -> None:
-        pass
+        await self.persistence_store.close()
 
     # OpenAI responses
     async def get_openai_response(

llama_stack/providers/inline/batches/reference/batches.py

@@ -129,6 +129,7 @@ class ReferenceBatchesImpl(Batches):
             # don't cancel tasks - just let them stop naturally on shutdown
             # cancelling would mark batches as "cancelled" in the database
             logger.info(f"Shutdown initiated with {len(self._processing_tasks)} active batch processing tasks")
+        await self.kvstore.close()
 
     # TODO (SECURITY): this currently works w/ configured api keys, not with x-llamastack-provider-data or with user policy restrictions
     async def create_batch(

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -64,7 +64,8 @@ class MetaReferenceEvalImpl(
             benchmark = Benchmark.model_validate_json(benchmark)
             self.benchmarks[benchmark.identifier] = benchmark
 
-    async def shutdown(self) -> None: ...
+    async def shutdown(self) -> None:
+        await self.kvstore.close()
 
     async def register_benchmark(self, task_def: Benchmark) -> None:
         # Store in kvstore

llama_stack/providers/utils/kvstore/api.py

@@ -19,3 +19,7 @@ class KVStore(Protocol):
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]: ...
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
+
+    async def close(self) -> None:
+        """Close any persistent connections. Optional method for cleanup."""
+        ...

llama_stack/providers/utils/kvstore/kvstore.py

@@ -43,6 +43,10 @@ class InmemoryKVStoreImpl(KVStore):
     async def delete(self, key: str) -> None:
         del self._store[key]
 
+    async def close(self) -> None:
+        """No-op for in-memory store."""
+        pass
+
 
 async def kvstore_impl(config: KVStoreConfig) -> KVStore:
     if config.type == KVStoreType.redis.value:

tests/unit/fixtures.py

@@ -18,6 +18,7 @@ async def sqlite_kvstore(tmp_path):
     kvstore = SqliteKVStoreImpl(kvstore_config)
     await kvstore.initialize()
     yield kvstore
+    await kvstore.close()
 
 
 @pytest.fixture(scope="function")

tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py

@@ -46,12 +46,9 @@ async def test_initialize_index(vector_index):
 
 
 async def test_add_chunks_query_vector(vector_index, sample_chunks, sample_embeddings):
-    vector_index.delete()
-    vector_index.initialize()
     await vector_index.add_chunks(sample_chunks, sample_embeddings)
     resp = await vector_index.query_vector(sample_embeddings[0], k=1, score_threshold=-1)
     assert resp.chunks[0].content == sample_chunks[0].content
-    vector_index.delete()
 
 
 async def test_chunk_id_conflict(vector_index, sample_chunks, embedding_dimension):