Merge branch 'main' of ssh://git.biggo.com:222/Funmula/dive-mcp-host

Author: LouisFunmula

dive_mcp_host/host/host.py

@@ -95,6 +95,7 @@ def __init__(
         )
         self._store = store_manager
         self._exit_stack: AsyncExitStack | None = None
+        self._lock: asyncio.Lock = asyncio.Lock()
 
     async def _run_in_context(self) -> AsyncGenerator[Self, None]:
         async with AsyncExitStack() as stack:
@@ -198,45 +199,48 @@ async def reload(
         Chats can be resumed after reload by using the same chat_id.
         """
         # NOTE: Do Not restart MCP Servers when there is on-going query.
-        if self._exit_stack is None:
-            raise RuntimeError("Host not initialized")
-
-        # Update config
-        old_config = self._config
-        self._config = new_config
-
-        try:
-            # Reload model if needed
-            if old_config.llm != new_config.llm:
-                self._model = None
-                await self._init_models()
-
-            await self._tool_manager.reload(
-                new_configs=new_config.mcp_servers, force=force_mcp
-            )
-
-            # Reload checkpointer if needed
-            if old_config.checkpointer != new_config.checkpointer:
-                if self._checkpointer is not None:
-                    await self._exit_stack.aclose()
-                    self._checkpointer = None
-
-                if new_config.checkpointer:
-                    checkpointer = get_checkpointer(str(new_config.checkpointer.uri))
-                    self._checkpointer = await self._exit_stack.enter_async_context(
-                        checkpointer
-                    )
-                    await self._checkpointer.setup()
-
-            # Call the reloader function to handle service restart
-            if reloader:
-                await reloader()
-
-        except Exception as e:
-            # Restore old config if reload fails
-            self._config = old_config
-            logging.error("Failed to reload host: %s", e)
-            raise
+        async with self._lock:
+            if self._exit_stack is None:
+                raise RuntimeError("Host not initialized")
+
+            # Update config
+            old_config = self._config
+            self._config = new_config
+
+            try:
+                # Reload model if needed
+                if old_config.llm != new_config.llm:
+                    self._model = None
+                    await self._init_models()
+
+                await self._tool_manager.reload(
+                    new_configs=new_config.mcp_servers, force=force_mcp
+                )
+
+                # Reload checkpointer if needed
+                if old_config.checkpointer != new_config.checkpointer:
+                    if self._checkpointer is not None:
+                        await self._exit_stack.aclose()
+                        self._checkpointer = None
+
+                    if new_config.checkpointer:
+                        checkpointer = get_checkpointer(
+                            str(new_config.checkpointer.uri)
+                        )
+                        self._checkpointer = await self._exit_stack.enter_async_context(
+                            checkpointer
+                        )
+                        await self._checkpointer.setup()
+
+                # Call the reloader function to handle service restart
+                if reloader:
+                    await reloader()
+
+            except Exception as e:
+                # Restore old config if reload fails
+                self._config = old_config
+                logging.error("Failed to reload host: %s", e)
+                raise
 
     @property
     def config(self) -> HostConfig: