loki-mode 7.78.0 → 7.80.0

package/lokistore/factory.py

@@ -0,0 +1,187 @@
+"""
+LokiStore factory: select and construct a backend from config or env.
+
+Defaults to local. With nothing configured, get_store() returns a LocalStore
+rooted at the project `.loki/` directory, honoring LOKI_DIR / TARGET_DIR
+exactly like the rest of the codebase. Cloud backends are constructed only when
+explicitly selected, and only then is their SDK imported.
+
+Config precedence: an explicit config dict overrides env, env overrides the
+local default. Recognized config keys mirror the env vars (without the
+LOKI_STORAGE_ prefix): backend, bucket, prefix, region, base_dir.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional
+
+from .base import LokiStore, StoreError
+from .local import LocalStore
+
+
+def resolve_local_base(explicit_base: Optional[str] = None) -> str:
+    """
+    Resolve the local `.loki/` base directory exactly like run.sh and the
+    dashboard do.
+
+    Resolution order, matching autonomy/run.sh
+    (`${LOKI_DIR:-${TARGET_DIR:-.}/.loki}`):
+      1. explicit_base argument, if given
+      2. $LOKI_DIR, if set (used verbatim, like the shell)
+      3. $TARGET_DIR/.loki, if TARGET_DIR is set
+      4. ./.loki  (current working directory)
+
+    Returns:
+        The resolved base directory as a string path. It is NOT created here;
+        LocalStore creates parents lazily on first write.
+    """
+    if explicit_base:
+        return explicit_base
+    loki_dir = os.environ.get("LOKI_DIR")
+    if loki_dir:
+        return loki_dir
+    target_dir = os.environ.get("TARGET_DIR", ".")
+    return os.path.join(target_dir, ".loki")
+
+
+def _config_from_env() -> Dict[str, Any]:
+    """Read the LOKI_STORAGE_* env vars into a config dict (omitting unset)."""
+    cfg: Dict[str, Any] = {}
+    backend = os.environ.get("LOKI_STORAGE_BACKEND")
+    if backend:
+        cfg["backend"] = backend
+    bucket = os.environ.get("LOKI_STORAGE_BUCKET")
+    if bucket:
+        cfg["bucket"] = bucket
+    prefix = os.environ.get("LOKI_STORAGE_PREFIX")
+    if prefix:
+        cfg["prefix"] = prefix
+    region = os.environ.get("LOKI_STORAGE_REGION")
+    if region:
+        cfg["region"] = region
+    return cfg
+
+
+def build_store(config: Optional[Dict[str, Any]] = None) -> LokiStore:
+    """
+    Construct a LokiStore from an explicit config dict (no env fallback).
+
+    Use this when a caller has already assembled config from its own source.
+    Prefer get_store() for the standard env-aware path.
+
+    Recognized keys:
+      backend : "local" (default) | "s3" | "gcs" | "azure-blob"
+      bucket  : bucket/container name (cloud backends)
+      prefix  : key prefix within the bucket (optional)
+      region  : region (s3, optional)
+      base_dir: local base directory (local backend only; overrides resolution)
+    """
+    config = dict(config or {})
+    backend = (config.get("backend") or "local").strip().lower()
+
+    if backend in ("local", "", "file", "filesystem"):
+        return LocalStore(resolve_local_base(config.get("base_dir")))
+
+    bucket = config.get("bucket")
+    prefix = config.get("prefix")
+    region = config.get("region")
+
+    if backend in ("s3", "aws", "aws-s3"):
+        from .cloud import S3Store
+
+        return S3Store(bucket=bucket, prefix=prefix, region=region)
+
+    if backend in ("gcs", "gcp", "google", "google-cloud-storage"):
+        from .cloud import GCSStore
+
+        return GCSStore(bucket=bucket, prefix=prefix, region=region)
+
+    if backend in ("azure-blob", "azure", "azureblob", "azure_blob"):
+        from .cloud import AzureBlobStore
+
+        return AzureBlobStore(bucket=bucket, prefix=prefix, region=region)
+
+    raise StoreError(
+        f"unknown storage backend: {backend!r} "
+        "(expected one of: local, s3, gcs, azure-blob)"
+    )
+
+
+def get_store(config: Optional[Dict[str, Any]] = None) -> LokiStore:
+    """
+    Return a LokiStore, defaulting to local.
+
+    Resolution:
+      1. Start from the LOKI_STORAGE_* env vars.
+      2. Overlay any explicit config dict (config keys win over env).
+      3. If no backend is selected, default to local.
+
+    With a clean environment and no config, this returns a LocalStore rooted at
+    the project `.loki/` -- zero new deps, zero behavior change.
+    """
+    merged = _config_from_env()
+    if config:
+        merged.update({k: v for k, v in config.items() if v is not None})
+    return build_store(merged)
+
+
+# ---------------------------------------------------------------------------
+# Metadata backend selector (stub for this release).
+#
+# Blobs (state, checkpoints, artifacts) go through LokiStore above. Structured
+# metadata (the dashboard's task/run index) lives in a relational store. Today
+# that is the existing async-SQLite database in dashboard/database.py at
+# $LOKI_DATA_DIR/dashboard.db. A postgres backend can be added later for shared
+# multi-instance fleets; the selector below is the seam for that, defaulting to
+# sqlite so there is no behavior change now.
+# ---------------------------------------------------------------------------
+
+def get_metadata_backend(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    """
+    Resolve the metadata (relational) backend selection.
+
+    This is intentionally a thin, documented descriptor rather than a live
+    connection: the existing dashboard/database.py owns the sqlite engine, and
+    a postgres implementation is deferred. Returning a descriptor lets future
+    callers branch on the choice without this module importing SQLAlchemy or
+    any database driver.
+
+    Selection (config["metadata_backend"] overrides
+    $LOKI_METADATA_BACKEND; default "sqlite"):
+      sqlite   -> reuse $LOKI_DATA_DIR/dashboard.db (default, implemented today
+                  in dashboard/database.py)
+      postgres -> read $LOKI_METADATA_URL (NOT YET IMPLEMENTED; selecting it
+                  returns the descriptor with implemented=False so callers can
+                  fail with a clear message until the impl lands)
+
+    Returns:
+        A descriptor dict: {backend, implemented, dsn, note}.
+    """
+    config = dict(config or {})
+    backend = (
+        config.get("metadata_backend")
+        or os.environ.get("LOKI_METADATA_BACKEND")
+        or "sqlite"
+    ).strip().lower()
+
+    if backend in ("sqlite", "", "default"):
+        data_dir = os.environ.get("LOKI_DATA_DIR", os.path.expanduser("~/.loki"))
+        return {
+            "backend": "sqlite",
+            "implemented": True,
+            "dsn": os.path.join(data_dir, "dashboard.db"),
+            "note": "reuses dashboard/database.py async-sqlite engine",
+        }
+
+    if backend in ("postgres", "postgresql", "pg"):
+        return {
+            "backend": "postgres",
+            "implemented": False,
+            "dsn": config.get("metadata_url") or os.environ.get("LOKI_METADATA_URL"),
+            "note": "postgres metadata backend is planned, not yet implemented",
+        }
+
+    raise StoreError(
+        f"unknown metadata backend: {backend!r} (expected: sqlite, postgres)"
+    )

package/lokistore/local.py

@@ -0,0 +1,219 @@
+"""
+LocalStore: the default, always-available, zero-dependency LokiStore backend.
+
+Backed by a base directory (the project `.loki/` by default). Behavior is
+byte-identical to today's direct `.loki/` file writes:
+
+- Atomic writes via a temp file in the same directory + os.replace (an atomic
+  rename on the same filesystem), so a reader never sees a torn file.
+- Best-effort fcntl advisory locking for concurrent writers of the same key,
+  reentrant per-thread to avoid self-deadlock, using PERSISTENT lock files
+  (never unlinked on release) to avoid the flock+unlink inode-replacement race
+  documented in memory/storage.py.
+- Path-traversal guard: a key can never escape the base directory.
+
+This mirrors the house idiom in memory/storage.py and dashboard/registry.py so
+local users get exactly the durability they have today, with no new deps.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+import threading
+from pathlib import Path
+from typing import List, Union
+
+from .base import LokiStore, normalize_key, read_source_bytes
+
+try:
+    import fcntl  # POSIX only; absent on Windows
+except ImportError:  # pragma: no cover - exercised only on Windows
+    fcntl = None  # type: ignore[assignment]
+
+
+class LocalStore(LokiStore):
+    """Filesystem-backed LokiStore rooted at a base directory."""
+
+    def __init__(self, base_dir: Union[str, os.PathLike]):
+        """
+        Args:
+            base_dir: Root directory for all keys. Created on first write.
+                      Typically the project `.loki/` directory; the factory
+                      resolves this honoring LOKI_DIR / TARGET_DIR.
+        """
+        self._base = Path(base_dir).expanduser()
+        # Reentrant lock tracking, mirroring memory/storage.py: a thread that
+        # already holds the lock for a path skips re-acquiring it, so nested
+        # operations on the same key do not deadlock.
+        self._held_locks: threading.local = threading.local()
+
+    @property
+    def base_dir(self) -> Path:
+        """The resolved root directory for this store."""
+        return self._base
+
+    # -- internal helpers ---------------------------------------------------
+
+    def _resolve(self, key: str) -> Path:
+        """
+        Map a normalized key to an absolute path inside the base dir, with a
+        realpath-based defense in depth so even a symlink inside the base
+        cannot redirect a read or write outside it.
+        """
+        clean = normalize_key(key)
+        full = self._base / clean
+
+        real_base = os.path.realpath(self._base)
+
+        def _under_base(real_path: str) -> bool:
+            return real_path == real_base or real_path.startswith(
+                real_base + os.sep
+            )
+
+        # Guard the parent so a key can never be created outside the base, even
+        # when the target itself does not exist yet (the write path).
+        real_parent = os.path.realpath(full.parent)
+        if not _under_base(real_parent):
+            raise ValueError(f"key escapes store base directory: {key!r}")
+
+        # Guard the full target too: a leaf symlink placed AT the key path and
+        # pointing outside the base would otherwise be followed on read
+        # (get/get_to/exists), leaking an arbitrary file. realpath resolves the
+        # leaf symlink, so a target whose real location is outside the base is
+        # rejected. This is a no-op for ordinary files and for not-yet-existing
+        # keys (whose realpath stays inside the already-checked parent).
+        real_full = os.path.realpath(full)
+        if not _under_base(real_full):
+            raise ValueError(f"key escapes store base directory: {key!r}")
+
+        return full
+
+    @contextlib.contextmanager
+    def _file_lock(self, path: Path):
+        """
+        Reentrant, best-effort exclusive advisory lock around a single key's
+        write. Uses a persistent sibling ".lock" file (never unlinked on
+        release) to avoid the flock+unlink inode race. Degrades to a no-op
+        where fcntl is unavailable (Windows); the atomic rename still
+        guarantees no torn reads, only lost-update protection is best-effort.
+        """
+        lock_path = path.with_suffix(path.suffix + ".lock")
+        lock_key = str(lock_path)
+
+        if not hasattr(self._held_locks, "paths"):
+            self._held_locks.paths = set()
+
+        # Reentrant: this thread already holds it -> no-op.
+        if lock_key in self._held_locks.paths:
+            yield
+            return
+
+        if fcntl is None:
+            # No advisory locking available; proceed (atomic rename still safe).
+            yield
+            return
+
+        lock_path.parent.mkdir(parents=True, exist_ok=True)
+        lock_file = None
+        try:
+            lock_file = open(lock_path, "w")
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+            self._held_locks.paths.add(lock_key)
+            yield
+        finally:
+            self._held_locks.paths.discard(lock_key)
+            if lock_file is not None:
+                try:
+                    fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+                finally:
+                    lock_file.close()
+                # Deliberately do NOT unlink the lock file (see module
+                # docstring and memory/storage.py for the inode-race rationale).
+
+    def _atomic_write_bytes(self, path: Path, payload: bytes) -> None:
+        """Write payload to path atomically (temp in same dir + os.replace)."""
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with self._file_lock(path):
+            fd, tmp_path = tempfile.mkstemp(
+                dir=str(path.parent), prefix=".tmp_", suffix=".part"
+            )
+            try:
+                with os.fdopen(fd, "wb") as f:
+                    f.write(payload)
+                    f.flush()
+                    os.fsync(f.fileno())
+                os.replace(tmp_path, str(path))
+            except BaseException:
+                with contextlib.suppress(OSError):
+                    os.unlink(tmp_path)
+                raise
+
+    # -- LokiStore interface ------------------------------------------------
+
+    def put(self, key: str, data: Union[bytes, bytearray, str, os.PathLike]) -> None:
+        payload = read_source_bytes(data)
+        path = self._resolve(key)
+        self._atomic_write_bytes(path, payload)
+
+    def get(self, key: str) -> bytes:
+        path = self._resolve(key)
+        if not path.is_file():
+            raise FileNotFoundError(f"no such key: {key!r}")
+        with self._file_lock(path):
+            with open(path, "rb") as f:
+                return f.read()
+
+    def get_to(self, key: str, dest_path: Union[str, os.PathLike]) -> None:
+        payload = self.get(key)
+        dest = Path(dest_path).expanduser()
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        # Reuse the atomic-write helper so the destination is never torn.
+        self._atomic_write_bytes(dest, payload)
+
+    def exists(self, key: str) -> bool:
+        return self._resolve(key).is_file()
+
+    def list(self, prefix: str = "") -> List[str]:
+        # Normalize the prefix to a path under the base. An empty prefix lists
+        # everything under the base.
+        if prefix:
+            clean_prefix = normalize_key(prefix)
+            search_root = self._base / clean_prefix
+        else:
+            search_root = self._base
+
+        if not search_root.exists():
+            return []
+
+        base_str = os.path.realpath(self._base)
+        results: List[str] = []
+
+        if search_root.is_file():
+            # Prefix pointed directly at a key.
+            rel = os.path.relpath(str(search_root), base_str).replace(os.sep, "/")
+            return [rel]
+
+        for root, _dirs, files in os.walk(search_root):
+            for name in files:
+                full = os.path.join(root, name)
+                # Skip internal lock and temp files so they never surface as keys.
+                if name.endswith(".lock") or name.startswith(".tmp_"):
+                    continue
+                rel = os.path.relpath(full, base_str).replace(os.sep, "/")
+                results.append(rel)
+
+        results.sort()
+        return results
+
+    def delete(self, key: str) -> bool:
+        path = self._resolve(key)
+        if not path.is_file():
+            return False
+        with self._file_lock(path):
+            try:
+                os.unlink(path)
+                return True
+            except FileNotFoundError:
+                return False

