task-69f8e929da860fb4533faa2a/memory_server.py

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

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

from fastmcp import FastMCP

SAFE_KEY_PATTERN = re.compile(r"^[A-Za-z0-9_\-:.]+$")


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

    def _validate_key(self, key: str) -> bool:
        return bool(key) and bool(SAFE_KEY_PATTERN.match(key))

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

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

    def _entry(self, key: str, value: Any) -> dict[str, Any]:
        return {
            "value": value,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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