feat(cli): Remove client cache and fix termination issues

Author: Davidyz

src/vectorcode/cli_utils.py

@@ -622,7 +622,7 @@ class LockManager:
     __singleton: "LockManager"
 
     def __new__(cls) -> "LockManager":
-        if cls.__singleton is None:
+        if not hasattr(cls, "__singleton") or cls.__singleton is None:
             cls.__singleton = super().__new__(cls)
             cls.__singleton.__locks = {}
         return cls.__singleton

src/vectorcode/common.py

@@ -114,30 +114,24 @@ async def start_server(configs: Config):
     return process
 
 
-__CLIENT_CACHE: dict[str, AsyncClientAPI] = {}
-
-
 async def get_client(configs: Config) -> AsyncClientAPI:
-    client_entry = configs.db_url
-    if __CLIENT_CACHE.get(client_entry) is None:
-        settings: dict[str, Any] = {"anonymized_telemetry": False}
-        if isinstance(configs.db_settings, dict):
-            valid_settings = {
-                k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
-            }
-            settings.update(valid_settings)
-        parsed_url = urlparse(configs.db_url)
-        settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
-        settings["chroma_server_http_port"] = parsed_url.port or 8000
-        settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
-        settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
-        settings_obj = Settings(**settings)
-        __CLIENT_CACHE[client_entry] = await chromadb.AsyncHttpClient(
-            settings=settings_obj,
-            host=str(settings_obj.chroma_server_host),
-            port=int(settings_obj.chroma_server_http_port or 8000),
-        )
-    return __CLIENT_CACHE[client_entry]
+    settings: dict[str, Any] = {"anonymized_telemetry": False}
+    if isinstance(configs.db_settings, dict):
+        valid_settings = {
+            k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
+        }
+        settings.update(valid_settings)
+    parsed_url = urlparse(configs.db_url)
+    settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
+    settings["chroma_server_http_port"] = parsed_url.port or 8000
+    settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
+    settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
+    settings_obj = Settings(**settings)
+    return await chromadb.AsyncHttpClient(
+        settings=settings_obj,
+        host=str(settings_obj.chroma_server_host),
+        port=int(settings_obj.chroma_server_http_port or 8000),
+    )
 
 
 def get_collection_name(full_path: str) -> str:
@@ -277,7 +271,7 @@ class ClientManager:
     __clients: dict[str, _ClientModel]
 
     def __new__(cls) -> "ClientManager":
-        if cls.__singleton is None:
+        if not hasattr(cls, "__singleton") or cls.__singleton is None:
             cls.__singleton = super().__new__(cls)
             cls.__singleton.__clients = {}
         return cls.__singleton

src/vectorcode/mcp_main.py

@@ -40,7 +40,7 @@
     get_project_config,
     load_config_file,
 )
-from vectorcode.common import ClientManager, get_client, get_collection, get_collections
+from vectorcode.common import ClientManager, get_collection, get_collections
 from vectorcode.subcommands.prompt import prompt_by_categories
 from vectorcode.subcommands.query import get_query_result_files
 
@@ -238,7 +238,7 @@ async def mcp_server():
         default_project_root = project_root
         default_config = await get_project_config(project_root)
         default_config.project_root = project_root
-        default_client = await get_client(default_config)
+        default_client = (await ClientManager().get_client(default_config)).client
         try:
             default_collection = await get_collection(default_client, default_config)
             logger.info("Collection initialised for %s.", project_root)
@@ -295,9 +295,16 @@ def parse_cli_args(args: Optional[list[str]] = None) -> MCPConfig:
 
 
 async def run_server():  # pragma: nocover
-    mcp = await mcp_server()
-    await mcp.run_stdio_async()
-    return 0
+    try:
+        mcp = await mcp_server()
+        await mcp.run_stdio_async()
+    finally:
+        termination_tasks: list[asyncio.Task] = []
+        for p in ClientManager().get_processes():
+            p.terminate()
+            termination_tasks.append(asyncio.create_task(p.wait()))
+        await asyncio.gather(*termination_tasks)
+        return 0
 
 
 def main():  # pragma: nocover