diff --git a/memory_server.py b/memory_server.py new file mode 100644 index 0000000..50dabfb --- /dev/null +++ b/memory_server.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import fnmatch +import json +from datetime import datetime +from pathlib import Path +from typing import Any, Optional + +from fastmcp import FastMCP + + +def validate_token(value: str, field: str, allow_colon: bool = False) -> str: + value = value.strip() + if not value: + raise ValueError(f"{field} must not be empty") + + forbidden = ["..", "/", "\\"] + if not allow_colon: + forbidden.append(":") + + for marker in forbidden: + if marker in value: + raise ValueError(f"{field} contains forbidden sequence: {marker}") + return value + + +class MemoryServer: + def __init__(self, storage_path: Path | None = None): + self.mcp = FastMCP("Memory-Server") + self.storage_path = storage_path or Path("./memory_data.json") + self._register_tools() + + def _load_memory(self) -> dict: + if not self.storage_path.exists(): + return {} + with open(self.storage_path, "r", encoding="utf-8") as f: + return json.load(f) + + def _save_memory(self, data: dict): + 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 _make_record(self, value: Any) -> dict[str, Any]: + return { + "value": value, + "timestamp": datetime.now().isoformat(timespec="seconds"), + } + + def _save_entry(self, key: str, value: Any) -> bool: + key = validate_token(key, "key", allow_colon=True) + data = self._load_memory() + data[key] = self._make_record(value) + self._save_memory(data) + return True + + def _get_entry(self, key: str) -> Optional[dict]: + key = validate_token(key, "key", allow_colon=True) + data = self._load_memory() + record = data.get(key) + if record is None: + return None + return {"key": key, **record} + + def _delete_entry(self, key: str) -> bool: + key = validate_token(key, "key", allow_colon=True) + data = self._load_memory() + if key not in data: + return False + del data[key] + self._save_memory(data) + return True + + def _list_entry_keys(self, pattern: str = "*") -> list[str]: + data = self._load_memory() + return sorted(key for key in data if fnmatch.fnmatch(key, pattern)) + + def _save_namespaced(self, key: str, value: Any, namespace: str = "default") -> bool: + namespace = validate_token(namespace, "namespace") + key = validate_token(key, "key") + return self._save_entry(f"{namespace}:{key}", value) + + def _get_namespace_entries(self, namespace: str = "default") -> list[dict]: + namespace = validate_token(namespace, "namespace") + prefix = f"{namespace}:" + data = self._load_memory() + result = [] + for key in sorted(data): + if key.startswith(prefix): + result.append({"key": key, **data[key]}) + return result + + def _register_tools(self) -> None: + @self.mcp.tool() + def save(key: str, value: Any) -> bool: + """Save a serializable value under a key.""" + + return self._save_entry(key, value) + + @self.mcp.tool() + def get(key: str) -> Optional[dict]: + """Return a saved value with its metadata or None.""" + + return self._get_entry(key) + + @self.mcp.tool() + def delete(key: str) -> bool: + """Delete a key from memory if it exists.""" + + return self._delete_entry(key) + + @self.mcp.tool() + def list_keys(pattern: str = "*") -> list[str]: + """List saved keys filtered by a wildcard pattern.""" + + return self._list_entry_keys(pattern) + + @self.mcp.tool() + def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool: + """Save a value inside a namespace using namespace:key format.""" + + return self._save_namespaced(key, value, namespace) + + @self.mcp.tool() + def get_by_namespace(namespace: str = "default") -> list[dict]: + """Return all records that belong to the namespace.""" + + return self._get_namespace_entries(namespace) + + +if __name__ == "__main__": + server = MemoryServer() + server.mcp.run( + transport="stdio", + show_banner=False, + log_level="ERROR", + )