langgraph-runtime-inmem 0.36.0rc2__tar.gz → 0.37.0.dev2__tar.gz

{langgraph_runtime_inmem-0.36.0rc2 → langgraph_runtime_inmem-0.37.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-runtime-inmem
-Version: 0.36.0rc2
+Version: 0.37.0.dev2
 Summary: Inmem implementation for the LangGraph API server.
 Author-email: Will Fu-Hinthorn <will@langchain.dev>
 License: Elastic-2.0

{langgraph_runtime_inmem-0.36.0rc2 → langgraph_runtime_inmem-0.37.0.dev2}/langgraph_runtime_inmem/__init__.py

@@ -1,6 +1,7 @@
 from langgraph_runtime_inmem import (
     checkpoint,
     database,
+    encryption,
     lifespan,
     metrics,
     ops,
@@ -10,11 +11,12 @@ from langgraph_runtime_inmem import (
     store,
 )
 
-__version__ = "0.36.0rc2"
+__version__ = "0.37.0.dev2"
 __all__ = [
     "ops",
     "database",
     "checkpoint",
+    "encryption",
     "lifespan",
     "retry",
     "store",

langgraph_runtime_inmem-0.37.0.dev2/langgraph_runtime_inmem/_persistence.py

@@ -0,0 +1,88 @@
+"""Periodic flushing for all PersistentDict stores."""
+
+from __future__ import annotations
+
+import functools
+import logging
+import os
+import threading
+import weakref
+
+from langgraph.checkpoint.memory import PersistentDict
+
+logger = logging.getLogger(__name__)
+
+_stores: dict[str, weakref.ref[PersistentDict]] = {}
+# Held while changing _stores and while syncing one entry, so once
+# unregister_persistent_dict() returns no sync of that dict is running or will start.
+_stores_lock = threading.Lock()
+_flush_thread: tuple[threading.Event, threading.Thread] | None = None
+_flush_interval: int = 10
+DISABLE_FILE_PERSISTENCE = (
+    os.getenv("LANGGRAPH_DISABLE_FILE_PERSISTENCE", "false").lower() == "true"
+)
+
+
+def register_persistent_dict(d: PersistentDict) -> None:
+    """Register a PersistentDict for periodic flushing."""
+    if DISABLE_FILE_PERSISTENCE:
+        return
+    global _flush_thread
+    with _stores_lock:
+        _stores[d.filename] = weakref.ref(d)
+        if _flush_thread is None:
+            logger.info("Starting dev persistence flush loop")
+            stop_event = threading.Event()
+            _flush_thread = (
+                stop_event,
+                threading.Thread(
+                    target=functools.partial(_flush_loop, stop_event), daemon=True
+                ),
+            )
+            _flush_thread[1].start()
+
+
+def unregister_persistent_dict(d: PersistentDict) -> None:
+    """Stop periodic flushing of ``d`` if it's still the dict registered for its file."""
+    with _stores_lock:
+        ref = _stores.get(d.filename)
+        if ref is not None and ref() is d:
+            _stores.pop(d.filename, None)
+
+
+def stop_flush_loop() -> None:
+    """Stop the background flush thread."""
+    global _flush_thread
+    if _flush_thread is not None:
+        logger.info("Stopping dev persistence flush loop")
+        _flush_thread[0].set()
+        _flush_thread[1].join()
+        _flush_thread = None
+
+
+def _flush_once() -> None:
+    with _stores_lock:
+        entries = list(_stores.items())
+    for filename, ref in entries:
+        with _stores_lock:
+            # Skip entries unregistered or replaced since the snapshot.
+            if _stores.get(filename) is not ref:
+                continue
+            # Compare to None: an empty dict is falsy but still needs flushing.
+            if (store := ref()) is None:
+                _stores.pop(filename, None)
+                logger.debug(
+                    "Dropped garbage-collected dev persistence file %s", filename
+                )
+                continue
+            try:
+                store.sync()
+            except Exception:
+                # One bad file must not stop the others (or later retries) from flushing.
+                logger.exception("Failed to flush dev persistence file %s", filename)
+
+
+def _flush_loop(stop_event: threading.Event) -> None:
+    while not stop_event.wait(timeout=_flush_interval):
+        _flush_once()
+    logger.info("dev persistence flush loop exiting")

{langgraph_runtime_inmem-0.36.0rc2 → langgraph_runtime_inmem-0.37.0.dev2}/langgraph_runtime_inmem/checkpoint.py

@@ -54,6 +54,7 @@ class InMemorySaver(InMemorySaverBase):
         self,
         *,
         serde: SerializerProtocol | None = None,
+        persist: bool = True,
         __persistence_hook__: Callable[[PersistentDict], None] | None = None,
     ) -> None:
         self.filename = os.path.join(".langgraph_api", ".langgraph_checkpoint.")
@@ -68,8 +69,6 @@ class InMemorySaver(InMemorySaverBase):
                 os.mkdir(".langgraph_api")
             thisfname = self.filename + str(i) + ".pckl"
             d = PersistentDict(*args, filename=thisfname)
-            if __persistence_hook__:
-                __persistence_hook__(d)
 
             try:
                 d.load()
@@ -94,13 +93,21 @@ class InMemorySaver(InMemorySaverBase):
                     os.remove(self.filename)
                 except Exception:
                     pass
+            # Register only once loaded, so a flush can't write the still-empty
+            # dict over the file.
+            if __persistence_hook__:
+                __persistence_hook__(d)
             return d
 
         from langgraph_api.serde import Serializer  # noqa: PLC0415
 
+        # persist=False: the saver borrows another saver's dicts instead of owning
+        # files. Owning them would load the pickles and take over their flush
+        # registrations (keyed by filename).
+        use_files = persist and not DISABLE_FILE_PERSISTENCE
         super().__init__(
             serde=serde if serde is not None else Serializer(),
-            factory=factory if not DISABLE_FILE_PERSISTENCE else defaultdict,
+            factory=factory if use_files else defaultdict,
         )
 
     def put(
@@ -286,9 +293,11 @@ def Checkpointer(*args, unpack_hook=None, **kwargs):
         else:
             ext_hook = unpack_hook
 
+        # A separate saver for its JSON-friendly serde and per-request latest_iter;
+        # the data itself is MEMORY's.
         saver = InMemorySaver(
             serde=Serializer(__unpack_ext_hook__=ext_hook),
-            __persistence_hook__=register_persistent_dict,
+            persist=False,
             **kwargs,
         )
         saver.writes = MEMORY.writes

{langgraph_runtime_inmem-0.36.0rc2 → langgraph_runtime_inmem-0.37.0.dev2}/langgraph_runtime_inmem/database.py

@@ -1,10 +1,13 @@
 import asyncio
+import contextlib
+import glob
 import os
+import pickle
 import uuid
 from collections import defaultdict
 from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import TYPE_CHECKING, Any, NotRequired
 from uuid import UUID
 
@@ -154,6 +157,96 @@ async def connect(
     yield InMemConnectionProto()
 
 
+def _load_ops_file() -> None:
+    """Load the saved ops file into ``GLOBAL_STORE``, setting it aside if unloadable."""
+    if not os.path.exists(OPS_FILENAME):
+        return
+
+    # Not GLOBAL_STORE.load(): it treats a truncated file as empty.
+    try:
+        with open(OPS_FILENAME, "rb") as f:
+            data = pickle.load(f)
+        if not isinstance(data, dict):
+            raise TypeError(f"expected a dict, found {type(data).__name__}")
+    except Exception as e:
+        if os.path.getsize(OPS_FILENAME) == 0:
+            return
+        _set_aside_ops_file(e)
+        return
+    GLOBAL_STORE.update(data)
+
+
+def _set_aside_ops_file(reason: Exception) -> None:
+    """Rename an unloadable ops file to ``<file>.bak-<UTC timestamp>`` and log why."""
+    from langgraph_api.graph import graph_file_module_names  # noqa: PLC0415
+
+    stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
+    backup = f"{OPS_FILENAME}.bak-{stamp}"
+    n = 1
+    while os.path.exists(backup):
+        backup = f"{OPS_FILENAME}.bak-{stamp}-{n}"
+        n += 1
+    # Don't catch: carrying on would let the next flush overwrite the file.
+    os.replace(OPS_FILENAME, backup)
+    with contextlib.suppress(FileNotFoundError):
+        os.remove(RETRY_COUNTER_FILENAME)
+    try:
+        with open(f"{backup}.reason", "w") as f:
+            f.write(f"{type(reason).__name__}: {reason}")
+    except OSError as e:
+        logger.warning("Could not save why %s was set aside: %s", backup, e)
+
+    message = (
+        f"Could not load saved threads, runs and crons from {OPS_FILENAME}\n"
+        f"  Reason: {type(reason).__name__}: {reason}\n"
+        "  The server is starting without them. The file was NOT deleted; it was moved to:\n"
+        f"    {backup}\n"
+    )
+    if (
+        isinstance(reason, ModuleNotFoundError)
+        and reason.name in graph_file_module_names()
+    ):
+        message += (
+            f"  It references a class from your graph file (module '{reason.name}'), saved by an older\n"
+            "  version of langgraph-api. It can't be restored; newer versions no longer save\n"
+            "  data this way, so this won't happen again."
+        )
+    else:
+        message += (
+            "  To restore it: fix the cause, stop the server, and rename that file back to\n"
+            f"  {os.path.basename(OPS_FILENAME)}. Threads created since this start will be lost.\n"
+            "  This usually means a class stored in thread state was renamed, moved, or\n"
+            "  failed to import (e.g. a syntax error in its module)."
+        )
+    logger.error(message, error_type=type(reason).__name__, backup=backup)
+
+
+def _warn_about_ops_backups() -> None:
+    """Log each ops-file backup left by an earlier start, with why it was set aside."""
+    paths = glob.glob(f"{glob.escape(OPS_FILENAME)}.bak-*")
+    backups = sorted(p for p in paths if not p.endswith(".reason"))
+    for reason in paths:
+        if reason.endswith(".reason") and reason.removesuffix(".reason") not in backups:
+            os.remove(reason)
+    if not backups:
+        return
+
+    lines = []
+    for backup in backups:
+        try:
+            with open(f"{backup}.reason") as f:
+                cause = f.read()
+        except OSError:
+            cause = "unknown cause"
+        lines.append(f"  {backup}: {cause}")
+    logger.warning(
+        f"{len(backups)} unloadable ops file backup(s). Rename one to "
+        f"{os.path.basename(OPS_FILENAME)} once its cause is fixed, or delete it:\n"
+        + "\n".join(lines),
+        backups=backups,
+    )
+
+
 async def start_pool() -> None:
     if store._STORE_CONFIG is None:
         from langgraph_api import config as langgraph_config  # noqa: PLC0415
@@ -164,24 +257,8 @@ async def start_pool() -> None:
 
     if not os.path.exists(".langgraph_api"):
         os.mkdir(".langgraph_api")
-    if os.path.exists(OPS_FILENAME):
-        try:
-            GLOBAL_STORE.load()
-        except ModuleNotFoundError:
-            logger.error(
-                "Unable to load cached data - your code has changed in a way that's incompatible with the cache."
-                "\nThis usually happens when you've:"
-                "\n  - Renamed or moved classes"
-                "\n  - Changed class structures"
-                "\n  - Pulled updates that modified class definitions in a way that's incompatible with the cache"
-                "\n\nRemoving invalid cache data stored at path: .langgraph_api"
-            )
-            await asyncio.to_thread(os.remove, OPS_FILENAME)
-            await asyncio.to_thread(os.remove, RETRY_COUNTER_FILENAME)
-        except Exception as e:
-            logger.error("Failed to load cached data: %s", str(e))
-            await asyncio.to_thread(os.remove, OPS_FILENAME)
-            await asyncio.to_thread(os.remove, RETRY_COUNTER_FILENAME)
+    _warn_about_ops_backups()
+    _load_ops_file()
     for k in ["runs", "threads", "assistant_versions", "assistants"]:
         if not GLOBAL_STORE.get(k):
             GLOBAL_STORE[k] = []
@@ -201,6 +278,7 @@ async def start_pool() -> None:
 
 
 async def stop_pool() -> None:
+    # Stop flushing before closing: close() clears the dicts, and a later flush would write {}.
     stop_flush_loop()
     await asyncio.to_thread(GLOBAL_STORE.close)
     await asyncio.to_thread(GLOBAL_RETRY_COUNTER.close)

langgraph_runtime_inmem-0.37.0.dev2/langgraph_runtime_inmem/encryption.py

@@ -0,0 +1,232 @@
+"""Structural JSON encryption owned by the Python in-memory runtime.
+
+Unlike Postgres persistence, this remains here while local persistence is
+Python-owned.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from langgraph_api.encryption.aes_json import AesEncryptionInstance
+    from langgraph_api.encryption.custom import JsonEncryptionWrapper, ModelType
+
+NESTED_ENCRYPTED_SUBFIELDS: dict[tuple[str, str], tuple[str, ...]] = {
+    ("run", "kwargs"): ("input", "config", "context", "command"),
+    ("run", "config"): ("configurable", "metadata"),
+    ("cron", "payload"): ("metadata", "context", "input", "config"),
+    ("cron", "config"): ("configurable", "metadata"),
+    ("assistant", "config"): ("configurable",),
+    ("thread", "config"): ("configurable",),
+}
+
+ENCRYPTION_FIELDS: dict[str, tuple[str, ...]] = {
+    "thread": ("metadata", "config", "values", "interrupts", "error"),
+    "run": ("metadata", "kwargs"),
+    "assistant": ("metadata", "config", "context"),
+    "cron": ("metadata", "payload"),
+    "store": ("value",),
+}
+
+NEVER_ENCRYPT_FIELDS = frozenset(
+    {
+        "thread_id",
+        "run_id",
+        "assistant_id",
+        "graph_id",
+        "checkpoint_id",
+        "task_id",
+        "__pregel_checkpointer",
+        "__pregel_resuming",
+        "__pregel_durability",
+        "__pregel_stream",
+        "__pregel_task_id",
+        "__pregel_checkpoint_ns",
+        "__after_seconds__",
+        "__request_start_time_ms__",
+        "__encryption_context__",
+        "__blob_encryption_context__",
+        "langgraph_version",
+        "langgraph_api_version",
+        "langgraph_plan",
+        "langgraph_host",
+        "langgraph_api_url",
+        "langgraph_request_id",
+        "langgraph_auth_user_id",
+        "langgraph_auth_permissions",
+    }
+)
+
+NEVER_ENCRYPT_PATHS = frozenset(
+    {
+        "run.kwargs.config.recursion_limit",
+        "run.kwargs.config.max_concurrency",
+        "run.kwargs.temporary",
+        "run.kwargs.config.configurable.ttl",
+        "checkpoint.metadata.source",
+        "checkpoint.metadata.step",
+        "checkpoint.metadata.parents",
+        "checkpoint.metadata.run_attempt",
+        "checkpoint.metadata.counters_since_delta_snapshot",
+    }
+)
+
+
+def should_skip_encryption(key: str, path: str) -> bool:
+    return key in NEVER_ENCRYPT_FIELDS or f"{path}.{key}" in NEVER_ENCRYPT_PATHS
+
+
+def extract_blob_encryption_context(
+    data: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+    """Return checkpoint blob context carried by an encrypted resource."""
+    if data is None:
+        return None
+
+    from langgraph_api.encryption.shared import (  # noqa: PLC0415
+        BLOB_ENCRYPTION_CONTEXT_KEY,
+        ENCRYPTION_CONTEXT_KEY,
+    )
+
+    return data.get(BLOB_ENCRYPTION_CONTEXT_KEY) or data.get(ENCRYPTION_CONTEXT_KEY)
+
+
+def encryption_fields(model_type: ModelType) -> list[str]:
+    """Return fields structurally encrypted by the local runtime."""
+    return list(ENCRYPTION_FIELDS[model_type])
+
+
+async def _decrypt_field(
+    obj: dict[str, Any],
+    field_name: str,
+    encryption_instance: JsonEncryptionWrapper | AesEncryptionInstance | None,
+    model_type: ModelType,
+) -> tuple[str, Any]:
+    from langgraph_api.encryption.middleware import (  # noqa: PLC0415
+        decrypt_json_if_needed,
+    )
+
+    if not obj.get(field_name):
+        return field_name, obj.get(field_name)
+
+    decrypted = await decrypt_json_if_needed(
+        obj[field_name], encryption_instance, model_type, field=field_name
+    )
+    nested_fields = NESTED_ENCRYPTED_SUBFIELDS.get((model_type, field_name), ())
+    if decrypted is not None:
+        results = await asyncio.gather(
+            *(
+                _decrypt_field(decrypted, name, encryption_instance, model_type)
+                for name in nested_fields
+                if isinstance(decrypted.get(name), dict)
+            )
+        )
+        for name, value in results:
+            decrypted[name] = value
+    return field_name, decrypted
+
+
+async def decrypt_fields(
+    obj: dict[str, Any],
+    model_type: ModelType,
+    fields: list[str],
+    encryption_instance: JsonEncryptionWrapper | AesEncryptionInstance | None,
+) -> None:
+    """Decrypt configured fields, including local-runtime nested structure."""
+    results = await asyncio.gather(
+        *(
+            _decrypt_field(obj, name, encryption_instance, model_type)
+            for name in fields
+            if name in obj
+        )
+    )
+    for name, value in results:
+        obj[name] = value
+
+
+async def _encrypt_field(
+    data: Mapping[str, Any],
+    field_name: str,
+    encryption_instance: JsonEncryptionWrapper | AesEncryptionInstance | None,
+    model_type: ModelType,
+    path: str | None = None,
+) -> tuple[str, Any]:
+    from langgraph_api.encryption.middleware import (  # noqa: PLC0415
+        encrypt_json_if_needed,
+    )
+
+    if field_name not in data or data[field_name] is None:
+        return field_name, data.get(field_name)
+
+    field_data = data[field_name]
+    if model_type == "thread" and field_name == "values":
+        from langgraph_api.serde import json_dumpb, json_loads  # noqa: PLC0415
+
+        field_data = json_loads(json_dumpb(field_data))
+    current_path = f"{path}.{field_name}" if path else f"{model_type}.{field_name}"
+    nested_fields = NESTED_ENCRYPTED_SUBFIELDS.get((model_type, field_name), ())
+    extracted: dict[str, Any] = {}
+    if nested_fields:
+        if not isinstance(field_data, dict):
+            raise TypeError(
+                f"'{field_name}' must be a dict for encryption, "
+                f"got {type(field_data).__name__}"
+            )
+        for name in nested_fields:
+            value = field_data.get(name)
+            if isinstance(value, dict) and value:
+                extracted[name] = value
+        if extracted:
+            field_data = {k: v for k, v in field_data.items() if k not in extracted}
+
+    encrypted = await encrypt_json_if_needed(
+        field_data,
+        encryption_instance,
+        model_type,
+        field=field_name,
+        path=current_path,
+        should_skip=should_skip_encryption,
+    )
+    if extracted and isinstance(encrypted, dict):
+        results = await asyncio.gather(
+            *(
+                _encrypt_field(
+                    {name: value}, name, encryption_instance, model_type, current_path
+                )
+                for name, value in extracted.items()
+            )
+        )
+        for name, value in results:
+            encrypted[name] = value
+            marker = value.get("__encryption_context__")
+            if (
+                "__encryption_context__" not in encrypted
+                and isinstance(marker, dict)
+                and marker.get("__langgraph_encryption_type__") == "aes"
+            ):
+                encrypted["__encryption_context__"] = marker
+    return field_name, encrypted
+
+
+async def encrypt_fields(
+    data: Mapping[str, Any],
+    model_type: ModelType,
+    fields: list[str],
+    encryption_instance: JsonEncryptionWrapper | AesEncryptionInstance | None,
+) -> dict[str, Any]:
+    """Encrypt configured fields, preserving local-runtime nested structure."""
+    result = dict(data)
+    encrypted_fields = await asyncio.gather(
+        *(
+            _encrypt_field(data, name, encryption_instance, model_type)
+            for name in fields
+            if name in data
+        )
+    )
+    for name, value in encrypted_fields:
+        result[name] = value
+    return result