fix: memory_server.py — build_server() + main.py для автопроверки

This commit is contained in:
2026-05-27 08:21:12 +00:00
parent c80e104057
commit 8cde94ff2c
+52 -44
View File
@@ -4,13 +4,13 @@ from __future__ import annotations
import fnmatch import fnmatch
import json import json
import re import re
from datetime import datetime, timezone from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import Any, Optional from typing import Any, Optional
from fastmcp import FastMCP from fastmcp import FastMCP
SAFE_KEY_PATTERN = re.compile(r"^[A-Za-z0-9_\-:.]+$") _UNSAFE_KEY = re.compile(r"\.\.|/|\\")
class MemoryServer: class MemoryServer:
@@ -20,30 +20,29 @@ class MemoryServer:
self._register_tools() self._register_tools()
def _validate_key(self, key: str) -> bool: def _validate_key(self, key: str) -> bool:
return bool(key) and bool(SAFE_KEY_PATTERN.match(key)) return bool(key) and _UNSAFE_KEY.search(key) is None
def _load_memory(self) -> dict[str, dict[str, Any]]: def _load_memory(self) -> dict:
if not self.storage_path.exists(): if not self.storage_path.exists():
return {} return {}
with open(self.storage_path, encoding="utf-8") as file: with open(self.storage_path, encoding="utf-8") as f:
return json.load(file) return json.load(f)
def _save_memory(self, data: dict[str, dict[str, Any]]) -> None: def _save_memory(self, data: dict) -> None:
self.storage_path.parent.mkdir(parents=True, exist_ok=True) self.storage_path.parent.mkdir(parents=True, exist_ok=True)
with open(self.storage_path, "w", encoding="utf-8") as file: with open(self.storage_path, "w", encoding="utf-8") as f:
json.dump(data, file, indent=2, ensure_ascii=False) json.dump(data, f, indent=2, ensure_ascii=False)
def _entry(self, key: str, value: Any) -> dict[str, Any]: def _timestamp(self) -> str:
return { return datetime.now().isoformat(timespec="seconds")
"value": value,
"timestamp": datetime.now(timezone.utc).isoformat(), def _make_entry(self, value: Any) -> dict:
} return {"value": value, "timestamp": self._timestamp()}
def _register_tools(self) -> None: def _register_tools(self) -> None:
mcp = self.mcp mcp = self.mcp
server = self
@mcp.tool @mcp.tool()
def save(key: str, value: Any) -> bool: def save(key: str, value: Any) -> bool:
"""Сохраняет значение по ключу в память сервера. """Сохраняет значение по ключу в память сервера.
@@ -54,17 +53,17 @@ class MemoryServer:
Returns: Returns:
True при успешном сохранении, False иначе. True при успешном сохранении, False иначе.
""" """
if not server._validate_key(key): if not self._validate_key(key):
return False return False
try: try:
data = server._load_memory() data = self._load_memory()
data[key] = server._entry(key, value) data[key] = self._make_entry(value)
server._save_memory(data) self._save_memory(data)
return True return True
except (OSError, TypeError, ValueError): except (TypeError, ValueError, OSError):
return False return False
@mcp.tool @mcp.tool()
def get(key: str) -> Optional[dict]: def get(key: str) -> Optional[dict]:
"""Возвращает значение по ключу с метаданными. """Возвращает значение по ключу с метаданными.
@@ -74,18 +73,18 @@ class MemoryServer:
Returns: Returns:
Словарь с полями key, value, timestamp или None если ключ не найден. Словарь с полями key, value, timestamp или None если ключ не найден.
""" """
if not server._validate_key(key): if not self._validate_key(key):
return None return None
record = server._load_memory().get(key) entry = self._load_memory().get(key)
if record is None: if entry is None:
return None return None
return { return {
"key": key, "key": key,
"value": record.get("value"), "value": entry["value"],
"timestamp": record.get("timestamp"), "timestamp": entry["timestamp"],
} }
@mcp.tool @mcp.tool()
def delete(key: str) -> bool: def delete(key: str) -> bool:
"""Удаляет ключ из памяти сервера. """Удаляет ключ из памяти сервера.
@@ -95,16 +94,16 @@ class MemoryServer:
Returns: Returns:
True при успешном удалении, False если ключ не найден. True при успешном удалении, False если ключ не найден.
""" """
if not server._validate_key(key): if not self._validate_key(key):
return False return False
data = server._load_memory() data = self._load_memory()
if key not in data: if key not in data:
return False return False
del data[key] del data[key]
server._save_memory(data) self._save_memory(data)
return True return True
@mcp.tool @mcp.tool()
def list_keys(pattern: str = "*") -> list[str]: def list_keys(pattern: str = "*") -> list[str]:
"""Возвращает список всех ключей с поддержкой wildcard-паттерна. """Возвращает список всех ключей с поддержкой wildcard-паттерна.
@@ -114,11 +113,13 @@ class MemoryServer:
Returns: Returns:
Список совпадающих ключей. Список совпадающих ключей.
""" """
keys = list(server._load_memory().keys()) keys = list(self._load_memory().keys())
return sorted(key for key in keys if fnmatch.fnmatch(key, pattern)) return sorted(k for k in keys if fnmatch.fnmatch(k, pattern))
@mcp.tool @mcp.tool()
def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool: def save_with_namespace(
key: str, value: Any, namespace: str = "default"
) -> bool:
"""Сохраняет значение с указанием пространства имён. """Сохраняет значение с указанием пространства имён.
Args: Args:
@@ -129,10 +130,10 @@ class MemoryServer:
Returns: Returns:
True при успехе, False иначе. True при успехе, False иначе.
""" """
full_key = f"{namespace}:{key}" storage_key = f"{namespace}:{key}"
return save(full_key, value) return save(storage_key, value)
@mcp.tool @mcp.tool()
def get_by_namespace(namespace: str = "default") -> list[dict]: def get_by_namespace(namespace: str = "default") -> list[dict]:
"""Возвращает все ключи из указанного namespace. """Возвращает все ключи из указанного namespace.
@@ -143,24 +144,31 @@ class MemoryServer:
Список словарей с метаданными всех ключей namespace. Список словарей с метаданными всех ключей namespace.
""" """
prefix = f"{namespace}:" prefix = f"{namespace}:"
data = self._load_memory()
result: list[dict] = [] result: list[dict] = []
for storage_key, record in server._load_memory().items(): for storage_key, entry in data.items():
if not storage_key.startswith(prefix): if not storage_key.startswith(prefix):
continue continue
short_key = storage_key[len(prefix) :] short_key = storage_key[len(prefix) :]
result.append( result.append(
{ {
"key": short_key, "key": short_key,
"value": record.get("value"), "value": entry["value"],
"timestamp": record.get("timestamp"), "timestamp": entry["timestamp"],
"storage_key": storage_key,
} }
) )
return sorted(result, key=lambda item: item["key"]) return sorted(result, key=lambda item: item["key"])
_ = (save, get, delete, list_keys, save_with_namespace, get_by_namespace)
def build_server(storage_path: str | Path = "./memory_data.json") -> MemoryServer:
"""Сборка MCP-сервера — точка входа для автопроверки."""
return MemoryServer(storage_path=storage_path)
if __name__ == "__main__": if __name__ == "__main__":
server = MemoryServer() server = build_server()
server.mcp.run( server.mcp.run(
transport="stdio", transport="stdio",
show_banner=False, show_banner=False,