@trac3er/oh-my-god 2.1.8 → 2.2.2

package/build/lib/runtime/mcp_memory_server.py

@@ -2,21 +2,24 @@ from __future__ import annotations
 
 import json
 import os
+import hashlib
+from pathlib import Path
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from datetime import datetime, timezone
 from typing import Any
 
 _MCP_IMPORT_ERROR: ModuleNotFoundError | None = None
 
 try:
-    from fastmcp import FastMCP
-    from starlette.requests import Request
-    from starlette.responses import JSONResponse
+    from fastmcp import FastMCP as FastMCPImpl
+    from starlette.requests import Request as RequestImpl
+    from starlette.responses import JSONResponse as JSONResponseImpl
 except ModuleNotFoundError as exc:
     _MCP_IMPORT_ERROR = exc
-    Request = Any
+    RequestImpl = Any
 
-    class JSONResponse(dict):
+    class JSONResponseFallback(dict):
         def __init__(self, content: dict[str, Any]):
             super().__init__(content)
 
@@ -26,7 +29,7 @@ except ModuleNotFoundError as exc:
 
         return decorator
 
-    class FastMCP:  # type: ignore[override]
+    class FastMCPFallback:  # type: ignore[override]
         def __init__(self, *_args: Any, **_kwargs: Any) -> None:
             self._import_error = _MCP_IMPORT_ERROR
 
@@ -37,11 +40,32 @@ except ModuleNotFoundError as exc:
         def run(self, *_args: Any, **_kwargs: Any) -> None:
             raise RuntimeError("fastmcp and starlette are required to run the OMG memory server") from self._import_error
 
-    FastMCP.__module__ = "fastmcp"
+    FastMCPFallback.__module__ = "fastmcp"
+    FastMCPImpl = FastMCPFallback
+    JSONResponseImpl = JSONResponseFallback
 
 from runtime.memory_store import MemoryStore, MemoryStoreFullError
 
-_store = MemoryStore()
+_store = MemoryStore(store_path=str(Path.home() / ".omg" / "shared-memory" / "store.json"))
+
+
+class _HybridExportBundle(list[dict[str, Any]]):
+    def __init__(self, items: list[dict[str, Any]], bundle: dict[str, Any]) -> None:
+        super().__init__(items)
+        self._bundle = bundle
+
+    def __getitem__(self, key: Any) -> Any:
+        if isinstance(key, str):
+            return self._bundle[key]
+        return super().__getitem__(key)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self._bundle.get(key, default)
+
+    def __contains__(self, item: object) -> bool:
+        if isinstance(item, str):
+            return item in self._bundle
+        return super().__contains__(item)
 
 
 def _load_state() -> None:
@@ -61,12 +85,12 @@ async def lifespan(_: object) -> AsyncIterator[None]:
         _save_state()
 
 
-mcp = FastMCP("OMG Memory Server", lifespan=lifespan)
+mcp = FastMCPImpl("OMG Memory Server", lifespan=lifespan)
 
 
 @mcp.custom_route("/health", methods=["GET"])
-async def health(_: Request) -> JSONResponse:
-    return JSONResponse({"status": "ok", "version": "2.0.8"})
+async def health(_: Any) -> Any:
+    return JSONResponseImpl({"status": "ok", "version": "2.0.8"})
 
 
 @mcp.tool()
@@ -75,26 +99,39 @@ def memory_store(
     content: str,
     source_cli: str,
     tags: list[str] | None = None,
+    namespace: str = "default",
+    retention_days: int | None = None,
 ) -> dict[str, Any]:
     try:
-        return _store.add(key=key, content=content, source_cli=source_cli, tags=tags)
+        return _store.add(
+            key=key,
+            content=content,
+            source_cli=source_cli,
+            tags=tags,
+            namespace=namespace,
+            retention_days=retention_days,
+        )
     except MemoryStoreFullError as exc:
         return {"error": str(exc)}
+    except ValueError as exc:
+        return {"error": str(exc)}
 
 
 @mcp.tool()
 def memory_search(
     query: str,
     source_cli: str | None = None,
+    namespace: str | None = None,
 ) -> list[dict[str, Any]]:
-    return _store.search(query=query, source_cli=source_cli)
+    return _store.search(query=query, source_cli=source_cli, namespace=namespace)
 
 
 @mcp.tool()
 def memory_list(
     source_cli: str | None = None,
+    namespace: str | None = None,
 ) -> list[dict[str, Any]]:
-    return _store.list_all(source_cli=source_cli)
+    return _store.list_all(source_cli=source_cli, namespace=namespace)
 
 
 @mcp.tool()
@@ -105,13 +142,61 @@ def memory_delete(item_id: str) -> dict[str, Any]:
 
 @mcp.tool()
 def memory_import(items: list[dict[str, Any]]) -> dict[str, int]:
-    count = _store.import_items(items)
+    count = _store.import_items(items, quarantined=False)
     return {"imported": count}
 
 
 @mcp.tool()
-def memory_export() -> list[dict[str, Any]]:
-    return _store.export_all()
+def memory_import_bundle(bundle: dict[str, Any]) -> dict[str, Any]:
+    if str(bundle.get("format", "")) != "omg.memory.export.v1":
+        return {"imported": 0, "error": "bundle format is invalid"}
+    payload = str(bundle.get("encrypted_payload", ""))
+    integrity = bundle.get("integrity", {})
+    expected_sha = ""
+    if isinstance(integrity, dict):
+        expected_sha = str(integrity.get("sha256", ""))
+    if not payload or hashlib.sha256(payload.encode("utf-8")).hexdigest() != expected_sha:
+        return {"imported": 0, "error": "bundle integrity check failed"}
+    if not payload.startswith("enc:v1:"):
+        return {"imported": 0, "error": "bundle payload must be encrypted"}
+    decoded = _store._decrypt_text(payload, purpose="export-bundle")  # noqa: SLF001
+    if not decoded:
+        return {"imported": 0, "error": "bundle decryption failed"}
+    try:
+        raw_items = json.loads(decoded)
+    except (TypeError, ValueError, json.JSONDecodeError):
+        return {"imported": 0, "error": "bundle payload is not valid JSON"}
+    if not isinstance(raw_items, list):
+        return {"imported": 0, "error": "bundle payload must be a list"}
+    count = _store.import_items(raw_items, quarantined=True)
+    return {"imported": count, "quarantined": count}
+
+
+@mcp.tool()
+def memory_export(namespace: str | None = None) -> list[dict[str, Any]] | dict[str, Any]:
+    items = _store.list_all(namespace=namespace)
+    serialized = json.dumps(items, separators=(",", ":"), ensure_ascii=True)
+    encrypted_payload = _store._encrypt_text(serialized, purpose="export-bundle")  # noqa: SLF001
+    bundle = {
+        "format": "omg.memory.export.v1",
+        "encrypted_payload": encrypted_payload,
+        "integrity": {
+            "sha256": hashlib.sha256(encrypted_payload.encode("utf-8")).hexdigest(),
+            "algorithm": "sha256",
+        },
+        "metadata": {
+            "count": len(items),
+            "exported_at": datetime.now(timezone.utc).isoformat(),
+            "namespace": namespace,
+        },
+    }
+    return _HybridExportBundle(items, bundle)
+
+
+@mcp.tool()
+def memory_promote(item_id: str) -> dict[str, Any]:
+    promoted = _store.promote_item(item_id)
+    return {"promoted": promoted, "id": item_id}
 
 
 @mcp.resource("memory://all")