Python MCP или как подключить свою LLM ко всему миру - Продвинутые сценарии

from abc import ABC, abstractmethod
from typing import Optional, AsyncIterator
import aiohttp
import anthropic
import openai
 
class LLMProvider(ABC):
   """Базовый класс для провайдеров"""
   
   @abstractmethod
   async def generate(
       self,
       prompt: str,
       max_tokens: int = 1000,
       temperature: float = 0.7
   ) -> str:
       pass
   
   @abstractmethod
   async def stream_generate(
       self,
       prompt: str,
       max_tokens: int = 1000
   ) -> AsyncIterator[str]:
       pass
 
class OpenAIProvider(LLMProvider):
   def __init__(self, api_key: str, model: str = "gpt-4"):
       self.client = openai.AsyncOpenAI(api_key=api_key)
       self.model = model
   
   async def generate(
       self,
       prompt: str,
       max_tokens: int = 1000,
       temperature: float = 0.7
   ) -> str:
       response = await self.client.chat.completions.create(
           model=self.model,
           messages=[{"role": "user", "content": prompt}],
           max_tokens=max_tokens,
           temperature=temperature
       )
       return response.choices[0].message.content
   
   async def stream_generate(
       self,
       prompt: str,
       max_tokens: int = 1000
   ) -> AsyncIterator[str]:
       stream = await self.client.chat.completions.create(
           model=self.model,
           messages=[{"role": "user", "content": prompt}],
           max_tokens=max_tokens,
           stream=True
       )
       async for chunk in stream:
           if chunk.choices[0].delta.content:
               yield chunk.choices[0].delta.content
 
class AnthropicProvider(LLMProvider):
   def __init__(self, api_key: str, model: str = "claude-3-opus-20240229"):
       self.client = anthropic.AsyncAnthropic(api_key=api_key)
       self.model = model
   
   async def generate(
       self,
       prompt: str,
       max_tokens: int = 1000,
       temperature: float = 0.7
   ) -> str:
       message = await self.client.messages.create(
           model=self.model,
           max_tokens=max_tokens,
           temperature=temperature,
           messages=[{"role": "user", "content": prompt}]
       )
       return message.content[0].text
   
   async def stream_generate(
       self,
       prompt: str,
       max_tokens: int = 1000
   ) -> AsyncIterator[str]:
       async with self.client.messages.stream(
           model=self.model,
           max_tokens=max_tokens,
           messages=[{"role": "user", "content": prompt}]
       ) as stream:
           async for text in stream.text_stream:
               yield text
 
# Фабрика для создания провайдера по имени
class ProviderFactory:
   @staticmethod
   def create(provider_type: str, **kwargs) -> LLMProvider:
       providers = {
           "openai": OpenAIProvider,
           "anthropic": AnthropicProvider,
       }
       
       if provider_type not in providers:
           raise ValueError(f"Неизвестный провайдер: {provider_type}")
       
       return providers[provider_type](**kwargs)

from mcp.types import Tool, TextContent
 
# Провайдер инициализируется из конфига
current_provider = ProviderFactory.create(
   provider_type="anthropic",
   api_key="sk-ant-..."
)
 
@server.list_tools()
async def handle_list_tools():
   return [
       Tool(
           name="generate_response",
           description="Генерирует ответ используя текущую LLM",
           inputSchema={
               "type": "object",
               "properties": {
                   "prompt": {"type": "string"},
                   "max_tokens": {"type": "integer", "default": 1000},
                   "temperature": {"type": "number", "default": 0.7}
               },
               "required": ["prompt"]
           }
       )
   ]
 
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
   if name == "generate_response":
       prompt = arguments["prompt"]
       max_tokens = arguments.get("max_tokens", 1000)
       temperature = arguments.get("temperature", 0.7)
       
       response = await current_provider.generate(
           prompt=prompt,
           max_tokens=max_tokens,
           temperature=temperature
       )
       
       return [TextContent(type="text", text=response)]

class FallbackProvider(LLMProvider):
   def __init__(self, primary: LLMProvider, fallback: LLMProvider):
       self.primary = primary
       self.fallback = fallback
   
   async def generate(self, prompt: str, **kwargs) -> str:
       try:
           return await self.primary.generate(prompt, **kwargs)
       except Exception as e:
           logger.warning(f"Primary провайдер упал: {e}, переключаюсь на fallback")
           return await self.fallback.generate(prompt, **kwargs)

import asyncpg
from typing import Optional, List, Dict, Any
import json
from contextlib import asynccontextmanager
 
class PostgresAdapter:
    """Асинхронный адаптер для работы с PostgreSQL"""
    
    def __init__(
        self,
        host: str,
        port: int,
        database: str,
        user: str,
        password: str,
        min_pool_size: int = 5,
        max_pool_size: int = 20
    ):
        self.connection_params = {
            "host": host,
            "port": port,
            "database": database,
            "user": user,
            "password": password
        }
        self.min_pool_size = min_pool_size
        self.max_pool_size = max_pool_size
        self.pool: Optional[asyncpg.Pool] = None
    
    async def initialize(self):
        """Создаём пул соединений"""
        self.pool = await asyncpg.create_pool(
            **self.connection_params,
            min_size=self.min_pool_size,
            max_size=self.max_pool_size,
            command_timeout=60  # таймаут 60 секунд на запрос
        )
    
    async def close(self):
        """Закрываем пул"""
        if self.pool:
            await self.pool.close()
    
    @asynccontextmanager
    async def acquire(self):
        """Контекстный менеджер для получения соединения из пула"""
        async with self.pool.acquire() as connection:
            yield connection
    
    async def execute_query(
        self, 
        query: str, 
        params: Optional[List] = None
    ) -> List[Dict[str, Any]]:
        """Выполняет SELECT запрос и возвращает результаты"""
        async with self.acquire() as conn:
            # Выполняем запрос
            rows = await conn.fetch(query, *(params or []))
            
            # Конвертируем Record объекты в словари
            results = [dict(row) for row in rows]
            return results
    
    async def execute_command(
        self,
        query: str,
        params: Optional[List] = None
    ) -> str:
        """Выполняет INSERT/UPDATE/DELETE и возвращает статус"""
        async with self.acquire() as conn:
            result = await conn.execute(query, *(params or []))
            return result
    
    async def get_table_schema(self, table_name: str) -> Dict[str, List[Dict]]:
        """Получает схему таблицы для генерации промптов"""
        query = """
            SELECT 
                column_name,
                data_type,
                is_nullable,
                column_default
            FROM information_schema.columns
            WHERE table_name = $1
            ORDER BY ordinal_position
        """
        
        async with self.acquire() as conn:
            rows = await conn.fetch(query, table_name)
            columns = [
                {
                    "name": row["column_name"],
                    "type": row["data_type"],
                    "nullable": row["is_nullable"] == "YES",
                    "default": row["column_default"]
                }
                for row in rows
            ]
            return {"table": table_name, "columns": columns}
    
    async def get_available_tables(self) -> List[str]:
        """Возвращает список всех доступных таблиц"""
        query = """
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public'
            AND table_type = 'BASE TABLE'
            ORDER BY table_name
        """
        
        async with self.acquire() as conn:
            rows = await conn.fetch(query)
            return [row["table_name"] for row in rows]

from mcp.types import Tool, TextContent
 
