refactor(cli): Refactor client termination to ClientManager

Author: Zhe Yu

src/vectorcode/common.py

@@ -293,3 +293,10 @@ async def get_client(self, configs: Config) -> _ClientModel:
 
     def get_processes(self) -> list[Process]:
         return [i.process for i in self.__clients.values() if i.process is not None]
+
+    async def kill_servers(self):
+        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)

src/vectorcode/lsp_main.py

@@ -43,7 +43,7 @@
     get_project_config,
     parse_cli_args,
 )
-from vectorcode.common import get_client, get_collection, try_server
+from vectorcode.common import ClientManager, get_collection, try_server
 from vectorcode.subcommands.ls import get_collection_list
 from vectorcode.subcommands.query import build_query_results
 
@@ -114,7 +114,7 @@ async def execute_command(ls: LanguageServer, args: list[str]):
                 parsed_args.project_root
             ].merge_from(parsed_args)
             final_configs.pipe = True
-            client = await get_client(final_configs)
+            client = (await ClientManager().get_client(final_configs)).client
             if final_configs.action in {CliAction.vectorise, CliAction.query}:
                 collection = await get_collection(
                     client=client,
@@ -123,7 +123,7 @@ async def execute_command(ls: LanguageServer, args: list[str]):
                 )
         else:
             final_configs = parsed_args
-            client = await get_client(parsed_args)
+            client = (await ClientManager().get_client(final_configs)).client
             collection = None
         logger.info("Merged final configs: %s", final_configs)
         progress_token = str(uuid.uuid4())
@@ -266,9 +266,11 @@ async def lsp_start() -> int:
         logger.info(f"{DEFAULT_PROJECT_ROOT=}")
 
     logger.info("Parsed LSP server CLI arguments: %s", args)
-    await asyncio.to_thread(server.start_io)
-
-    return 0
+    try:
+        await asyncio.to_thread(server.start_io)
+    finally:
+        await ClientManager().kill_servers()
+        return 0
 
 
 def main():  # pragma: nocover

src/vectorcode/main.py

@@ -100,9 +100,7 @@ async def async_main():
         return_val = 1
         logger.error(traceback.format_exc())
     finally:
-        for p in ClientManager().get_processes():
-            p.terminate()
-            await p.wait()
+        await ClientManager().kill_servers()
         return return_val
 
 

src/vectorcode/mcp_main.py

@@ -299,11 +299,7 @@ async def run_server():  # pragma: nocover
         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)
+        await ClientManager().kill_servers()
         return 0