mcphost-api/commons/mcp_manager.py

import time
import threading
import asyncio
from typing import List, Optional
import pexpect
from loguru import logger

from commons.settings import settings
from commons.response_cleaners import clean_response
from commons.logging_utils import log_performance


class Config:
    """Configuration constants for MCPHost management"""
    SPAWN_TIMEOUT = 60
    ECHO_DELAY = 0.5
    READ_TIMEOUT = 0.1
    RESPONSE_WAIT_TIME = 1
    CHUNK_SIZE = 1000
    MAX_READ_SIZE = 10000
    PROMPT_INDICATOR = "Enter your prompt"


class MCPHostManager:
    """Manages MCP process lifecycle and communication"""

    def __init__(self):
        self.child: Optional[pexpect.spawn] = None
        self.config = Config()
        self.lock = threading.Lock()

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown()

    @log_performance
    def start(self) -> bool:
        """Start the mcphost process"""
        command = self._build_command()
        logger.info("Starting mcphost: {}", ' '.join(command))

        try:
            self.child = pexpect.spawn(
                ' '.join(command),
                timeout=self.config.SPAWN_TIMEOUT,
                encoding='utf-8'
            )
            self.child.setecho(False)
            return self._wait_for_ready()
        except Exception as e:
            logger.error("Error starting mcphost: {}", e)
            logger.exception("Full traceback:")
            return False

    @log_performance
    async def send_prompt_async(self, prompt: str) -> str:
        """Send a prompt to mcphost and get the response (async wrapper)"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.send_prompt, prompt)

    def send_prompt(self, prompt: str) -> str:
        """Send a prompt to mcphost and get the response"""
        with self.lock:
            if not self._is_alive():
                logger.warning("MCPHost not running, attempting to restart...")
                if not self.start():
                    return "Error: Failed to restart MCPHost"

            try:
                self._clear_pending_output()
                self._send_command(prompt)
                response = self._collect_response()
                return clean_response(response, prompt)
            except Exception as e:
                logger.exception("Exception in send_prompt: {}", str(e))
                return f"Error: {str(e)}"

    def shutdown(self):
        """Shutdown mcphost gracefully"""
        if self._is_alive():
            logger.info("Shutting down mcphost...")
            self.child.sendcontrol('c')
            self.child.close()
            logger.info("MCPHost shutdown complete")

    def _is_alive(self) -> bool:
        """Check if the process is running"""
        return self.child is not None and self.child.isalive()

    def _build_command(self) -> List[str]:
        """Build the command to start mcphost"""
        command = [
            settings.mcphost_path,
            '--config', settings.mcphost_config,
            '--model', settings.mcphost_model,
            '--openai-url', settings.openai_url,
            '--openai-api-key', settings.openai_api_key
        ]

        if settings.debug:
            command.insert(1, '--debug')

        return command

    def _wait_for_ready(self) -> bool:
        """Wait for the process to be ready"""
        try:
            self.child.expect(self.config.PROMPT_INDICATOR)
            logger.success("MCPHost started and ready")
            self._clear_buffer()
            return True
        except Exception as e:
            logger.error("Error waiting for prompt: {}", e)
            return False

    def _clear_buffer(self):
        """Clear any remaining output in the buffer"""
        time.sleep(self.config.ECHO_DELAY)
        try:
            self.child.read_nonblocking(
                size=self.config.MAX_READ_SIZE,
                timeout=self.config.READ_TIMEOUT
            )
        except:
            pass

    def _clear_pending_output(self):
        """Clear any pending output from the process"""
        try:
            self.child.read_nonblocking(
                size=self.config.MAX_READ_SIZE,
                timeout=self.config.READ_TIMEOUT
            )
        except:
            pass

    def _send_command(self, prompt: str):
        """Send a command to the process"""
        logger.debug("Sending prompt: {}", prompt)
        self.child.send(prompt)
        self.child.send('\r')

        # Wait for the model to process
        time.sleep(self.config.RESPONSE_WAIT_TIME)

    def _collect_response(self) -> str:
        """Collect response from the process"""
        response = ""
        response_complete = False

        with logger.catch(message="Error during response collection"):
            while not response_complete:
                try:
                    chunk = self.child.read_nonblocking(
                        size=self.config.CHUNK_SIZE,
                        timeout=3
                    )

                    if chunk:
                        response += chunk
                        logger.trace("Received chunk: {}", chunk[:50] + "..." if len(chunk) > 50 else chunk)
                        if self.config.PROMPT_INDICATOR in chunk:
                            response_complete = True
                    else:
                        break

                except pexpect.TIMEOUT:
                    if response and self.config.PROMPT_INDICATOR in response:
                        response_complete = True
                    elif response:
                        logger.debug("Waiting for more response data...")
                        time.sleep(1)
                        continue
                    else:
                        break
                except Exception as e:
                    logger.error("Error reading response: {}", e)
                    break

        logger.debug("Collected response length: {} characters", len(response))
        return response