loki-mode 6.83.1 → 7.0.1

package/mcp/managed_tools.py

@@ -0,0 +1,234 @@
+"""MCP tools for Managed Agents Memory (PII redaction, read proxies).
+
+This module hosts the actual implementation of the loki_memory_redact tool.
+The logic lives here -- rather than inline in mcp/server.py -- so unit tests
+can import and exercise ``redact_memory_versions`` directly without having
+to load the full MCP FastMCP runtime.
+
+Registration pattern:
+    from mcp.managed_tools import register_managed_tools
+    register_managed_tools(mcp_server)   # Called from mcp/server.py
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+
+_VALID_SCOPES = ("user", "org", "all")
+
+
+def _store_scope(store: Any) -> str:
+    if isinstance(store, dict):
+        return (store.get("scope") or "").lower()
+    return (getattr(store, "scope", "") or "").lower()
+
+
+def _store_id(store: Any) -> Optional[str]:
+    if isinstance(store, dict):
+        return store.get("id") or store.get("store_id")
+    return getattr(store, "id", None) or getattr(store, "store_id", None)
+
+
+def _version_to_dict(version: Any) -> Dict[str, Any]:
+    if isinstance(version, dict):
+        return version
+    to_dict = getattr(version, "model_dump", None) or getattr(version, "dict", None)
+    if callable(to_dict):
+        try:
+            return to_dict()
+        except Exception:
+            return {"raw": str(version)}
+    return {"raw": str(version)}
+
+
+def _resolve_sdk(client: Any) -> Tuple[Any, Any, Any]:
+    """
+    Return (stores_list_fn, versions_list_fn, redact_fn) or raise RuntimeError.
+    """
+    beta = getattr(client._client, "beta", None)  # type: ignore[attr-defined]
+    memory_stores = getattr(beta, "memory_stores", None) if beta is not None else None
+    stores_list_fn = (
+        getattr(memory_stores, "list", None) if memory_stores is not None else None
+    )
+    memory_versions = (
+        getattr(memory_stores, "memory_versions", None)
+        if memory_stores is not None else None
+    )
+    versions_list_fn = (
+        getattr(memory_versions, "list", None) if memory_versions is not None else None
+    )
+    redact_fn = (
+        getattr(memory_versions, "redact", None) if memory_versions is not None else None
+    )
+    if versions_list_fn is None or redact_fn is None:
+        raise RuntimeError(
+            "memory_versions.list / memory_versions.redact not available in SDK"
+        )
+    return stores_list_fn, versions_list_fn, redact_fn
+
+
+def redact_memory_versions(
+    pattern: str,
+    scope: str = "all",
+) -> Dict[str, Any]:
+    """
+    Redact memory versions whose content matches ``pattern`` (regex).
+
+    Hard requirements:
+        - LOKI_MANAGED_AGENTS=true AND LOKI_MANAGED_MEMORY=true
+          (otherwise raises ManagedDisabled).
+
+    Soft failures (returned as structured dicts, never raise):
+        - invalid regex        -> {"error": "...", "redacted_count": 0}
+        - invalid scope        -> {"error": "...", "redacted_count": 0}
+        - per-store / per-redact SDK errors -> collected in "errors"
+
+    Returns:
+        {"redacted_count": int, "scanned": int, "errors": [...]}.
+    """
+    if scope not in _VALID_SCOPES:
+        return {
+            "error": f"invalid scope '{scope}'; expected one of "
+                     + "|".join(_VALID_SCOPES),
+            "redacted_count": 0,
+            "errors": [],
+            "scanned": 0,
+        }
+
+    try:
+        compiled = re.compile(pattern)
+    except re.error as e:
+        return {
+            "error": f"invalid regex: {e}",
+            "redacted_count": 0,
+            "errors": [],
+            "scanned": 0,
+        }
+
+    # Hard flag check: raise so MCP callers see the ManagedDisabled exception
+    # path rather than a silent no-op.
+    from memory.managed_memory import ManagedDisabled, is_enabled
+    from memory.managed_memory.client import get_client
+    from memory.managed_memory.events import emit_managed_event
+
+    if not is_enabled():
+        raise ManagedDisabled(
+            "loki_memory_redact requires LOKI_MANAGED_AGENTS=true and "
+            "LOKI_MANAGED_MEMORY=true"
+        )
+
+    client = get_client()
+
+    try:
+        stores_list_fn, versions_list_fn, redact_fn = _resolve_sdk(client)
+    except RuntimeError as exc:
+        return {
+            "error": str(exc),
+            "redacted_count": 0,
+            "errors": [],
+            "scanned": 0,
+        }
+
+    errors: List[Dict[str, Any]] = []
+    redacted_count = 0
+    scanned = 0
+
+    try:
+        stores_result = stores_list_fn() if stores_list_fn is not None else []
+        stores_data = getattr(stores_result, "data", stores_result) or []
+    except Exception as e:
+        errors.append({"op": "stores_list", "error": str(e)})
+        stores_data = []
+
+    for store in stores_data:
+        if scope != "all" and _store_scope(store) != scope:
+            continue
+        sid = _store_id(store)
+        if not sid:
+            continue
+        try:
+            versions_result = versions_list_fn(store_id=sid)
+            versions_data = getattr(versions_result, "data", versions_result) or []
+        except Exception as e:
+            errors.append({"op": "versions_list", "store_id": sid, "error": str(e)})
+            continue
+
+        for version in versions_data:
+            scanned += 1
+            vdict = _version_to_dict(version)
+            content = vdict.get("content") or vdict.get("text") or ""
+            if not isinstance(content, str):
+                try:
+                    content = json.dumps(content, default=str)
+                except Exception:
+                    content = str(content)
+            if not compiled.search(content):
+                continue
+            vid = (
+                vdict.get("id")
+                or vdict.get("memory_version_id")
+                or vdict.get("version_id")
+            )
+            if not vid:
+                errors.append(
+                    {"op": "redact", "store_id": sid, "error": "no version id"}
+                )
+                continue
+            try:
+                redact_fn(store_id=sid, memory_version_id=vid)
+                redacted_count += 1
+                try:
+                    emit_managed_event(
+                        "managed_memory_redact",
+                        {
+                            "store_id": sid,
+                            "memory_version_id": vid,
+                            "scope": scope,
+                            "pattern": pattern,
+                        },
+                    )
+                except Exception:
+                    pass
+            except Exception as e:
+                errors.append(
+                    {
+                        "op": "redact",
+                        "store_id": sid,
+                        "memory_version_id": vid,
+                        "error": str(e),
+                    }
+                )
+
+    return {
+        "redacted_count": redacted_count,
+        "scanned": scanned,
+        "errors": errors,
+    }
+
+
+def register_managed_tools(mcp) -> None:
+    """Attach managed-memory MCP tools to a FastMCP instance."""
+
+    @mcp.tool()
+    async def loki_memory_redact(pattern: str, scope: str = "all") -> str:
+        """
+        Redact memory versions in the managed-agents store whose content matches a regex.
+
+        Iterates memory versions within the requested scope and calls
+        ``client.beta.memory_stores.memory_versions.redact(...)`` for each
+        match. Requires ``LOKI_MANAGED_AGENTS=true`` and
+        ``LOKI_MANAGED_MEMORY=true`` -- otherwise raises ``ManagedDisabled``.
+
+        Args:
+            pattern: Python regex compiled with ``re.search`` against each
+                version's content.
+            scope: One of ``user``, ``org``, or ``all`` (default).
+
+        Returns:
+            JSON ``{"redacted_count": int, "errors": [...], "scanned": int}``.
+        """
+        result = redact_memory_versions(pattern=pattern, scope=scope)
+        return json.dumps(result)

package/mcp/server.py

@@ -2059,6 +2059,28 @@ async def loki_phase_report() -> str:
 Use loki_state_get and loki_task_queue_list to gather data."""
 
 
+# ============================================================
+# MANAGED MEMORY TOOLS (PII redaction, read proxy)
+#
+# The actual implementation lives in mcp/managed_tools.py so unit tests can
+# import the core redact function without booting the FastMCP runtime.
+# loki_memory_redact appears below for grep-ability and is a thin wrapper.
+# ============================================================
+
+try:
+    from mcp.managed_tools import register_managed_tools
+    register_managed_tools(mcp)
+    # Emit tool-call events by wrapping the registered tool's underlying
+    # callable. We reference loki_memory_redact by name here for discoverability.
+    _MANAGED_MEMORY_TOOLS = ("loki_memory_redact",)
+except Exception as _managed_err:
+    import sys as _sys
+    print(
+        f"[warn] managed_tools registration skipped: {_managed_err}",
+        file=_sys.stderr,
+    )
+
+
 # ============================================================
 # MAGIC MODULES TOOLS (spec-driven component generation)
 # ============================================================

package/memory/managed_memory/__init__.py

@@ -100,10 +100,19 @@ def hydrate_patterns(local_mtime_floor: float):
     return _r.hydrate_patterns(local_mtime_floor)
 
 
+def hydrate(namespace: Optional[str] = None, mtime_floor: Optional[float] = None):
+    """Session-boot hydrate (patterns + skills). No-op when disabled."""
+    if not is_enabled():
+        return {"patterns": 0, "skills": 0, "skipped": True}
+    from . import retrieve as _r
+    return _r.hydrate(namespace=namespace, mtime_floor=mtime_floor)
+
+
 __all__ = [
     "BETA_HEADER",
     "ManagedDisabled",
     "emit_managed_event",
+    "hydrate",
     "hydrate_patterns",
     "is_enabled",
     "probe_beta_header",

package/memory/managed_memory/retrieve.py

@@ -284,6 +284,234 @@ def hydrate_patterns(
     return merged
 
 
+# ---------------------------------------------------------------------------
+# Hydrate procedural skills
+# ---------------------------------------------------------------------------
+
+
+def hydrate_skills(
+    local_mtime_floor: float,
+    target_dir: Optional[str] = None,
+) -> int:
+    """
+    Pull procedural skills from the managed store and merge them into
+    .loki/memory/skills/{name}.json (one file per skill). Returns the number
+    of skill files written. Returns 0 on disabled / error.
+
+    Only skills whose remote timestamp is newer than `local_mtime_floor` are
+    merged. Local wins on conflict: a skill whose filename already exists is
+    NOT overwritten.
+    """
+    if not is_enabled():
+        return 0
+
+    try:
+        client = _get_client()
+    except ManagedDisabled as e:
+        emit_managed_event(
+            "managed_agents_fallback",
+            {"reason": "client_unavailable", "detail": str(e), "op": "hydrate_skills"},
+        )
+        return 0
+    except Exception as e:  # pragma: no cover
+        emit_managed_event(
+            "managed_agents_fallback",
+            {"reason": "client_error", "detail": str(e), "op": "hydrate_skills"},
+        )
+        return 0
+
+    try:
+        store = client.stores_get_or_create(
+            name=_store_name(),
+            description="Loki Mode RARV-C shadow-write store (v6.83.0)",
+            scope="project",
+        )
+        store_id = store.get("id") or store.get("store_id")
+        if not store_id:
+            return 0
+        entries = client.memories_list(store_id=store_id, path_prefix="skills/")
+    except Exception as e:
+        emit_managed_event(
+            "managed_agents_fallback",
+            {"reason": "list_error", "detail": str(e), "op": "hydrate_skills"},
+        )
+        return 0
+
+    target_dir = target_dir or os.environ.get("LOKI_TARGET_DIR") or os.getcwd()
+    skills_dir = Path(target_dir) / ".loki" / "memory" / "skills"
+    skills_dir.mkdir(parents=True, exist_ok=True)
+
+    merged = 0
+    for e in entries:
+        content = e.get("content")
+        if not content:
+            continue
+        try:
+            skill = json.loads(content)
+        except (TypeError, json.JSONDecodeError):
+            continue
+        sid = skill.get("id") or skill.get("skill_id")
+        name = skill.get("name") or sid
+        if not name:
+            continue
+
+        # Sanitize filename (mirror MemoryStorage.save_skill).
+        safe_name = "".join(
+            c if c.isalnum() or c in "-_" else "_" for c in str(name)
+        )
+        skill_path = skills_dir / f"{safe_name}.json"
+        if skill_path.exists():
+            # Local wins on conflict.
+            continue
+
+        # Optional mtime gate.
+        ts = skill.get("updated_at") or skill.get("created_at")
+        if ts and local_mtime_floor:
+            try:
+                if isinstance(ts, (int, float)) and float(ts) < local_mtime_floor:
+                    continue
+            except (TypeError, ValueError):
+                pass
+
+        try:
+            from memory.storage import MemoryStorage  # type: ignore
+
+            storage = MemoryStorage(str(skills_dir.parent))
+            storage._atomic_write(skill_path, skill)
+        except Exception:
+            import tempfile
+
+            fd, tmp = tempfile.mkstemp(
+                dir=str(skills_dir), prefix=".tmp_", suffix=".json"
+            )
+            try:
+                with os.fdopen(fd, "w", encoding="utf-8") as f:
+                    json.dump(skill, f, indent=2, default=str)
+                os.replace(tmp, skill_path)
+            except Exception as ex:
+                if os.path.exists(tmp):
+                    os.unlink(tmp)
+                emit_managed_event(
+                    "managed_agents_fallback",
+                    {"reason": "atomic_write_failed", "detail": str(ex), "op": "hydrate_skills"},
+                )
+                continue
+        merged += 1
+
+    emit_managed_event(
+        "managed_memory_hydrate_skills",
+        {"merged": merged, "candidates": len(entries)},
+    )
+    return merged
+
+
+# ---------------------------------------------------------------------------
+# Session hydrate (patterns + skills) with idempotency guard
+# ---------------------------------------------------------------------------
+
+
+_HYDRATE_SENTINEL = ".loki/managed/hydrate.lock"
+
+
+def _already_hydrated_this_session(target_dir: str) -> bool:
+    """Idempotent: once we write the sentinel file, a second hydrate is no-op."""
+    sentinel = Path(target_dir) / _HYDRATE_SENTINEL
+    return sentinel.exists()
+
+
+def _mark_hydrated(target_dir: str) -> None:
+    sentinel = Path(target_dir) / _HYDRATE_SENTINEL
+    sentinel.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        sentinel.write_text(str(int(time.time())), encoding="utf-8")
+    except OSError:
+        pass
+
+
+def hydrate(
+    namespace: Optional[str] = None,
+    mtime_floor: Optional[float] = None,
+    target_dir: Optional[str] = None,
+) -> Dict[str, int]:
+    """
+    Session-boot hydrate: pull semantic patterns AND procedural skills from
+    the managed store and merge them into local .loki/memory/. Emits a single
+    `managed_memory_hydrate` event with counts.
+
+    Args:
+        namespace: Optional logical namespace label; reserved for multi-tenant
+            stores (not yet used by the backend). Included in the event for
+            observability.
+        mtime_floor: Only merge remote entries updated after this epoch
+            timestamp. Defaults to 0.0 (pull everything not already local).
+        target_dir: Override .loki root; defaults to LOKI_TARGET_DIR or cwd.
+
+    Returns:
+        {"patterns": N, "skills": M, "skipped": bool}. Disabled flags / errors
+        return {"patterns": 0, "skills": 0, "skipped": True/False}.
+
+    Idempotent: a second call within the same session (while the lock file
+    exists) short-circuits and returns zero counts with skipped=True.
+    """
+    target_dir = target_dir or os.environ.get("LOKI_TARGET_DIR") or os.getcwd()
+
+    if not is_enabled():
+        return {"patterns": 0, "skills": 0, "skipped": True}
+
+    if _already_hydrated_this_session(target_dir):
+        emit_managed_event(
+            "managed_memory_hydrate",
+            {
+                "patterns": 0,
+                "skills": 0,
+                "skipped": True,
+                "reason": "already_hydrated_this_session",
+                "namespace": namespace or "",
+            },
+        )
+        return {"patterns": 0, "skills": 0, "skipped": True}
+
+    floor = float(mtime_floor) if mtime_floor is not None else 0.0
+
+    patterns_merged = 0
+    skills_merged = 0
+    try:
+        patterns_merged = hydrate_patterns(
+            local_mtime_floor=floor, target_dir=target_dir
+        )
+    except Exception as e:  # pragma: no cover - defensive
+        emit_managed_event(
+            "managed_agents_fallback",
+            {"reason": "hydrate_patterns_error", "detail": str(e), "op": "hydrate"},
+        )
+    try:
+        skills_merged = hydrate_skills(
+            local_mtime_floor=floor, target_dir=target_dir
+        )
+    except Exception as e:  # pragma: no cover - defensive
+        emit_managed_event(
+            "managed_agents_fallback",
+            {"reason": "hydrate_skills_error", "detail": str(e), "op": "hydrate"},
+        )
+
+    _mark_hydrated(target_dir)
+
+    emit_managed_event(
+        "managed_memory_hydrate",
+        {
+            "patterns": patterns_merged,
+            "skills": skills_merged,
+            "skipped": False,
+            "namespace": namespace or "",
+        },
+    )
+    return {
+        "patterns": patterns_merged,
+        "skills": skills_merged,
+        "skipped": False,
+    }
+
+
 # ---------------------------------------------------------------------------
 # Module CLI
 # ---------------------------------------------------------------------------
@@ -319,7 +547,15 @@ def _main(argv: Optional[list] = None) -> int:
             floor = 0.0
             if args.since_seconds and args.since_seconds > 0:
                 floor = time.time() - args.since_seconds
-            hydrate_patterns(local_mtime_floor=floor)
+            # Phase 2: session-boot hydrate covers patterns + skills and is
+            # idempotent (sentinel-guarded). Prints a one-line summary to
+            # stdout so callers can log counts without parsing JSON.
+            result = hydrate(mtime_floor=floor)
+            print(
+                f"[managed] hydrate patterns={result.get('patterns', 0)} "
+                f"skills={result.get('skills', 0)} "
+                f"skipped={result.get('skipped', False)}"
+            )
             return 0
 
         query = args.query or ""

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "loki-mode",
-  "version": "6.83.1",
+  "version": "7.0.1",
   "description": "Loki Mode by Autonomi - Multi-agent autonomous startup system for Claude Code, Codex CLI, and Gemini CLI",
   "keywords": [
     "agent",
@@ -67,6 +67,7 @@
     "VERSION",
     "autonomy/",
     "providers/",
+    "agents/",
     "skills/",
     "references/",
     "docs/**/*.md",
@@ -105,7 +106,8 @@
     "test:visual": "node --experimental-vm-modules node_modules/jest/bin/jest.js dashboard-ui/tests/visual-regression.test.js",
     "test:parity": "node --experimental-vm-modules dashboard-ui/scripts/check-parity.js",
     "test:parity:json": "node --experimental-vm-modules dashboard-ui/scripts/check-parity.js --json",
-    "test:dashboard": "npm run test:visual && npm run test:parity"
+    "test:dashboard": "npm run test:visual && npm run test:parity",
+    "test:integration": "bash tests/integration/run_integration_suite.sh"
   },
   "engines": {
     "node": ">=18.0.0"