# Инициализируем адаптер при старте сервера
db_adapter = PostgresAdapter(
    host="localhost",
    port=5432,
    database="analytics",
    user="analytics_user",
    password="секретный_пароль"
)
 
@server.list_tools()
async def handle_list_tools():
    """Регистрируем инструменты для работы с БД"""
    return [
        Tool(
            name="query_database",
            description="Выполняет SQL SELECT запрос и возвращает результаты",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "SQL запрос для выполнения"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Максимальное количество строк",
                        "default": 100
                    }
                },
                "required": ["query"]
            }
        ),
        Tool(
            name="get_table_info",
            description="Получает структуру таблицы - колонки, типы данных",
            inputSchema={
                "type": "object",
                "properties": {
                    "table_name": {
                        "type": "string",
                        "description": "Название таблицы"
                    }
                },
                "required": ["table_name"]
            }
        ),
        Tool(
            name="list_tables",
            description="Возвращает список всех доступных таблиц в БД",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        )
    ]
 
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    """Обрабатываем вызовы инструментов БД"""
    
    if name == "query_database":
        query = arguments["query"]
        limit = arguments.get("limit", 100)
        
        # Добавляем LIMIT если его нет
        if "limit" not in query.lower():
            query = f"{query.rstrip(';')} LIMIT {limit}"
        
        try:
            results = await db_adapter.execute_query(query)
            
            # Форматируем результаты в читаемый вид
            if not results:
                return [TextContent(
                    type="text",
                    text="Запрос выполнен успешно, но не вернул данных"
                )]
            
            return [TextContent(
                type="text",
                text=json.dumps(results, ensure_ascii=False, indent=2)
            )]
            
        except asyncpg.PostgresSyntaxError as e:
            return [TextContent(
                type="text",
                text=f"Синтаксическая ошибка в SQL: {str(e)}"
            )]
        except Exception as e:
            return [TextContent(
                type="text",
                text=f"Ошибка выполнения запроса: {str(e)}"
            )]
    
    elif name == "get_table_info":
        table_name = arguments["table_name"]
        
        try:
            schema = await db_adapter.get_table_schema(table_name)
            return [TextContent(
                type="text",
                text=json.dumps(schema, ensure_ascii=False, indent=2)
            )]
        except Exception as e:
            return [TextContent(
                type="text",
                text=f"Не удалось получить схему таблицы: {str(e)}"
            )]
    
    elif name == "list_tables":
        try:
            tables = await db_adapter.get_available_tables()
            return [TextContent(
                type="text",
                text=f"Доступные таблицы: {', '.join(tables)}"
            )]
        except Exception as e:
            return [TextContent(
                type="text",
                text=f"Ошибка получения списка таблиц: {str(e)}"
            )]
 
# Инициализация и cleanup
async def startup():
    """Вызывается при старте сервера"""
    await db_adapter.initialize()
 
async def shutdown():
    """Вызывается при остановке сервера"""
    await db_adapter.close()

from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from typing import List, Dict, Any, Optional
import json
 
class MongoAdapter:
    """Асинхронный адаптер для MongoDB"""
    
    def __init__(
        self,
        connection_string: str,
        database_name: str,
        max_pool_size: int = 50
    ):
        self.client = AsyncIOMotorClient(
            connection_string,
            maxPoolSize=max_pool_size,
            serverSelectionTimeoutMS=5000
        )
        self.db: AsyncIOMotorDatabase = self.client[database_name]
    
    async def find_documents(
        self,
        collection: str,
        filter_query: Dict[str, Any],
        projection: Optional[Dict[str, int]] = None,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """Находит документы по запросу"""
        coll = self.db[collection]
        
        cursor = coll.find(
            filter=filter_query,
            projection=projection
        ).limit(limit)
        
        documents = await cursor.to_list(length=limit)
        
        # Конвертируем ObjectId в строки для JSON-сериализации
        for doc in documents:
            if '_id' in doc:
                doc['_id'] = str(doc['_id'])
        
        return documents
    
    async def aggregate(
        self,
        collection: str,
        pipeline: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Выполняет aggregation pipeline"""
        coll = self.db[collection]
        
        cursor = coll.aggregate(pipeline)
        results = await cursor.to_list(length=None)
        
        for doc in results:
            if '_id' in doc:
                doc['_id'] = str(doc['_id'])
        
        return results
    
    async def get_collections(self) -> List[str]:
        """Возвращает список коллекций"""
        return await self.db.list_collection_names()
    
    async def get_collection_stats(
        self,
        collection: str
    ) -> Dict[str, Any]:
        """Получает статистику коллекции"""
        return await self.db.command("collStats", collection)

from aioch import Client as ClickHouseClient
from typing import List, Dict, Any
 
class ClickHouseAdapter:
    """Адаптер для ClickHouse"""
    
    def __init__(self, host: str, port: int = 9000):
        self.client = ClickHouseClient(host=host, port=port)
    
    async def execute_query(
        self,
        query: str
    ) -> List[Dict[str, Any]]:
        """Выполняет аналитический запрос"""
        # ClickHouse возвращает данные как список кортежей
        rows = await self.client.execute(query, with_column_types=True)
        
        if not rows:
            return []
        
        # Извлекаем названия и типы колонок
        columns = [(name, type_) for name, type_ in rows[1]]
        data = rows[0]
        
        # Преобразуем в список словарей
        results = []
        for row in data:
            doc = {}
            for (col_name, col_type), value in zip(columns, row):
                doc[col_name] = value
            results.append(doc)
        
        return results

class UnifiedDatabaseAdapter:
    """Объединяет все адаптеры БД"""
    
    def __init__(self):
        self.postgres = PostgresAdapter(...)
        self.mongo = MongoAdapter(...)
        self.clickhouse = ClickHouseAdapter(...)
        self.adapters = {
            "postgres": self.postgres,
            "mongodb": self.mongo,
            "clickhouse": self.clickhouse
        }
    
    async def query(
        self,
        database_type: str,
        query: str,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """Универсальный метод запроса"""
        adapter = self.adapters.get(database_type)
        if not adapter:
            raise ValueError(f"Неизвестная БД: {database_type}")
        
        if database_type == "postgres":
            return await adapter.execute_query(query)
        elif database_type == "mongodb":
            # Парсим MongoDB query из JSON
            collection = kwargs.get("collection")
            filter_query = json.loads(query)
            return await adapter.find_documents(collection, filter_query)
        elif database_type == "clickhouse":
            return await adapter.execute_query(query)

import aiofiles
import aiofiles.os
from pathlib import Path
from typing import List, Optional, AsyncIterator
import asyncio
 
class FileSystemAdapter:
    """Асинхронный адаптер для работы с файловой системой"""
    
    def __init__(self, base_path: str, max_file_size: int = 10_000_000):
        self.base_path = Path(base_path).resolve()
        self.max_file_size = max_file_size  # ограничение 10MB
    
    def _validate_path(self, file_path: str) -> Path:
        """Проверяет что путь безопасен - внутри base_path"""
        full_path = (self.base_path / file_path).resolve()
        
        # Защита от path traversal атак
        if not str(full_path).startswith(str(self.base_path)):
            raise ValueError(f"Доступ за пределы base_path запрещён: {file_path}")
        
        return full_path
    
    async def read_file(self, file_path: str) -> str:
        """Читает содержимое файла"""
        path = self._validate_path(file_path)
        
        # Проверяем размер перед чтением
        stat = await aiofiles.os.stat(path)
        if stat.st_size > self.max_file_size:
            raise ValueError(f"Файл слишком большой: {stat.st_size} байт")
        
        async with aiofiles.open(path, mode='r', encoding='utf-8') as f:
            return await f.read()
    
    async def write_file(self, file_path: str, content: str) -> None:
        """Записывает контент в файл"""
        path = self._validate_path(file_path)
        
        # Создаём директории если не существуют
        await aiofiles.os.makedirs(path.parent, exist_ok=True)
        
        async with aiofiles.open(path, mode='w', encoding='utf-8') as f:
            await f.write(content)
    
    async def list_files(
        self,
        directory: str = "",
        pattern: str = "*",
        recursive: bool = False
    ) -> List[str]:
        """Перечисляет файлы в директории"""
        dir_path = self._validate_path(directory)
        
        if not await aiofiles.os.path.isdir(dir_path):
            raise ValueError(f"Не директория: {directory}")
        
        files = []
        search_pattern = "**/" + pattern if recursive else pattern
        
        for path in dir_path.glob(search_pattern):
            if path.is_file():
                # Возвращаем относительный путь
                rel_path = path.relative_to(self.base_path)
                files.append(str(rel_path))
        
        return sorted(files)
    
    async def search_in_files(
        self,
        query: str,
        directory: str = "",
        extensions: Optional[List[str]] = None
    ) -> List[dict]:
        """Ищет строку в файлах"""
        dir_path = self._validate_path(directory)
        results = []
        
        # Собираем файлы для поиска
        files_to_search = []
        for path in dir_path.rglob("*"):
            if not path.is_file():
                continue
            
            if extensions and path.suffix not in extensions:
                continue
            
            files_to_search.append(path)
        
        # Ищем параллельно
        search_tasks = [
            self._search_in_file(path, query) 
            for path in files_to_search
        ]
        
        search_results = await asyncio.gather(*search_tasks)
        
        # Фильтруем пустые результаты
        for result in search_results:
            if result["matches"]:
                results.append(result)
        
        return results
    
    async def _search_in_file(self, path: Path, query: str) -> dict:
        """Ищет в конкретном файле"""
        matches = []
        
        try:
            async with aiofiles.open(path, mode='r', encoding='utf-8') as f:
                line_num = 0
                async for line in f:
                    line_num += 1
                    if query.lower() in line.lower():
                        matches.append({
                            "line": line_num,
                            "content": line.strip()
                        })
        except (UnicodeDecodeError, PermissionError):
            # Пропускаем бинарные файлы и файлы без доступа
            pass
        
        rel_path = path.relative_to(self.base_path)
        return {
            "file": str(rel_path),
            "matches": matches
        }

import aiohttp
import asyncio
from typing import Optional, Dict, Any
from tenacity import retry, stop_after_attempt, wait_exponential
import time
 
class APIClient:
    """Универсальный асинхронный клиент для внешних API"""
    
    def __init__(
        self,
        base_url: str,
        api_key: Optional[str] = None,
        timeout: int = 30,
        rate_limit: int = 100  # запросов в минуту
    ):
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        
        # Rate limiting через token bucket
        self.rate_limit = rate_limit
        self.tokens = rate_limit
        self.last_update = time.time()
        self.lock = asyncio.Lock()
    
    async def _check_rate_limit(self):
        """Проверяет и обновляет rate limit"""
        async with self.lock:
            now = time.time()
            elapsed = now - self.last_update
            
            # Добавляем токены пропорционально времени
            self.tokens = min(
                self.rate_limit,
                self.tokens + (elapsed * self.rate_limit / 60)
            )
            self.last_update = now
            
            if self.tokens < 1:
                # Ждём пока появится токен
                wait_time = (1 - self.tokens) * 60 / self.rate_limit
                await asyncio.sleep(wait_time)
                self.tokens = 1
            
            self.tokens -= 1
    
    def _prepare_headers(self) -> Dict[str, str]:
        """Формирует заголовки с аутентификацией"""
        headers = {
            "Content-Type": "application/json",
            "User-Agent": "MCP-Server/1.0"
        }
        
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        
        return headers
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        reraise=True
    )
    async def request(
        self,
        method: str,
        endpoint: str,
        data: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Выполняет HTTP запрос с retry'ями"""
        await self._check_rate_limit()
        
        url = f"{self.base_url}/{endpoint.lstrip('/')}"
        headers = self._prepare_headers()
        
        async with aiohttp.ClientSession(timeout=self.timeout) as session:
            async with session.request(
                method=method,
                url=url,
                json=data,
                params=params,
                headers=headers
            ) as response:
                # Проверяем статус
                if response.status >= 400:
                    error_text = await response.text()
                    raise aiohttp.ClientResponseError(
                        request_info=response.request_info,
                        history=response.history,
                        status=response.status,
                        message=f"API ошибка: {error_text}"
                    )
                
                # Для 204 No Content
                if response.status == 204:
                    return {"success": True}
                
                return await response.json()
    
    async def get(self, endpoint: str, params: Optional[Dict] = None):
        """GET запрос"""
        return await self.request("GET", endpoint, params=params)
    
    async def post(self, endpoint: str, data: Dict[str, Any]):
        """POST запрос"""
        return await self.request("POST", endpoint, data=data)
    
    async def put(self, endpoint: str, data: Dict[str, Any]):
        """PUT запрос"""
        return await self.request("PUT", endpoint, data=data)
    
    async def delete(self, endpoint: str):
        """DELETE запрос"""
        return await self.request("DELETE", endpoint)

import asyncio
from typing import Dict, Any
import json
 
class PipelineStateManager:
    """Управляет состоянием между этапами pipeline"""
    
    def __init__(self):
        self.storage: Dict[str, Any] = {}
        self.lock = asyncio.Lock()
    
    async def store(self, key: str, data: Any) -> str:
        """Сохраняет данные и возвращает ключ"""
        async with self.lock:
            self.storage[key] = data
            return key
    
    async def retrieve(self, key: str) -> Any:
        """Извлекает данные по ключу"""
        async with self.lock:
            return self.storage.get(key)
    
    async def cleanup(self, key: str):
        """Удаляет данные после использования"""
        async with self.lock:
            self.storage.pop(key, None)
 
# Глобальный менеджер состояния
state_manager = PipelineStateManager()
 
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    if name == "extract_transactions":
        # Извлекаем данные из БД
        transactions = await db.query(arguments["sql"])
        
        # Сохраняем в state manager
        state_key = f"transactions_{time.time()}"
        await state_manager.store(state_key, transactions)
        
        return [TextContent(
            type="text",
            text=json.dumps({
                "state_key": state_key,
                "count": len(transactions),
                "sample": transactions[:5]  # показываем примеры
            })
        )]
    
    elif name == "analyze_patterns":
        # Получаем данные из предыдущего этапа
        state_key = arguments["state_key"]
        transactions = await state_manager.retrieve(state_key)
        
        # Анализируем паттерны
        patterns = find_purchase_patterns(transactions)
        
        # Сохраняем результат анализа
        result_key = f"patterns_{time.time()}"
        await state_manager.store(result_key, patterns)
        
        # Чистим промежуточные данные
        await state_manager.cleanup(state_key)
        
        return [TextContent(
            type="text",
            text=json.dumps({
                "state_key": result_key,
                "patterns_found": len(patterns)
            })
        )]

# Создание ограниченного пользователя в PostgreSQL
CREATE USER mcp_readonly WITH PASSWORD 'сложный_пароль';
GRANT CONNECT ON DATABASE analytics TO mcp_readonly;
GRANT USAGE ON SCHEMA public TO mcp_readonly;
GRANT SELECT ON TABLE orders, products, customers TO mcp_readonly;
 
# Запрещаем всё остальное явно
REVOKE ALL ON ALL TABLES IN SCHEMA public FROM mcp_readonly;
REVOKE CREATE ON SCHEMA public FROM mcp_readonly;

import sqlparse
from sqlparse.sql import Statement
from sqlparse.tokens import Keyword, DML
 
def validate_query_safety(query: str) -> tuple[bool, str]:
    """Проверяет что запрос безопасен для выполнения"""
    try:
        parsed = sqlparse.parse(query)
        if not parsed:
            return False, "Пустой запрос"
        
        statement: Statement = parsed[0]
        
        # Проверяем тип запроса
        first_token = statement.token_first(skip_ws=True, skip_cm=True)
        if not first_token:
            return False, "Невалидный SQL"
        
        # Разрешены только SELECT
        if first_token.ttype != DML or first_token.value.upper() != 'SELECT':
            return False, f"Запрещённый тип запроса: {first_token.value}"
        
        # Ищем опасные ключевые слова во всём запросе
        dangerous_keywords = {
            'DROP', 'DELETE', 'TRUNCATE', 'ALTER', 
            'CREATE', 'INSERT', 'UPDATE', 'GRANT', 'REVOKE'
        }
        
        query_upper = query.upper()
        for keyword in dangerous_keywords:
            if keyword in query_upper:
                return False, f"Обнаружено запрещённое слово: {keyword}"
        
        return True, "OK"
        
    except Exception as e:
        return False, f"Ошибка парсинга: {str(e)}"
 
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    if name == "query_database":
        query = arguments["query"]
        
        # Валидируем перед выполнением
        is_safe, message = validate_query_safety(query)
        if not is_safe:
            return [TextContent(
                type="text",
                text=f"Запрос отклонён: {message}"
            )]
        
        # Выполняем только после проверки
        results = await db.execute_query(query)
        return [TextContent(type="text", text=json.dumps(results))]

from pathlib import Path
 
class SecureFileSystem:
    """Файловая система с жёсткими ограничениями"""
    
    def __init__(self, allowed_paths: list[str]):
        # Резолвим пути к абсолютным для надёжного сравнения
        self.allowed_paths = [Path(p).resolve() for p in allowed_paths]
        
        # Максимальные размеры
        self.max_file_size = 10 * 1024 * 1024  # 10MB
        self.max_files_per_request = 100
    
    def _is_path_allowed(self, path: str) -> tuple[bool, Path]:
        """Проверяет разрешён ли доступ к пути"""
        try:
            requested_path = Path(path).resolve()
        except (ValueError, OSError):
            return False, None
        
        # Проверяем что путь внутри разрешённых директорий
        for allowed in self.allowed_paths:
            try:
                # relative_to выбросит ValueError если пути не связаны
                requested_path.relative_to(allowed)
                return True, requested_path
            except ValueError:
                continue
        
        return False, None
    
    async def read_file(self, path: str) -> str:
        """Безопасное чтение файла"""
        is_allowed, resolved_path = self._is_path_allowed(path)
        
        if not is_allowed:
            raise PermissionError(
                f"Доступ к файлу запрещён: {path}"
            )
        
        # Проверяем размер
        file_size = await aiofiles.os.path.getsize(resolved_path)
        if file_size > self.max_file_size:
            raise ValueError(
                f"Файл слишком большой: {file_size} байт "
                f"(лимит {self.max_file_size})"
            )
        
        # Читаем только после всех проверок
        async with aiofiles.open(resolved_path, 'r') as f:
            return await f.read()

from pydantic import BaseModel, Field, validator, root_validator
from typing import Optional, List
from datetime import datetime
import re
 
class DatabaseQueryParams(BaseModel):
    """Параметры для запроса к базе данных"""
    
    query: str = Field(
        ...,
        min_length=10,
        max_length=10000,
        description="SQL запрос для выполнения"
    )
    
    limit: int = Field(
        default=100,
        ge=1,  # greater or equal
        le=10000,  # less or equal
        description="Максимум строк в результате"
    )
    
    timeout: Optional[int] = Field(
        default=30,
        ge=1,
        le=300,
        description="Таймаут в секундах"
    )
    
    @validator('query')
    def validate_query_format(cls, v):
        """Проверяем формат SQL запроса"""
        # Удаляем лишние пробелы
        v = ' '.join(v.split())
        
        # Проверяем что начинается с SELECT
        if not v.strip().upper().startswith('SELECT'):
            raise ValueError("Разрешены только SELECT запросы")
        
        # Проверяем на опасные конструкции
        dangerous_patterns = [
            r';\s*DROP',
            r';\s*DELETE',
            r';\s*TRUNCATE',
            r'--',  # SQL комментарии
            r'/\*.*\*/',  # блочные комментарии
        ]
        
        for pattern in dangerous_patterns:
            if re.search(pattern, v, re.IGNORECASE):
                raise ValueError(
                    f"Обнаружена потенциально опасная конструкция"
                )
        
        return v
    
    @root_validator
    def validate_timeout_for_large_limits(cls, values):
        """Для больших лимитов требуем больший таймаут"""
        limit = values.get('limit')
        timeout = values.get('timeout')
        
        if limit > 1000 and timeout < 60:
            raise ValueError(
                "Для запросов больше 1000 строк таймаут должен быть минимум 60 сек"
            )
        
        return values
 
class FileOperationParams(BaseModel):
    """Параметры для операций с файлами"""
    
    path: str = Field(
        ...,
        min_length=1,
        max_length=500,
        description="Путь к файлу относительно базовой директории"
    )
    
    content: Optional[str] = Field(
        default=None,
        max_length=5_000_000,  # 5MB текста
        description="Содержимое для записи"
    )
    
    encoding: str = Field(
        default="utf-8",
        description="Кодировка файла"
    )
    
    @validator('path')
    def validate_path_safety(cls, v):
        """Проверяем безопасность пути"""
        # Запрещаем абсолютные пути
        if v.startswith('/') or ':' in v:
            raise ValueError("Абсолютные пути запрещены")
        
        # Запрещаем path traversal
        if '..' in v:
            raise ValueError("Путь не может содержать '..'")
        
        # Запрещаем скрытые файлы
        if any(part.startswith('.') for part in v.split('/')):
            raise ValueError("Доступ к скрытым файлам запрещён")
        
        return v.strip()
    
    @validator('encoding')
    def validate_encoding(cls, v):
        """Проверяем что кодировка поддерживается"""
        import codecs
        try:
            codecs.lookup(v)
        except LookupError:
            raise ValueError(f"Неподдерживаемая кодировка: {v}")
        return v

@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    try:
        if name == "query_database":
            # Валидация через Pydantic
            params = DatabaseQueryParams(**arguments)
            
            # Если дошли сюда - все проверки пройдены
            results = await db.execute_query(
                query=params.query,
                limit=params.limit,
                timeout=params.timeout
            )
            
            return [TextContent(
                type="text",
                text=json.dumps(results, ensure_ascii=False)
            )]
        
        elif name == "write_file":
            params = FileOperationParams(**arguments)
            
            await fs.write_file(
                path=params.path,
                content=params.content,
                encoding=params.encoding
            )
            
            return [TextContent(
                type="text",
                text=f"Файл записан: {params.path}"
            )]
    
    except ValidationError as e:
        # Pydantic возвращает структурированные ошибки
        errors = e.errors()
        error_messages = [
            f"{err['loc'][0]}: {err['msg']}"
            for err in errors
        ]
        
        return [TextContent(
            type="text",
            text=f"Ошибка валидации параметров:\n" + "\n".join(error_messages)
        )]

import unicodedata
 
def sanitize_text(text: str) -> str:
    """Очищает текст от проблемных символов"""
    # Нормализуем Unicode
    text = unicodedata.normalize('NFKC', text)
    
    # Убираем управляющие символы кроме переносов
    text = ''.join(
        char for char in text
        if unicodedata.category(char)[0] != 'C' or char in '\n\r\t'
    )
    
    # Приводим переносы к единому формату
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    
    # Убираем trailing whitespace с каждой строки
    lines = [line.rstrip() for line in text.split('\n')]
    
    return '\n'.join(lines)
 
class TextContentParams(BaseModel):
    content: str
    
    @validator('content', pre=True)
    def sanitize_content(cls, v):
        if isinstance(v, str):
            return sanitize_text(v)
        return v

from enum import Enum
from typing import Set, Dict, Optional
import jwt
from datetime import datetime, timedelta
 
class Permission(Enum):
    """Доступные разрешения"""
    READ_DATABASE = "db:read"
    WRITE_DATABASE = "db:write"
    READ_FILES = "files:read"
    WRITE_FILES = "files:write"
    EXECUTE_COMMANDS = "commands:execute"
    MANAGE_USERS = "users:manage"
    VIEW_LOGS = "logs:view"
 
class Role:
    """Роль с набором разрешений"""
    def __init__(self, name: str, permissions: Set[Permission]):
        self.name = name
        self.permissions = permissions
    
    def has_permission(self, permission: Permission) -> bool:
        return permission in self.permissions
 
# Предопределённые роли
ROLES = {
    "viewer": Role("viewer", {
        Permission.READ_DATABASE,
        Permission.READ_FILES,
        Permission.VIEW_LOGS
    }),
    "developer": Role("developer", {
        Permission.READ_DATABASE,
        Permission.WRITE_DATABASE,
        Permission.READ_FILES,
        Permission.WRITE_FILES,
        Permission.EXECUTE_COMMANDS,
        Permission.VIEW_LOGS
    }),
    "admin": Role("admin", set(Permission))  # все разрешения
}
 
class User:
    """Пользователь системы"""
    def __init__(
        self,
        user_id: str,
        username: str,
        roles: Set[str]
    ):
        self.user_id = user_id
        self.username = username
        self.roles = roles
    
    def get_permissions(self) -> Set[Permission]:
        """Собирает все разрешения из всех ролей"""
        permissions = set()
        for role_name in self.roles:
            role = ROLES.get(role_name)
            if role:
                permissions.update(role.permissions)
        return permissions
    
    def has_permission(self, permission: Permission) -> bool:
        """Проверяет наличие конкретного разрешения"""
        return permission in self.get_permissions()

class AuthManager:
    """Управление аутентификацией и авторизацией"""
    
    def __init__(self, secret_key: str):
        self.secret_key = secret_key
        self.users: Dict[str, User] = {}
    
    def create_token(self, user: User, expires_in: int = 3600) -> str:
        """Создаёт JWT токен для пользователя"""
        payload = {
            "user_id": user.user_id,
            "username": user.username,
            "roles": list(user.roles),
            "exp": datetime.utcnow() + timedelta(seconds=expires_in),
            "iat": datetime.utcnow()
        }
        
        return jwt.encode(payload, self.secret_key, algorithm="HS256")
    
    def verify_token(self, token: str) -> Optional[User]:
        """Проверяет токен и возвращает пользователя"""
        try:
            payload = jwt.decode(
                token, 
                self.secret_key, 
                algorithms=["HS256"]
            )
            
            return User(
                user_id=payload["user_id"],
                username=payload["username"],
                roles=set(payload["roles"])
            )
        
        except jwt.ExpiredSignatureError:
            return None
        except jwt.InvalidTokenError:
            return None
    
    def check_permission(
        self, 
        user: User, 
        required_permission: Permission
    ) -> tuple[bool, str]:
        """Проверяет наличие разрешения у пользователя"""
        if user.has_permission(required_permission):
            return True, "OK"
        
        return False, f"Недостаточно прав: требуется {required_permission.value}"

from functools import wraps
from typing import Callable
 
auth_manager = AuthManager(secret_key="супер-секретный-ключ-храните-в-env")
 
def require_permission(permission: Permission):
    """Декоратор для проверки прав доступа"""
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Извлекаем токен из контекста (передаётся клиентом)
            token = kwargs.get("_auth_token")
            if not token:
                raise PermissionError("Требуется аутентификация")
            
            # Проверяем токен
            user = auth_manager.verify_token(token)
            if not user:
                raise PermissionError("Невалидный токен")
            
            # Проверяем права
            has_access, message = auth_manager.check_permission(
                user, 
                permission
            )
            if not has_access:
                raise PermissionError(message)
            
            # Добавляем пользователя в kwargs для использования в функции
            kwargs["_current_user"] = user
            
            return await func(*args, **kwargs)
        
        return wrapper
    return decorator
 
@require_permission(Permission.WRITE_DATABASE)
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict, **kwargs):
    """Обработчик с проверкой прав"""
    current_user = kwargs.get("_current_user")
    
    if name == "modify_data":
        # Логируем кто выполнил операцию
        logger.info(
            f"Пользователь {current_user.username} модифицирует данные"
        )
        
        # Выполняем операцию
        result = await db.execute_command(arguments["query"])
        return [TextContent(type="text", text=result)]

import asyncio
import time
from typing import Optional
from dataclasses import dataclass
 
@dataclass
class RateLimitConfig:
  """Конфигурация rate limiter'а"""
  max_requests: int  # макс запросов
  time_window: float  # временное окно в секундах
  burst_size: Optional[int] = None  # разрешённый burst
 
class TokenBucketRateLimiter:
  """Token bucket алгоритм для rate limiting"""
  
  def __init__(self, config: RateLimitConfig):
      self.max_tokens = config.max_requests
      self.refill_rate = config.max_requests / config.time_window
      self.burst_size = config.burst_size or config.max_requests
      
      self.tokens = float(self.burst_size)
      self.last_update = time.time()
      self.lock = asyncio.Lock()
  
  async def acquire(self, tokens_needed: int = 1) -> bool:
      """Пытается получить токены"""
      async with self.lock:
          now = time.time()
          elapsed = now - self.last_update
          
          # Добавляем токены пропорционально времени
          self.tokens = min(
              self.burst_size,
              self.tokens + (elapsed * self.refill_rate)
          )
          self.last_update = now
          
          if self.tokens >= tokens_needed:
              self.tokens -= tokens_needed
              return True
          
          return False
  
  async def wait_and_acquire(self, tokens_needed: int = 1) -> None:
      """Ждёт пока появятся токены"""
      while True:
          if await self.acquire(tokens_needed):
              return
          
          # Вычисляем сколько ждать
          async with self.lock:
              tokens_deficit = tokens_needed - self.tokens
              wait_time = tokens_deficit / self.refill_rate
          
          await asyncio.sleep(min(wait_time, 1.0))

class RateLimitMiddleware:
  """Middleware для rate limiting запросов"""
  
  def __init__(self):
      # Разные лимиты для разных операций
      self.limiters = {
          "query_database": TokenBucketRateLimiter(
              RateLimitConfig(max_requests=50, time_window=60)
          ),
          "write_file": TokenBucketRateLimiter(
              RateLimitConfig(max_requests=20, time_window=60)
          ),
          "default": TokenBucketRateLimiter(
              RateLimitConfig(max_requests=100, time_window=60)
          )
      }
  
  def rate_limit(self, tool_name: str):
      """Декоратор для rate limiting"""
      def decorator(func):
          @wraps(func)
          async def wrapper(*args, **kwargs):
              limiter = self.limiters.get(
                  tool_name, 
                  self.limiters["default"]
              )
              
              # Пробуем получить токен
              if not await limiter.acquire():
                  raise Exception(
                      f"Rate limit exceeded для {tool_name}. "
                      "Повторите запрос через минуту"
                  )
              
              return await func(*args, **kwargs)
          return wrapper
      return decorator
 
rate_limiter = RateLimitMiddleware()
 
@rate_limiter.rate_limit("query_database")
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
  if name == "query_database":
      # Выполняем только если прошли rate limit
      results = await db.execute_query(arguments["query"])
      return [TextContent(type="text", text=json.dumps(results))]

import psutil
 
class AdaptiveRateLimiter:
  """Rate limiter который адаптируется под нагрузку"""
  
  def __init__(self, base_config: RateLimitConfig):
      self.base_limiter = TokenBucketRateLimiter(base_config)
      self.cpu_threshold = 80.0  # процент CPU
      self.memory_threshold = 85.0  # процент памяти
  
  async def acquire(self) -> bool:
      """Проверяет rate limit с учётом нагрузки"""
      # Проверяем ресурсы системы
      cpu_usage = psutil.cpu_percent(interval=0.1)
      mem_usage = psutil.virtual_memory().percent
      
      # Под высокой нагрузкой строже лимиты
      if cpu_usage > self.cpu_threshold or mem_usage > self.memory_threshold:
          # Требуем больше токенов = меньше пропускаем
          return await self.base_limiter.acquire(tokens_needed=3)
      
      return await self.base_limiter.acquire(tokens_needed=1)

from aiohttp import web
import ssl
 
def create_ssl_context(cert_file: str, key_file: str) -> ssl.SSLContext:
    """Создаёт SSL контекст для HTTPS"""
    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    context.load_cert_chain(cert_file, key_file)
    
    # Отключаем устаревшие протоколы
    context.minimum_version = ssl.TLSVersion.TLSv1_2
    
    # Только безопасные cipher suites
    context.set_ciphers(
        'ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20:!aNULL:!MD5:!DSS'
    )
    
    return context
 
async def run_server():
    """Запускаем MCP-сервер с HTTPS"""
    app = web.Application()
    # Настройка роутов и middleware
    
    ssl_context = create_ssl_context(
        cert_file="/path/to/cert.pem",
        key_file="/path/to/key.pem"
    )
    
    runner = web.AppRunner(app)
    await runner.setup()
    
    site = web.TCPSite(
        runner, 
        host='0.0.0.0', 
        port=443,
        ssl_context=ssl_context
    )
    
    await site.start()

from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
import os
import base64
 
class PayloadEncryption:
    """Шифрование данных на уровне приложения"""
    
    def __init__(self):
        # Генерируем RSA ключи для сервера
        self.private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048
        )
        self.public_key = self.private_key.public_key()
    
    def encrypt_data(self, data: str, client_public_key: rsa.RSAPublicKey) -> dict:
        """Шифрует данные для отправки клиенту"""
        # Генерируем случайный AES ключ
        aes_key = os.urandom(32)  # 256 бит
        iv = os.urandom(16)
        
        # Шифруем данные AES
        cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        
        # Паддинг до кратности 16 байт
        padded_data = data.encode() + b'\0' * (16 - len(data) % 16)
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        
        # Шифруем AES ключ RSA ключом клиента
        encrypted_key = client_public_key.encrypt(
            aes_key,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        
        return {
            "encrypted_data": base64.b64encode(encrypted_data).decode(),
            "encrypted_key": base64.b64encode(encrypted_key).decode(),
            "iv": base64.b64encode(iv).decode()
        }

import asyncio
from typing import TypeVar, Awaitable
 
T = TypeVar('T')
 
async def with_timeout(
    coro: Awaitable[T],
    timeout: float,
    operation_name: str = "операция"
) -> T:
    """Выполняет корутину с таймаутом и понятной ошибкой"""
    try:
        return await asyncio.wait_for(coro, timeout=timeout)
    except asyncio.TimeoutError:
        raise TimeoutError(
            f"Таймаут {operation_name}: превышен лимит {timeout}с"
        )
 
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    if name == "query_database":
        query = arguments["query"]
        
        try:
            # Даём запросу максимум 30 секунд
            results = await with_timeout(
                db.execute_query(query),
                timeout=30.0,
                operation_name=f"выполнения запроса"
            )
            
            return [TextContent(
                type="text",
                text=json.dumps(results)
            )]
            
        except TimeoutError as e:
            return [TextContent(
                type="text",
                text=f"Запрос выполняется слишком долго: {str(e)}"
            )]

import asyncpg
from aiohttp import ClientError, ClientTimeout
 
async def safe_tool_execution(tool_func, *args, **kwargs):
    """Безопасное выполнение инструмента с обработкой всех типов ошибок"""
    try:
        return await tool_func(*args, **kwargs)
        
    except asyncio.TimeoutError:
        return [TextContent(
            type="text",
            text="Операция превысила допустимое время выполнения. "
                 "Попробуйте упростить запрос или повторите позже."
        )]
    
    except asyncpg.PostgresSyntaxError as e:
        return [TextContent(
            type="text",
            text=f"Синтаксическая ошибка в SQL-запросе: {str(e)}"
        )]
    
    except asyncpg.PostgresError as e:
        return [TextContent(
            type="text",
            text=f"Ошибка базы данных: {str(e)}"
        )]
    
    except ClientTimeout:
        return [TextContent(
            type="text",
            text="Внешний API не ответил вовремя. "
                 "Сервис может быть перегружен."
        )]
    
    except ClientError as e:
        return [TextContent(
            type="text",
            text=f"Ошибка сетевого запроса: {str(e)}"
        )]
    
    except PermissionError as e:
        return [TextContent(
            type="text",
            text=f"Недостаточно прав: {str(e)}"
        )]
    
    except FileNotFoundError as e:
        return [TextContent(
            type="text",
            text=f"Файл не найден: {str(e)}"
        )]
    
    except ValueError as e:
        return [TextContent(
            type="text",
            text=f"Некорректные параметры: {str(e)}"
        )]
    
    except Exception as e:
        # Логируем неожиданные ошибки для отладки
        logger.exception(f"Неожиданная ошибка в {tool_func.__name__}")
        return [TextContent(
            type="text",
            text=f"Внутренняя ошибка сервера. "
                 f"Обратитесь к администратору. ID: {uuid.uuid4()}"
        )]

import pytest
from unittest.mock import AsyncMock, patch, MagicMock
import json
 
from server import (
    validate_query_safety,
    sanitize_text,
    handle_call_tool
)
 
class TestQueryValidation:
    """Тесты валидации SQL-запросов"""
    
    def test_valid_select_query(self):
        """Простой SELECT должен проходить"""
        query = "SELECT * FROM users WHERE id = 1"
        is_safe, message = validate_query_safety(query)
        
        assert is_safe
        assert message == "OK"
    
    def test_query_with_joins(self):
        """SELECT с JOIN разрешён"""
        query = """
            SELECT u.name, o.total 
            FROM users u 
            JOIN orders o ON u.id = o.user_id
        """
        is_safe, _ = validate_query_safety(query)
        assert is_safe
    
    def test_rejects_delete(self):
        """DELETE запросы блокируются"""
        query = "DELETE FROM users WHERE id = 1"
        is_safe, message = validate_query_safety(query)
        
        assert not is_safe
        assert "DELETE" in message
    
    def test_rejects_drop_table(self):
        """DROP TABLE категорически запрещён"""
        query = "DROP TABLE users"
        is_safe, message = validate_query_safety(query)
        
        assert not is_safe
        assert "DROP" in message
    
    def test_rejects_injection_attempt(self):
        """SQL-инъекция через комментарии"""
        query = "SELECT * FROM users -- ; DROP TABLE users"
        is_safe, _ = validate_query_safety(query)
        assert not is_safe

@pytest.mark.asyncio
class TestDatabaseAdapter:
    """Тесты адаптера БД с моками"""
    
    async def test_execute_query_success(self):
        """Успешное выполнение запроса"""
        mock_connection = AsyncMock()
        mock_connection.fetch.return_value = [
            {"id": 1, "name": "Иван"},
            {"id": 2, "name": "Мария"}
        ]
        
        adapter = PostgresAdapter(...)
        adapter.pool.acquire = AsyncMock(return_value=mock_connection)
        
        results = await adapter.execute_query("SELECT * FROM users")
        
        assert len(results) == 2
        assert results[0]["name"] == "Иван"
        mock_connection.fetch.assert_called_once()
    
    async def test_execute_query_timeout(self):
        """Обработка таймаута"""
        mock_connection = AsyncMock()
        mock_connection.fetch.side_effect = asyncio.TimeoutError()
        
        adapter = PostgresAdapter(...)
        adapter.pool.acquire = AsyncMock(return_value=mock_connection)
        
        with pytest.raises(asyncio.TimeoutError):
            await adapter.execute_query("SELECT * FROM huge_table")
    
    async def test_connection_pool_exhausted(self):
        """Поведение при исчерпании пула соединений"""
        adapter = PostgresAdapter(...)
        
        # Мокируем пул который не может дать соединение
        adapter.pool.acquire = AsyncMock(
            side_effect=asyncpg.TooManyConnectionsError("Pool exhausted")
        )
        
        with pytest.raises(asyncpg.TooManyConnectionsError):
            await adapter.execute_query("SELECT 1")

@pytest.mark.parametrize("input_text,expected_output", [
    ("  hello  ", "hello"),
    ("line1\r\nline2", "line1\nline2"),
    ("тест\u200bс\u200bневидимыми", "тестсневидимыми"),
    ("", ""),
    ("   \t\n   ", ""),
])
def test_sanitize_text(input_text, expected_output):
    """Тестируем очистку текста на разных входных данных"""
    result = sanitize_text(input_text)
    assert result == expected_output

@pytest.fixture
async def test_database():
    """Создаёт временную БД для тестов"""
    # Setup
    db_name = f"test_db_{uuid.uuid4().hex[:8]}"
    await create_database(db_name)
    
    adapter = PostgresAdapter(database=db_name, ...)
    await adapter.initialize()
    
    # Заполняем тестовыми данными
    await adapter.execute_command("""
        CREATE TABLE users (id SERIAL PRIMARY KEY, name TEXT);
        INSERT INTO users (name) VALUES ('Тестовый пользователь');
    """)
    
    yield adapter  # отдаём тесту
    
    # Teardown
    await adapter.close()
    await drop_database(db_name)
 
@pytest.mark.asyncio
async def test_user_query(test_database):
    """Тест использует фикстуру"""
    results = await test_database.execute_query(
        "SELECT * FROM users"
    )
    assert len(results) == 1

import pytest
import asyncio
from mcp.client import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from contextlib import AsyncExitStack
 
@pytest.fixture
async def mcp_session():
    """Фикстура для запуска сервера и подключения клиента"""
    exit_stack = AsyncExitStack()
    
    # Параметры для запуска сервера через stdio
    server_params = StdioServerParameters(
        command="python",
        args=["path/to/server.py"],
        env=None
    )
    
    # Запускаем сервер и подключаемся
    stdio_transport = await exit_stack.enter_async_context(
        stdio_client(server_params)
    )
    read_stream, write_stream = stdio_transport
    
    # Создаём клиентскую сессию
    session = await exit_stack.enter_async_context(
        ClientSession(read_stream, write_stream)
    )
    
    # Инициализируем соединение
    await session.initialize()
    
    yield session
    
    # Cleanup происходит автоматически через exit_stack
    await exit_stack.aclose()
 
@pytest.mark.asyncio
async def test_list_tools_integration(mcp_session):
    """Тест получения списка инструментов"""
    response = await mcp_session.list_tools()
    
    assert response.tools is not None
    assert len(response.tools) > 0
    
    # Проверяем что каждый инструмент имеет нужные поля
    for tool in response.tools:
        assert tool.name
        assert tool.description
        assert tool.inputSchema
        assert tool.inputSchema.get("type") == "object"
 
@pytest.mark.asyncio  
async def test_tool_call_integration(mcp_session):
    """Тест вызова инструмента end-to-end"""
    # Вызываем реальный инструмент
    result = await mcp_session.call_tool(
        name="get_weather",
        arguments={"city": "Moscow", "units": "celsius"}
    )
    
    assert result.content is not None
    assert len(result.content) > 0
    
    # Проверяем структуру ответа
    content = result.content[0]
    assert content.type == "text"
    assert "Moscow" in content.text
    assert "°C" in content.text or "celsius" in content.text.lower()
 
@pytest.mark.asyncio
async def test_concurrent_requests(mcp_session):
    """Проверяем что сервер обрабатывает параллельные запросы"""
    # Запускаем 10 запросов одновременно
    tasks = [
        mcp_session.call_tool(
            name="get_weather",
            arguments={"city": f"City{i}"}
        )
        for i in range(10)
    ]
    
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # Все запросы должны завершиться успешно
    for result in results:
        assert not isinstance(result, Exception)
        assert result.content
 
@pytest.mark.asyncio
async def test_error_handling_integration(mcp_session):
    """Тест обработки ошибок через весь стек"""
    # Вызываем инструмент с невалидными параметрами
    result = await mcp_session.call_tool(
        name="query_database",
        arguments={"query": "DELETE FROM users"}  # запрещённый запрос
    )
    
    # Сервер должен вернуть читаемое сообщение об ошибке
    assert result.content
    content_text = result.content[0].text
    assert "запрещ" in content_text.lower() or "DELETE" in content_text

from aiohttp import web
from aiohttp.test_utils import TestClient
 
@pytest.fixture
async def http_mcp_client(aiohttp_client):
    """Фикстура для HTTP клиента"""
    app = create_mcp_server_app()  # ваша функция создания app
    client = await aiohttp_client(app)
    return client
 
@pytest.mark.asyncio
async def test_http_tool_call(http_mcp_client):
    """Вызов инструмента через HTTP"""
    response = await http_mcp_client.post(
        '/mcp',
        json={
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "get_weather",
                "arguments": {"city": "London"}
            }
        }
    )
    
    assert response.status == 200
    data = await response.json()
    
    assert data["jsonrpc"] == "2.0"
    assert data["id"] == 1
    assert "result" in data

