diff --git a/memory_server.py b/memory_server.py new file mode 100644 index 0000000..7405969 --- /dev/null +++ b/memory_server.py @@ -0,0 +1,202 @@ +"""Memory Server for MCP protocol. + +This server exposes tools for storing, retrieving, and managing key/value pairs +with optional namespace support. The data is persisted to a JSON file with +metadata (timestamp). +""" + +import json +import fnmatch +from datetime import datetime +from pathlib import Path +from typing import Any, Optional, List + +from fastmcp import FastMCP + + +class MemoryServer: + def __init__(self): + self.mcp = FastMCP("Memory-Server") + self.storage_path = Path("./memory_data.json") + # Load existing data into memory for quick access + self._data = self._load_memory() + + def _load_memory(self) -> dict: + """Load memory from JSON file. + + Returns: + dict: Mapping of keys to value+timestamp dicts. + """ + 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): + """Persist memory to JSON file. + + Args: + data (dict): The memory mapping to write. + """ + 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) + + # ---------- Basic tools ---------- + @self.mcp.tool() + def save(self, key: str, value: Any) -> bool: + """Save a value under a key. + + Args: + key: Unique identifier. + value: Any JSON‑serializable value. + + Returns: + True on success. + """ + try: + self._data[key] = { + "value": value, + "timestamp": datetime.utcnow().isoformat() + } + self._save_memory(self._data) + return True + except Exception as e: + print(f"Error saving key {key}: {e}") + return False + + @self.mcp.tool() + def get(self, key: str) -> Optional[dict]: + """Retrieve a value by key. + + Args: + key: Identifier. + + Returns: + Dict with fields ``key``, ``value``, ``timestamp`` or None. + """ + entry = self._data.get(key) + if entry is None: + return None + return { + "key": key, + "value": entry["value"], + "timestamp": entry["timestamp"], + } + + @self.mcp.tool() + def delete(self, key: str) -> bool: + """Delete a key from memory. + + Args: + key: Identifier. + + Returns: + True if deleted, False if key not found. + """ + if key in self._data: + del self._data[key] + self._save_memory(self._data) + return True + return False + + @self.mcp.tool() + def list_keys(self, pattern: str = "*") -> List[str]: + """List all keys matching a wildcard pattern. + + Args: + pattern: Wildcard pattern (supports ``*`` and ``?``). + + Returns: + List of matching keys. + """ + return [k for k in self._data.keys() if fnmatch.fnmatch(k, pattern)] + + # ---------- Namespace tools ---------- + @self.mcp.tool() + def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool: + """Save a value under a namespaced key. + + Args: + key: Identifier. + value: Value to store. + namespace: Namespace prefix. + + Returns: + True on success. + """ + composite_key = f"{namespace}:{key}" + return self.save(composite_key, value) + + @self.mcp.tool() + def get_by_namespace(self, namespace: str = "default") -> List[dict]: + """Retrieve all entries in a namespace. + + Args: + namespace: Namespace to query. + + Returns: + List of dicts with ``key``, ``value``, ``timestamp``. + """ + prefix = f"{namespace}:" + results = [] + for k, v in self._data.items(): + if k.startswith(prefix): + results.append({ + "key": k, + "value": v["value"], + "timestamp": v["timestamp"], + }) + return results + + @self.mcp.tool() + def get_namespace_key(self, key: str, namespace: str = "default") -> Optional[dict]: + """Retrieve a single key within a namespace. + + Args: + key: Identifier without namespace. + namespace: Namespace. + + Returns: + Entry dict or None. + """ + composite_key = f"{namespace}:{key}" + return self.get(composite_key) + + @self.mcp.tool() + def delete_namespace_key(self, key: str, namespace: str = "default") -> bool: + """Delete a single key within a namespace. + + Args: + key: Identifier without namespace. + namespace: Namespace. + + Returns: + True if deleted. + """ + composite_key = f"{namespace}:{key}" + return self.delete(composite_key) + + @self.mcp.tool() + def list_namespace_keys(self, namespace: str = "default", pattern: str = "*") -> List[str]: + """List keys in a namespace matching a pattern. + + Args: + namespace: Namespace. + pattern: Wildcard pattern. + + Returns: + List of matching keys (including namespace prefix). + """ + prefix = f"{namespace}:" + return [k for k in self._data.keys() if k.startswith(prefix) and fnmatch.fnmatch(k, pattern)] + + +# ---------- Server entry point ---------- +if __name__ == "__main__": + server = MemoryServer() + server.mcp.run( + transport="stdio", + show_banner=False, + log_level="ERROR", + )