task-69f8e929da860fb4533faa2a/memory_server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Memory‑Server – MCP‑сервер для хранения ключ‑значение с поддержкой namespaces.

Запуск:
    python memory_server.py
"""

from __future__ import annotations

import json
import fnmatch
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastmcp import FastMCP


class MemoryServer:
    """
    Сервер, реализующий инструменты «save», «get» и т.д. через протокол MCP.
    Данные хранятся в JSON‑файле ``memory_data.json`` в формате:

        {
            "namespace:key": {"value": ..., "timestamp": "..."}
        }
    """

    def __init__(self, storage_path: str | Path = "./memory_data.json") -> None:
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path(storage_path)

        # При инициализации загружаем существующую память
        self._memory: Dict[str, Dict[str, Any]] = self._load_memory()

    # ------------------------------------------------------------------
    # Вспомогательные методы для работы с файлом
    # ------------------------------------------------------------------
    def _load_memory(self) -> Dict[str, Dict[str, Any]]:
        """Загружает память из JSON‑файла. Если файл не существует – возвращаем пустой dict."""
        if not self.storage_path.exists():
            return {}
        try:
            with self.storage_path.open("r", encoding="utf-8") as f:
                data = json.load(f)
            # Приводим к типу Dict[str, Dict[str, Any]]
            return {k: v for k, v in data.items()}
        except Exception as exc:
            print(f"[MemoryServer] Ошибка чтения файла {self.storage_path}: {exc}")
            return {}

    def _save_memory(self) -> None:
        """Сохраняет текущий словарь памяти в JSON‑файл."""
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        with self.storage_path.open("w", encoding="utf-8") as f:
            json.dump(self._memory, f, indent=2, ensure_ascii=False)

    # ------------------------------------------------------------------
    # 1️⃣ Базовые инструменты
    # ------------------------------------------------------------------
    @self.mcp.tool()
    def save(key: str, value: Any) -> bool:
        """
        Сохраняет значение по ключу в память сервера.

        Args:
            key: Идентификатор для сохранения (уникальный ключ).
            value: Любое сериализуемое значение.

        Returns:
            True при успешном сохранении, False иначе.
        """
        if not isinstance(key, str) or "/" in key or ".." in key:
            return False

        server = MemoryServer()
        timestamp = datetime.now(timezone.utc).isoformat()
        server._memory[key] = {"value": value, "timestamp": timestamp}
        server._save_memory()
        return True

    @self.mcp.tool()
    def get(key: str) -> Optional[Dict[str, Any]]:
        """
        Возвращает значение по ключу с метаданными.

        Args:
            key: Идентификатор для поиска.

        Returns:
            Словарь с полями {"key": ..., "value": ..., "timestamp": ...}
            или None если ключ не найден.
        """
        server = MemoryServer()
        entry = server._memory.get(key)
        if entry is None:
            return None
        return {
            "key": key,
            "value": entry["value"],
            "timestamp": entry["timestamp"],
        }

    @self.mcp.tool()
    def delete(key: str) -> bool:
        """
        Удаляет ключ из памяти сервера.

        Args:
            key: Идентификатор для удаления.

        Returns:
            True при успешном удалении, False если ключ не найден.
        """
        server = MemoryServer()
        if key in server._memory:
            del server._memory[key]
            server._save_memory()
            return True
        return False

    @self.mcp.tool()
    def list_keys(pattern: str = "*") -> List[str]:
        """
        Возвращает список всех ключей с поддержкой wildcard‑паттерна.

        Args:
            pattern: Паттерн для фильтрации (поддерживает * и ?).

        Returns:
            Список совпадающих ключей.
        """
        server = MemoryServer()
        return [k for k in server._memory if fnmatch.fnmatch(k, pattern)]

    # ------------------------------------------------------------------
    # 2️⃣ Инструменты с поддержкой namespaces
    # ------------------------------------------------------------------
    @self.mcp.tool()
    def save_with_namespace(
        key: str,
        value: Any,
        namespace: str = "default",
    ) -> bool:
        """
        Сохраняет значение с указанием пространства имён.

        Args:
            key: Идентификатор.
            value: Значение для сохранения.
            namespace: Пространство имён (по умолчанию 'default').

        Returns:
            True при успехе, False иначе.
        """
        if not isinstance(key, str) or "/" in key or ".." in key:
            return False
        full_key = f"{namespace}:{key}"
        server = MemoryServer()
        timestamp = datetime.now(timezone.utc).isoformat()
        server._memory[full_key] = {"value": value, "timestamp": timestamp}
        server._save_memory()
        return True

    @self.mcp.tool()
    def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]:
        """
        Возвращает все ключи из указанного namespace.

        Args:
            namespace: Пространство имён для чтения.

        Returns:
            Список словарей с метаданными всех ключей namespace.
        """
        server = MemoryServer()
        prefix = f"{namespace}:"
        result: List[Dict[str, Any]] = []
        for k, v in server._memory.items():
            if k.startswith(prefix):
                result.append(
                    {
                        "key": k,
                        "value": v["value"],
                        "timestamp": v["timestamp"],
                    }
                )
        return result


# ------------------------------------------------------------------
# 3️⃣ Запуск сервера
# ------------------------------------------------------------------
if __name__ == "__main__":
    server = MemoryServer()
    # Стандартный транспорт – stdio (stdin/stdout)
    server.mcp.run(
        transport="stdio",
        show_banner=False,
        log_level="ERROR",  # Минимум логов, чтобы не засорять консоль
    )