import logging
import json
from typing import Any
import time
 
logger = logging.getLogger("mcp.transport")
 
class TransportDebugger:
    """Отладочный слой для транспорта"""
    
    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self.message_counter = 0
    
    def log_outgoing(self, message: dict):
        """Логирует исходящее сообщение"""
        if not self.enabled:
            return
        
        self.message_counter += 1
        logger.debug(
            f"[OUT #{self.message_counter}] {json.dumps(message, ensure_ascii=False)[:500]}"
        )
    
    def log_incoming(self, message: dict):
        """Логирует входящее сообщение"""
        if not self.enabled:
            return
        
        self.message_counter += 1
        logger.debug(
            f"[IN #{self.message_counter}] {json.dumps(message, ensure_ascii=False)[:500]}"
        )
    
    def log_error(self, error: Exception, context: str):
        """Логирует ошибки транспорта"""
        logger.error(
            f"Транспортная ошибка в {context}: {type(error).__name__}: {str(error)}",
            exc_info=True
        )
 
debugger = TransportDebugger()
 
async def send_message(writer, message: dict):
    """Отправка с логированием"""
    debugger.log_outgoing(message)
    
    try:
        data = json.dumps(message).encode('utf-8')
        writer.write(data + b'\n')
        await writer.drain()
    except Exception as e:
        debugger.log_error(e, "отправка сообщения")
        raise

