360 lines
11 KiB
Python
360 lines
11 KiB
Python
import json
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
from http.server import HTTPServer, BaseHTTPRequestHandler
|
|
from typing import Optional, List
|
|
from functools import partial, wraps
|
|
|
|
import pexpect
|
|
from loguru import logger
|
|
|
|
from settings import settings
|
|
|
|
|
|
def log_performance(func):
|
|
"""Decorator to log function performance"""
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
start_time = time.time()
|
|
logger.debug("Starting {}", func.__name__)
|
|
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
duration = time.time() - start_time
|
|
logger.debug("{} completed in {:.2f}s", func.__name__, duration)
|
|
return result
|
|
except Exception as e:
|
|
duration = time.time() - start_time
|
|
logger.error("{} failed after {:.2f}s: {}", func.__name__, duration, str(e))
|
|
raise
|
|
|
|
return wrapper
|
|
|
|
|
|
# Configure loguru
|
|
logger.remove() # Remove default handler
|
|
|
|
# Console handler only
|
|
logger.add(
|
|
sys.stderr,
|
|
level="DEBUG" if settings.debug else "INFO",
|
|
format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
|
|
colorize=True,
|
|
backtrace=True,
|
|
diagnose=True
|
|
)
|
|
|
|
|
|
class Config:
|
|
"""Configuration constants for the application"""
|
|
SPAWN_TIMEOUT = 60
|
|
ECHO_DELAY = 0.5
|
|
READ_TIMEOUT = 0.1
|
|
RESPONSE_WAIT_TIME = 2
|
|
CHUNK_SIZE = 1000
|
|
MAX_READ_SIZE = 10000
|
|
PROMPT_INDICATOR = "Enter your prompt"
|
|
ANSI_PATTERN = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
|
TUI_BORDER = '┃'
|
|
SKIP_PATTERNS = ['alt+enter', 'Enter your prompt']
|
|
|
|
|
|
def clean_response(response: str, original_prompt: str) -> str:
|
|
"""Clean and format MCP response"""
|
|
if not response:
|
|
return ""
|
|
|
|
# Remove ANSI escape sequences
|
|
response = Config.ANSI_PATTERN.sub('', response)
|
|
|
|
lines = response.split('\n')
|
|
cleaned_lines = []
|
|
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
|
|
# Skip empty lines and original prompt
|
|
if not stripped or stripped == original_prompt:
|
|
continue
|
|
|
|
# Handle TUI decorations
|
|
if stripped.startswith(Config.TUI_BORDER):
|
|
content = stripped.strip(Config.TUI_BORDER).strip()
|
|
if content and content != original_prompt:
|
|
cleaned_lines.append(content)
|
|
continue
|
|
|
|
# Skip navigation hints
|
|
if any(pattern in line for pattern in Config.SKIP_PATTERNS):
|
|
continue
|
|
|
|
# Add non-empty, non-decoration lines
|
|
cleaned_lines.append(stripped)
|
|
|
|
return '\n'.join(cleaned_lines)
|
|
|
|
|
|
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
|
|
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
|
|
|
|
|
|
class SubprocessHandler(BaseHTTPRequestHandler):
|
|
"""HTTP request handler with dependency injection"""
|
|
|
|
def __init__(self, *args, manager: MCPHostManager, **kwargs):
|
|
self.manager = manager
|
|
super().__init__(*args, **kwargs)
|
|
|
|
@logger.catch
|
|
def do_POST(self):
|
|
"""Handle POST requests"""
|
|
try:
|
|
content_length = self._get_content_length()
|
|
input_data = self._read_request_body(content_length)
|
|
prompt = self._parse_prompt(input_data)
|
|
|
|
logger.info("Received prompt: {}", prompt)
|
|
response = self.manager.send_prompt(prompt)
|
|
logger.info("Response: {}", response)
|
|
|
|
self._send_success_response(response)
|
|
except ValueError as e:
|
|
logger.warning("Bad request: {}", str(e))
|
|
self._send_error_response(400, str(e))
|
|
except Exception as e:
|
|
logger.exception("Unexpected error in POST handler")
|
|
self._send_error_response(500, "Internal server error")
|
|
|
|
def do_GET(self):
|
|
"""Handle GET requests (not supported)"""
|
|
self._send_error_response(404, "Not Found")
|
|
|
|
def log_message(self, format, *args):
|
|
"""Override to prevent default logging"""
|
|
pass
|
|
|
|
def _get_content_length(self) -> int:
|
|
"""Get content length from headers"""
|
|
return int(self.headers.get('Content-Length', 0))
|
|
|
|
def _read_request_body(self, content_length: int) -> str:
|
|
"""Read request body"""
|
|
if content_length == 0:
|
|
raise ValueError("Empty request body")
|
|
return self.rfile.read(content_length).decode('utf-8')
|
|
|
|
def _parse_prompt(self, input_data: str) -> str:
|
|
"""Parse prompt from input data"""
|
|
try:
|
|
data = json.loads(input_data)
|
|
return data.get('prompt', input_data)
|
|
except json.JSONDecodeError:
|
|
return input_data
|
|
|
|
def _send_success_response(self, content: str):
|
|
"""Send successful response"""
|
|
self.send_response(200)
|
|
self.send_header('Content-type', 'text/plain')
|
|
self.end_headers()
|
|
self.wfile.write(content.encode('utf-8'))
|
|
|
|
def _send_error_response(self, code: int, message: str):
|
|
"""Send error response"""
|
|
self.send_response(code)
|
|
self.send_header('Content-type', 'text/plain')
|
|
self.end_headers()
|
|
self.wfile.write(message.encode('utf-8'))
|
|
|
|
|
|
def create_server(manager: MCPHostManager) -> HTTPServer:
|
|
"""Factory function to create HTTP server with manager"""
|
|
handler = partial(SubprocessHandler, manager=manager)
|
|
return HTTPServer((settings.host, settings.port), handler)
|
|
|
|
|
|
@logger.catch
|
|
def main():
|
|
"""Simple and clean main function"""
|
|
# Startup banner
|
|
logger.info("=" * 50)
|
|
logger.info("MCP Host Server v1.0")
|
|
logger.info("Debug Mode: {}", "ON" if settings.debug else "OFF")
|
|
logger.info("=" * 50)
|
|
|
|
try:
|
|
with MCPHostManager() as manager:
|
|
server = create_server(manager)
|
|
logger.success("Server started at {}:{}", settings.host, settings.port)
|
|
logger.info("Ready to accept requests.")
|
|
logger.info("Press Ctrl+C to shutdown")
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
logger.info("Shutdown signal received")
|
|
except Exception as e:
|
|
logger.error("Fatal error: {}", e)
|
|
sys.exit(1)
|
|
finally:
|
|
if 'server' in locals():
|
|
server.shutdown()
|
|
logger.success("Server shutdown complete")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |