Fix tiny-agents cli exit issues (#3125)

* first draft

* Proper exit events

* Update src/huggingface_hub/inference/_mcp/_cli_hacks.py

Co-authored-by: célina <[email protected]>

* Update src/huggingface_hub/inference/_mcp/_cli_hacks.py

Co-authored-by: célina <[email protected]>

* Update src/huggingface_hub/inference/_mcp/_cli_hacks.py

Co-authored-by: célina <[email protected]>

* make style

* comment

* exit_event required

---------

Co-authored-by: célina <[email protected]>

Author: Wauplin

src/huggingface_hub/inference/_mcp/_cli_hacks.py

@@ -0,0 +1,87 @@
+import asyncio
+import sys
+from functools import partial
+
+import typer
+
+
+def _patch_anyio_open_process():
+    """
+    Patch anyio.open_process to allow detached processes on Windows and Unix-like systems.
+
+    This is necessary to prevent the MCP client from being interrupted by Ctrl+C when running in the CLI.
+    """
+    import subprocess
+
+    import anyio
+
+    if getattr(anyio, "_tiny_agents_patched", False):
+        return
+    anyio._tiny_agents_patched = True
+
+    original_open_process = anyio.open_process
+
+    if sys.platform == "win32":
+        # On Windows, we need to set the creation flags to create a new process group
+
+        async def open_process_in_new_group(*args, **kwargs):
+            """
+            Wrapper for open_process to handle Windows-specific process creation flags.
+            """
+            # Ensure we pass the creation flags for Windows
+            kwargs.setdefault("creationflags", subprocess.CREATE_NEW_PROCESS_GROUP)
+            return await original_open_process(*args, **kwargs)
+
+        anyio.open_process = open_process_in_new_group
+    else:
+        # For Unix-like systems, we can use setsid to create a new session
+        async def open_process_in_new_group(*args, **kwargs):
+            """
+            Wrapper for open_process to handle Unix-like systems with start_new_session=True.
+            """
+            kwargs.setdefault("start_new_session", True)
+            return await original_open_process(*args, **kwargs)
+
+        anyio.open_process = open_process_in_new_group
+
+
+async def _async_prompt(exit_event: asyncio.Event, prompt: str = "» ") -> str:
+    """
+    Asynchronous prompt function that reads input from stdin without blocking.
+
+    This function is designed to work in an asynchronous context, allowing the event loop to gracefully stop it (e.g. on Ctrl+C).
+
+    Alternatively, we could use https://github.com/vxgmichel/aioconsole but that would be an additional dependency.
+    """
+    loop = asyncio.get_event_loop()
+
+    if sys.platform == "win32":
+        # Windows: Use run_in_executor to avoid blocking the event loop
+        # Degraded solution: this is not ideal as user will have to CTRL+C once more to stop the prompt (and it'll not be graceful)
+        return await loop.run_in_executor(None, partial(typer.prompt, prompt, prompt_suffix=" "))
+    else:
+        # UNIX-like: Use loop.add_reader for non-blocking stdin read
+        future = loop.create_future()
+
+        def on_input():
+            line = sys.stdin.readline()
+            loop.remove_reader(sys.stdin)
+            future.set_result(line)
+
+        print(prompt, end=" ", flush=True)
+        loop.add_reader(sys.stdin, on_input)  # not supported on Windows
+
+        # Wait for user input or exit event
+        # Wait until either the user hits enter or exit_event is set
+        await asyncio.wait(
+            [future, exit_event.wait()],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        # Check which one has been triggered
+        if exit_event.is_set():
+            future.cancel()
+            return ""
+
+        line = await future
+        return line.strip()

src/huggingface_hub/inference/_mcp/cli.py

@@ -1,13 +1,12 @@
 import asyncio
-import os
 import signal
 import traceback
-from functools import partial
 from typing import Any, Dict, List, Optional
 
 import typer
 from rich import print
 
+from ._cli_hacks import _async_prompt, _patch_anyio_open_process
 from .agent import Agent
 from .utils import _load_agent_config
 
@@ -25,11 +24,6 @@
 app.add_typer(run_cli, name="run")
 
 
-async def _ainput(prompt: str = "» ") -> str:
-    loop = asyncio.get_running_loop()
-    return await loop.run_in_executor(None, partial(typer.prompt, prompt, prompt_suffix=" "))
-
-
 async def run_agent(
     agent_path: Optional[str],
 ) -> None:
@@ -41,11 +35,14 @@ async def run_agent(
             Path to a local folder containing an `agent.json` and optionally a custom `PROMPT.md` file or a built-in agent stored in a Hugging Face dataset.
 
     """
+    _patch_anyio_open_process()  # Hacky way to prevent stdio connections to be stopped by Ctrl+C
+
     config, prompt = _load_agent_config(agent_path)
 
     servers: List[Dict[str, Any]] = config.get("servers", [])
 
     abort_event = asyncio.Event()
+    exit_event = asyncio.Event()
     first_sigint = True
 
     loop = asyncio.get_running_loop()
@@ -60,8 +57,7 @@ def _sigint_handler() -> None:
             return
 
         print("\n[red]Exiting...[/red]", flush=True)
-
-        os._exit(130)
+        exit_event.set()
 
     try:
         sigint_registered_in_loop = False
@@ -71,6 +67,7 @@ def _sigint_handler() -> None:
         except (AttributeError, NotImplementedError):
             # Windows (or any loop that doesn't support it) : fall back to sync
             signal.signal(signal.SIGINT, lambda *_: _sigint_handler())
+
         async with Agent(
             provider=config.get("provider"),
             model=config.get("model"),
@@ -86,8 +83,12 @@ def _sigint_handler() -> None:
             while True:
                 abort_event.clear()
 
+                # Check if we should exit
+                if exit_event.is_set():
+                    break
+
                 try:
-                    user_input = await _ainput()
+                    user_input = await _async_prompt(exit_event=exit_event)
                     first_sigint = True
                 except EOFError:
                     print("\n[red]EOF received, exiting.[/red]", flush=True)
@@ -103,6 +104,8 @@ def _sigint_handler() -> None:
                     async for chunk in agent.run(user_input, abort_event=abort_event):
                         if abort_event.is_set() and not first_sigint:
                             break
+                        if exit_event.is_set():
+                            break
 
                         if hasattr(chunk, "choices"):
                             delta = chunk.choices[0].delta

src/huggingface_hub/inference/_mcp/mcp_client.py

@@ -109,6 +109,11 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.client.__aexit__(exc_type, exc_val, exc_tb)
         await self.cleanup()
 
+    async def cleanup(self):
+        """Clean up resources"""
+        await self.client.close()
+        await self.exit_stack.aclose()
+
     @overload
     async def add_mcp_server(self, type: Literal["stdio"], **params: Unpack[StdioServerParameters_T]): ...
 
@@ -329,7 +334,3 @@ async def process_single_turn_with_tools(
             tool_message_as_obj = ChatCompletionInputMessage.parse_obj_as_instance(tool_message)
             messages.append(tool_message_as_obj)
             yield tool_message_as_obj
-
-    async def cleanup(self):
-        """Clean up resources"""
-        await self.exit_stack.aclose()