package/mcp/__init__.py

@@ -57,4 +57,4 @@ try:
 except ImportError:
     __all__ = ['mcp']
 
-__version__ = '7.78.0'
+__version__ = '7.80.0'

package/memory/retrieval.py

@@ -1483,6 +1483,153 @@ class MemoryRetrieval:
     # Private Helper Methods
     # -------------------------------------------------------------------------
 
+    # -------------------------------------------------------------------------
+    # Optional structure-aware tree retrieval (PageIndex pattern, OFF by default)
+    # -------------------------------------------------------------------------
+
+    def retrieve_tree(
+        self,
+        context: Dict[str, Any],
+        top_k: int = 5,
+        manifest: Optional[Dict[str, Any]] = None,
+        store: Optional[Any] = None,
+        llm: Optional[Callable[[str], str]] = None,
+    ) -> List[Dict[str, Any]]:
+        """Structure-aware tree retrieval over the code-index manifest.
+
+        Third, parallel, OPTIONAL retrieval path alongside keyword and vector.
+        It is NEVER reached unless a caller invokes it (directly or through the
+        LOKI_RETRIEVAL_MODE=tree dispatcher in retrieve_dispatch). The default
+        retrieve_task_aware path is byte-unchanged.
+
+        Builds (or loads from the LokiStore cache) a TOC tree from the code
+        index manifest, then reasons down it for the query. Degrades to a
+        deterministic keyword scorer when no LLM callable is available, and
+        further degrades to the existing keyword retrieval path when the
+        manifest itself is absent.
+
+        Args:
+            context: query context (goal, phase, action_type, files).
+            top_k: maximum number of results.
+            manifest: parsed code-index manifest. When None, it is loaded from
+                .loki/state/code-index-manifest.json (relative to the store).
+            store: a LokiStore for caching the built tree. When None, one is
+                built via lokistore.get_store() (local default, no new deps).
+            llm: optional LLM callable (prompt -> response) for reasoning
+                descent. When None, the keyword scorer is used.
+
+        Returns:
+            A ranked list of result dicts. Each carries "_source": "tree".
+            On any failure or a missing manifest, falls back to the existing
+            keyword retrieval so the caller always gets results.
+        """
+        # Local imports keep these optional modules off the default import path.
+        try:
+            from .tree_index import build_or_load_manifest_tree
+            from .tree_search import tree_search
+        except ImportError as exc:  # pragma: no cover - defensive
+            logger.warning("tree retrieval modules unavailable: %s", exc)
+            return self._tree_keyword_fallback(context, top_k)
+
+        if store is None:
+            try:
+                from lokistore import get_store
+
+                store = get_store()
+            except Exception as exc:  # noqa: BLE001 - degrade, never abort
+                logger.warning("could not obtain LokiStore for tree cache: %s", exc)
+                store = None
+
+        if manifest is None:
+            manifest = self._load_code_index_manifest(store)
+
+        if not manifest or not (manifest.get("files") or {}):
+            # No structure to reason over: fall back to keyword retrieval so
+            # the caller still gets results.
+            return self._tree_keyword_fallback(context, top_k)
+
+        query = self._build_query_from_context(context)
+
+        try:
+            if store is not None:
+                tree = build_or_load_manifest_tree(manifest, store)
+            else:
+                from .tree_index import build_tree_from_manifest
+
+                tree = build_tree_from_manifest(manifest)
+            return tree_search(tree, query, top_k=top_k, llm=llm)
+        except Exception as exc:  # noqa: BLE001 - degrade, never abort
+            logger.warning("tree retrieval failed (%s); using keyword fallback", exc)
+            return self._tree_keyword_fallback(context, top_k)
+
+    def retrieve_dispatch(
+        self,
+        context: Dict[str, Any],
+        top_k: int = 5,
+        token_budget: Optional[int] = None,
+        mode: Optional[str] = None,
+        **tree_kwargs: Any,
+    ) -> List[Dict[str, Any]]:
+        """Dispatch to a retrieval mode, defaulting to the existing path.
+
+        Mode resolution (first non-empty wins):
+          1. explicit `mode` argument
+          2. LOKI_RETRIEVAL_MODE env var
+          3. "task_aware" (the existing default path)
+
+        Only mode == "tree" diverges; every other value (including the default)
+        calls retrieve_task_aware UNCHANGED, so local devs who set nothing get
+        byte-identical behavior. Unknown modes also fall through to the default.
+        """
+        import os as _os
+
+        resolved = (mode or _os.environ.get("LOKI_RETRIEVAL_MODE") or "task_aware")
+        resolved = resolved.strip().lower()
+
+        if resolved == "tree":
+            return self.retrieve_tree(context, top_k=top_k, **tree_kwargs)
+
+        # Default and any unknown mode: existing behavior, untouched.
+        return self.retrieve_task_aware(
+            context, top_k=top_k, token_budget=token_budget
+        )
+
+    def _load_code_index_manifest(
+        self, store: Optional[Any]
+    ) -> Optional[Dict[str, Any]]:
+        """Load the code-index manifest, preferring the LokiStore.
+
+        Tries the store key "state/code-index-manifest.json" first (so it
+        honors LOKI_DIR / TARGET_DIR resolution), then a direct filesystem
+        read as a fallback. Returns None when no manifest is found.
+        """
+        manifest_key = "state/code-index-manifest.json"
+        if store is not None:
+            try:
+                if store.exists(manifest_key):
+                    raw = store.get(manifest_key)
+                    return json.loads(raw.decode("utf-8"))
+            except (FileNotFoundError, OSError, ValueError, UnicodeDecodeError):
+                pass
+        # Filesystem fallback relative to the configured base path.
+        candidate = Path(".loki/state/code-index-manifest.json")
+        try:
+            if candidate.is_file():
+                return json.loads(candidate.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            pass
+        return None
+
+    def _tree_keyword_fallback(
+        self, context: Dict[str, Any], top_k: int
+    ) -> List[Dict[str, Any]]:
+        """Fallback used by tree retrieval: the existing keyword path.
+
+        Reuses retrieve_task_aware so a tree-mode caller is never worse off
+        than the default mode when the manifest or an LLM is unavailable.
+        """
+        return self.retrieve_task_aware(context, top_k=top_k)
+
     def _build_query_from_context(self, context: Dict[str, Any]) -> str:
         """Build a query string from context dictionary."""
         parts = []