import asyncio
 
async def read_message(reader: asyncio.StreamReader) -> dict:
    """Читает одно JSON-RPC сообщение из потока"""
    try:
        line = await asyncio.wait_for(
            reader.readline(),
            timeout=30.0
        )
        
        if not line:
            raise EOFError("Поток закрыт")
        
        message = json.loads(line.decode('utf-8'))
        debugger.log_incoming(message)
        return message
        
    except asyncio.TimeoutError:
        logger.error("Таймаут чтения сообщения")
        raise
    except json.JSONDecodeError as e:
        logger.error(f"Невалидный JSON: {line[:200]}")
        debugger.log_error(e, "парсинг JSON")
        raise

import cProfile
import pstats
from io import StringIO
 
def profile_tool_execution(func):
    """Декоратор для профилирования инструмента"""
    async def wrapper(*args, **kwargs):
        profiler = cProfile.Profile()
        profiler.enable()
        
        try:
            result = await func(*args, **kwargs)
            return result
        finally:
            profiler.disable()
            
            # Выводим топ-20 самых медленных вызовов
            stream = StringIO()
            stats = pstats.Stats(profiler, stream=stream)
            stats.strip_dirs()
            stats.sort_stats('cumulative')
            stats.print_stats(20)
            
            logger.info(f"Профиль {func.__name__}:
{stream.getvalue()}")
    
    return wrapper
 
@profile_tool_execution
async def handle_expensive_tool(name: str, arguments: dict):
    """Инструмент под профайлером"""
    results = await db.execute_query(arguments["query"])
    processed = process_results(results)  # тут может быть узкое место
    return [TextContent(type="text", text=json.dumps(processed))]

# Профилирование живого процесса
py-spy record -o profile.svg --pid 12345 --duration 60
 
# Топ функций по CPU
py-spy top --pid 12345

import tracemalloc
import linecache
 
tracemalloc.start()
 
# Делаем снимок до операции
snapshot_before = tracemalloc.take_snapshot()
 
# Выполняем подозрительную операцию
await process_large_dataset()
 
# Снимок после
snapshot_after = tracemalloc.take_snapshot()
 
# Сравниваем
top_stats = snapshot_after.compare_to(snapshot_before, 'lineno')
 
print("Топ-10 мест с наибольшим приростом памяти:")
for stat in top_stats[:10]:
    print(f"{stat.size_diff / 1024 / 1024:.2f} MB - {stat.traceback}")

import asyncio
import warnings
 
# Включаем debug режим
asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
loop = asyncio.new_event_loop()
loop.set_debug(True)
 
# Предупреждения о медленных операциях (>100ms)
warnings.simplefilter('always', ResourceWarning)
loop.slow_callback_duration = 0.1
 
asyncio.set_event_loop(loop)

async def profile_query(query: str):
    """Анализ производительности SQL"""
    explained = f"EXPLAIN ANALYZE {query}"
    
    async with db.acquire() as conn:
        result = await conn.fetch(explained)
        
        for row in result:
            logger.info(row['QUERY PLAN'])

import structlog
import uuid
from datetime import datetime
 
# Настройка structlog
structlog.configure(
    processors=[
        structlog.stdlib.add_log_level,
        structlog.stdlib.add_logger_name,
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
        structlog.processors.JSONRenderer()
    ],
    context_class=dict,
    logger_factory=structlog.stdlib.LoggerFactory(),
    cache_logger_on_first_use=True,
)
 
logger = structlog.get_logger()
 
class RequestTracer:
    """Трассировка запросов через систему"""
    
    def __init__(self):
        self.active_requests = {}
    
    async def trace_request(self, func):
        """Декоратор для трассировки"""
        async def wrapper(*args, **kwargs):
            request_id = str(uuid.uuid4())[:8]
            
            # Контекст для всех логов этого запроса
            log = logger.bind(
                request_id=request_id,
                tool_name=func.__name__,
                timestamp_start=datetime.utcnow().isoformat()
            )
            
            log.info("запрос начат", arguments=str(kwargs)[:200])
            
            try:
                result = await func(*args, **kwargs)
                
                log.info(
                    "запрос завершён успешно",
                    result_size=len(str(result))
                )
                
                return result
                
            except Exception as e:
                log.error(
                    "запрос завершён с ошибкой",
                    error_type=type(e).__name__,
                    error_message=str(e)
                )
                raise
        
        return wrapper
 
tracer = RequestTracer()
 
@tracer.trace_request
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    log = logger.bind(tool=name)
    
    log.debug("валидация параметров", params=arguments)
    
    if name == "query_database":
        log.info("выполнение SQL", query=arguments["query"][:100])
        
        results = await db.execute_query(arguments["query"])
        
        log.info(
            "результаты получены",
            rows_count=len(results),
            execution_time_ms=42
        )
        
        return [TextContent(type="text", text=json.dumps(results))]