diff --git a/memory_server.py b/memory_server.py index 8960ee8..ed5f8c5 100644 --- a/memory_server.py +++ b/memory_server.py @@ -1,178 +1,139 @@ """ -Memory Server for MCP protocol. +Memory Server implementing MCP protocol using FastMCP. Provides basic memory operations with optional namespace support. """ import json -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, List, Optional from fastmcp import FastMCP +import fnmatch class MemoryServer: - """FastMCP based memory server. + """A simple MCP server that stores key/value pairs in a JSON file. - Stores data in a JSON file located at ``./memory_data.json``. - Each entry is stored as ``{namespace:key: {value, timestamp}}``. + The server exposes six tools: + * save – store a value under a key. + * get – retrieve a stored value. + * delete – remove a key. + * list_keys – list all keys with optional glob pattern. + * save_with_namespace – same as ``save`` but prefixes the key with a namespace. + * get_by_namespace – return all items belonging to a namespace. """ - def __init__(self) -> None: + def __init__(self, storage_path: str | Path = "./memory_data.json"): self.mcp = FastMCP("Memory-Server") - self.storage_path = Path("./memory_data.json") - - # --------------------------------------------------------------------- - # Internal helpers - # --------------------------------------------------------------------- - def _load_memory(self) -> Dict[str, Any]: - if not self.storage_path.exists(): - return {} - with self.storage_path.open("r", encoding="utf-8") as f: - return json.load(f) - - def _save_memory(self, data: Dict[str, Any]) -> None: + self.storage_path = Path(storage_path) + # Ensure directory exists for future writes self.storage_path.parent.mkdir(parents=True, exist_ok=True) - with self.storage_path.open("w", encoding="utf-8") as f: + + def _load(self) -> Dict[str, Any]: + """Load the entire memory store from disk. + + Returns an empty dict if file does not exist or is invalid. + """ + try: + if self.storage_path.exists(): + with self.storage_path.open("r", encoding="utf-8") as f: + return json.load(f) + except Exception: + # Corrupted file – start fresh + pass + return {} + + def _save(self, data: Dict[str, Any]) -> None: + """Persist the memory store to disk. + + The file is written atomically by writing to a temp file first. + """ + tmp = self.storage_path.with_suffix(".tmp") + with tmp.open("w", encoding="utf-8") as f: json.dump(data, f, indent=2, ensure_ascii=False) + tmp.replace(self.storage_path) - def _make_key(self, key: str, namespace: str = "default") -> str: - return f"{namespace}:{key}" + def _current_timestamp(self) -> str: + return datetime.now(timezone.utc).isoformat() - # --------------------------------------------------------------------- - # Tools exposed via FastMCP - # --------------------------------------------------------------------- - @self.mcp.tool() - def save(self, key: str, value: Any) -> bool: - """Save a value under ``default`` namespace. + @property + def mcp_tool(self): + # Helper to expose the FastMCP instance for decorators + return self.mcp - Args: - key: Identifier for the value. - value: Serializable data. + # ---------- Basic tools ---------- + @FastMCP.tool() + def save(key: str, value: Any) -> bool: + """Store *value* under *key*. - Returns: - True if saved successfully. + The value is JSON‑serialisable. Existing keys are overwritten. + Returns ``True`` on success. """ - return self.save_with_namespace(key, value, "default") - - @self.mcp.tool() - def get(self, key: str) -> Optional[Dict[str, Any]]: - """Retrieve a value from ``default`` namespace. - - Args: - key: Identifier to fetch. - - Returns: - Dictionary with ``value`` and ``timestamp`` or None if missing. - """ - return self.get_by_namespace("default").get(key) - - @self.mcp.tool() - def delete(self, key: str) -> bool: - """Delete a key from ``default`` namespace. - - Args: - key: Identifier to remove. - - Returns: - True if removed, False otherwise. - """ - return self.delete_from_namespace(key, "default") - - @self.mcp.tool() - def list_keys(self, pattern: str = "*") -> List[str]: - """List all keys in ``default`` namespace matching a glob pattern. - - Args: - pattern: Glob pattern (supports * and ?). - - Returns: - List of key names without namespace prefix. - """ - return self.list_keys_in_namespace("default", pattern) - - @self.mcp.tool() - def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool: - """Save a value under a specific namespace. - - Args: - key: Identifier for the value. - value: Serializable data. - namespace: Namespace name. - - Returns: - True if saved successfully. - """ - full_key = self._make_key(key, namespace) - data = self._load_memory() - data[full_key] = { - "value": value, - "timestamp": datetime.utcnow().isoformat() + "Z", - } - self._save_memory(data) + data = self._load() + data[key] = {"value": value, "timestamp": self._current_timestamp()} + self._save(data) return True - @self.mcp.tool() - def get_by_namespace(self, namespace: str = "default") -> Dict[str, Any]: - """Return all key/value pairs in a namespace. + @FastMCP.tool() + def get(key: str) -> Optional[Dict[str, Any]]: + """Retrieve the stored item for *key*. - Args: - namespace: Namespace name. - - Returns: - Mapping of key to dict with ``value`` and ``timestamp``. + Returns a dict with ``value`` and ``timestamp`` or ``None`` if not found. """ - data = self._load_memory() - prefix = f"{namespace}:" - result: Dict[str, Any] = {} - for k, v in data.items(): - if k.startswith(prefix): - key_name = k[len(prefix) :] - result[key_name] = v - return result + data = self._load() + return data.get(key) - @self.mcp.tool() - def delete_from_namespace(self, key: str, namespace: str = "default") -> bool: - """Delete a key from a specific namespace. + @FastMCP.tool() + def delete(key: str) -> bool: + """Remove *key* from the store. - Args: - key: Identifier to remove. - namespace: Namespace name. - - Returns: - True if removed, False otherwise. + Returns ``True`` if key existed and was removed, otherwise ``False``. """ - full_key = self._make_key(key, namespace) - data = self._load_memory() - if full_key in data: - del data[full_key] - self._save_memory(data) + data = self._load() + if key in data: + del data[key] + self._save(data) return True return False - @self.mcp.tool() - def list_keys_in_namespace(self, namespace: str = "default", pattern: str = "*") -> List[str]: - """List keys in a namespace matching a glob pattern. + @FastMCP.tool() + def list_keys(pattern: str = "*") -> List[str]: + """Return all keys matching *pattern*. - Args: - namespace: Namespace name. - pattern: Glob pattern. - - Returns: - List of key names without namespace prefix. + The pattern supports Unix shell wildcards (``*``, ``?``). """ - import fnmatch + data = self._load() + return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)] - data = self._load_memory() + # ---------- Namespace tools ---------- + @FastMCP.tool() + def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool: + """Store *value* under ``namespace:key``. + + The key is prefixed with the namespace and a colon. Existing entries are overwritten. + Returns ``True`` on success. + """ + namespaced_key = f"{namespace}:{key}" + return self.save(namespaced_key, value) + + @FastMCP.tool() + def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]: + """Return all items belonging to *namespace*. + + Each returned dict contains ``key`` (without namespace prefix), ``value`` and ``timestamp``. + """ + data = self._load() prefix = f"{namespace}:" - keys = [k[len(prefix) :] for k in data if k.startswith(prefix)] - return [k for k in keys if fnmatch.fnmatch(k, pattern)] - - # --------------------------------------------------------------------- - # Entry point - # --------------------------------------------------------------------- - def run(self) -> None: - self.mcp.run(transport="stdio", show_banner=False, log_level="ERROR") + result: List[Dict[str, Any]] = [] + for full_key, meta in data.items(): + if full_key.startswith(prefix): + key_without_ns = full_key[len(prefix) :] + item = {"key": key_without_ns, "value": meta["value"], "timestamp": meta["timestamp"]} + result.append(item) + return result +# ---------- Server entry point ---------- if __name__ == "__main__": server = MemoryServer() - server.run() + # Run with stdio transport – suitable for local testing and client usage. + server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")