task-69f8e929da860fb4533faa2a/memory_server.py

"""MCP-сервер для управления памятью агента (FastMCP + JSON-хранилище)."""
from __future__ import annotations

import fnmatch
import json
import re
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

from fastmcp import FastMCP

_UNSAFE_KEY = re.compile(r"\.\.|/|\\")


class MemoryServer:
    def __init__(self) -> None:
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path("./memory_data.json")
        self._register_tools()

    def _validate_key(self, key: str) -> bool:
        return bool(key) and _UNSAFE_KEY.search(key) is None

    def _load_memory(self) -> dict:
        if not self.storage_path.exists():
            return {}
        with open(self.storage_path, encoding="utf-8") as f:
            return json.load(f)

    def _save_memory(self, data: dict) -> None:
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.storage_path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _timestamp(self) -> str:
        return datetime.now().isoformat(timespec="seconds")

    def _make_entry(self, value: Any) -> dict:
        return {"value": value, "timestamp": self._timestamp()}

    def _register_tools(self) -> None:
        mcp = self.mcp

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

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

            Returns:
                True при успешном сохранении, False иначе.
            """
            if not self._validate_key(key):
                return False
            try:
                data = self._load_memory()
                data[key] = self._make_entry(value)
                self._save_memory(data)
                return True
            except (TypeError, ValueError, OSError):
                return False

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

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

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

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

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

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

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

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

            Returns:
                Список совпадающих ключей.
            """
            keys = list(self._load_memory().keys())
            return sorted(k for k in keys if fnmatch.fnmatch(k, pattern))

        @mcp.tool()
        def save_with_namespace(
            key: str, value: Any, namespace: str = "default"
        ) -> bool:
            """Сохраняет значение с указанием пространства имён.

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

            Returns:
                True при успехе, False иначе.
            """
            storage_key = f"{namespace}:{key}"
            return save(storage_key, value)

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

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

            Returns:
                Список словарей с метаданными всех ключей namespace.
            """
            prefix = f"{namespace}:"
            data = self._load_memory()
            result: list[dict] = []
            for storage_key, entry in data.items():
                if not storage_key.startswith(prefix):
                    continue
                short_key = storage_key[len(prefix) :]
                result.append(
                    {
                        "key": short_key,
                        "value": entry["value"],
                        "timestamp": entry["timestamp"],
                    }
                )
            return sorted(result, key=lambda item: item["key"])

        _ = (save, get, delete, list_keys, save_with_namespace, get_by_namespace)


if __name__ == "__main__":
    server = MemoryServer()
    server.mcp.run(
        transport="stdio",
        show_banner=False,
        log_level="ERROR",
    )