140 lines
4.7 KiB
Python
140 lines
4.7 KiB
Python
"""
|
||
Memory Server implementing MCP protocol using FastMCP.
|
||
|
||
Provides basic memory operations with optional namespace support.
|
||
"""
|
||
import json
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from fastmcp import FastMCP
|
||
import fnmatch
|
||
|
||
class MemoryServer:
|
||
"""A simple MCP server that stores key/value pairs in a JSON file.
|
||
|
||
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, storage_path: str | Path = "./memory_data.json"):
|
||
self.mcp = FastMCP("Memory-Server")
|
||
self.storage_path = Path(storage_path)
|
||
# Ensure directory exists for future writes
|
||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
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 _current_timestamp(self) -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
@property
|
||
def mcp_tool(self):
|
||
# Helper to expose the FastMCP instance for decorators
|
||
return self.mcp
|
||
|
||
# ---------- Basic tools ----------
|
||
@FastMCP.tool()
|
||
def save(key: str, value: Any) -> bool:
|
||
"""Store *value* under *key*.
|
||
|
||
The value is JSON‑serialisable. Existing keys are overwritten.
|
||
Returns ``True`` on success.
|
||
"""
|
||
data = self._load()
|
||
data[key] = {"value": value, "timestamp": self._current_timestamp()}
|
||
self._save(data)
|
||
return True
|
||
|
||
@FastMCP.tool()
|
||
def get(key: str) -> Optional[Dict[str, Any]]:
|
||
"""Retrieve the stored item for *key*.
|
||
|
||
Returns a dict with ``value`` and ``timestamp`` or ``None`` if not found.
|
||
"""
|
||
data = self._load()
|
||
return data.get(key)
|
||
|
||
@FastMCP.tool()
|
||
def delete(key: str) -> bool:
|
||
"""Remove *key* from the store.
|
||
|
||
Returns ``True`` if key existed and was removed, otherwise ``False``.
|
||
"""
|
||
data = self._load()
|
||
if key in data:
|
||
del data[key]
|
||
self._save(data)
|
||
return True
|
||
return False
|
||
|
||
@FastMCP.tool()
|
||
def list_keys(pattern: str = "*") -> List[str]:
|
||
"""Return all keys matching *pattern*.
|
||
|
||
The pattern supports Unix shell wildcards (``*``, ``?``).
|
||
"""
|
||
data = self._load()
|
||
return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]
|
||
|
||
# ---------- 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}:"
|
||
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()
|
||
# Run with stdio transport – suitable for local testing and client usage.
|
||
server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
|