904 lines
31 KiB
Python
904 lines
31 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
import dataclasses
|
|
import decimal
|
|
import importlib
|
|
import json
|
|
import logging
|
|
import pathlib
|
|
import pickle
|
|
import re
|
|
import sys
|
|
from collections import deque
|
|
from collections.abc import Callable, Iterable, Sequence
|
|
from datetime import date, datetime, time, timedelta, timezone
|
|
from enum import Enum
|
|
from inspect import isclass
|
|
from ipaddress import (
|
|
IPv4Address,
|
|
IPv4Interface,
|
|
IPv4Network,
|
|
IPv6Address,
|
|
IPv6Interface,
|
|
IPv6Network,
|
|
)
|
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
|
from uuid import UUID
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import ormsgpack
|
|
from langchain_core.load.load import Reviver
|
|
|
|
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
|
|
from langgraph.checkpoint.serde.base import SerializerProtocol
|
|
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
|
|
from langgraph.checkpoint.serde.types import (
|
|
SendProtocol,
|
|
_DeltaSnapshot,
|
|
)
|
|
from langgraph.store.base import Item
|
|
|
|
if TYPE_CHECKING:
|
|
from langgraph.checkpoint.serde._msgpack import (
|
|
AllowedMsgpackModules,
|
|
)
|
|
|
|
LC_REVIVER = Reviver(allowed_objects="core")
|
|
EMPTY_BYTES = b""
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Dedup log warnings across process lifetime; cap bounds state if types are
|
|
# dynamically generated (also acts as a circuit breaker on warning volume).
|
|
# Dedup is best-effort: racing threads may each emit once for the same key,
|
|
# and warnings are silently dropped once _MAX_WARNED_TYPES is reached.
|
|
_MAX_WARNED_TYPES = 1000
|
|
_warned_unregistered_types: set[tuple[str, str]] = set()
|
|
_warned_blocked_types: set[tuple[str, str]] = set()
|
|
|
|
|
|
def _is_safe_json_type(id_list: list[str]) -> bool:
|
|
"""Return True if an lc=2 id refers to a type in SAFE_MSGPACK_TYPES.
|
|
|
|
Safe types bypass the ``allowed_json_modules`` gate so that old "json" format
|
|
checkpoints (written before the msgpack migration) can be resumed without
|
|
requiring users to configure an explicit allowlist.
|
|
"""
|
|
if len(id_list) < 2:
|
|
return False
|
|
module_name = ".".join(id_list[:-1])
|
|
return (module_name, id_list[-1]) in _lg_msgpack.SAFE_MSGPACK_TYPES
|
|
|
|
|
|
def _warn_once(
|
|
seen: set[tuple[str, str]], key: tuple[str, str], msg: str, *args: object
|
|
) -> None:
|
|
if key in seen or len(seen) >= _MAX_WARNED_TYPES:
|
|
return
|
|
seen.add(key)
|
|
logger.warning(msg, *args)
|
|
|
|
|
|
class JsonPlusSerializer(SerializerProtocol):
|
|
"""Serializer that uses ormsgpack, with optional fallbacks.
|
|
|
|
!!! warning
|
|
|
|
Security note: This serializer is intended for use within the `BaseCheckpointSaver`
|
|
class and called within the Pregel loop. It should not be used on untrusted
|
|
python objects. If an attacker can write directly to your checkpoint database,
|
|
they may be able to trigger code execution when data is deserialized.
|
|
|
|
Set the environment variable ``LANGGRAPH_STRICT_MSGPACK=true`` to restrict
|
|
deserialization to a built-in allowlist of safe types. You can also pass
|
|
an explicit ``allowed_msgpack_modules`` to the constructor.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
pickle_fallback: bool = False,
|
|
allowed_json_modules: Iterable[tuple[str, ...]] | Literal[True] | None = None,
|
|
allowed_msgpack_modules: (
|
|
AllowedMsgpackModules | Literal[True] | None
|
|
) = _lg_msgpack._SENTINEL,
|
|
__unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
|
|
) -> None:
|
|
if allowed_msgpack_modules is _lg_msgpack._SENTINEL:
|
|
if _lg_msgpack.STRICT_MSGPACK_ENABLED:
|
|
# Strict: only SAFE_MSGPACK_TYPES are allowed.
|
|
allowed_msgpack_modules = None
|
|
else:
|
|
# Permissive (default): all types allowed with a warning.
|
|
# Set LANGGRAPH_STRICT_MSGPACK=true to lock this down.
|
|
allowed_msgpack_modules = True
|
|
self.pickle_fallback = pickle_fallback
|
|
self._allowed_json_modules: set[tuple[str, ...]] | Literal[True] | None = (
|
|
_normalize_allowlist(allowed_json_modules)
|
|
)
|
|
self._allowed_msgpack_modules = _normalize_allowlist(allowed_msgpack_modules)
|
|
|
|
self._custom_unpack_ext_hook = __unpack_ext_hook__ is not None
|
|
self._unpack_ext_hook = (
|
|
__unpack_ext_hook__
|
|
if __unpack_ext_hook__ is not None
|
|
else _create_msgpack_ext_hook(self._allowed_msgpack_modules)
|
|
)
|
|
|
|
def with_msgpack_allowlist(
|
|
self, extra_allowlist: Iterable[tuple[str, ...] | type]
|
|
) -> JsonPlusSerializer:
|
|
"""Return a new serializer with a merged msgpack allowlist."""
|
|
base_allowlist = self._allowed_msgpack_modules
|
|
if base_allowlist is True or base_allowlist is False:
|
|
return self
|
|
elif base_allowlist:
|
|
base_allowlist = set(base_allowlist)
|
|
else:
|
|
base_allowlist = set()
|
|
extra = _normalize_module_keys(tuple(extra_allowlist))
|
|
merged = base_allowlist | extra
|
|
if merged == base_allowlist:
|
|
return self
|
|
allowed_msgpack_modules: AllowedMsgpackModules | Literal[True] | None
|
|
if merged:
|
|
allowed_msgpack_modules = tuple(merged)
|
|
elif isinstance(self._allowed_msgpack_modules, set):
|
|
allowed_msgpack_modules = tuple(self._allowed_msgpack_modules)
|
|
else:
|
|
allowed_msgpack_modules = self._allowed_msgpack_modules
|
|
|
|
clone = copy.copy(self)
|
|
clone._allowed_json_modules = _normalize_allowlist(self._allowed_json_modules)
|
|
clone._allowed_msgpack_modules = _normalize_allowlist(allowed_msgpack_modules)
|
|
if not clone._custom_unpack_ext_hook:
|
|
clone._unpack_ext_hook = _create_msgpack_ext_hook(
|
|
clone._allowed_msgpack_modules
|
|
)
|
|
return clone
|
|
|
|
def _encode_constructor_args(
|
|
self,
|
|
constructor: Callable | type[Any],
|
|
*,
|
|
method: None | str | Sequence[None | str] = None,
|
|
args: Sequence[Any] | None = None,
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
out = {
|
|
"lc": 2,
|
|
"type": "constructor",
|
|
"id": (*constructor.__module__.split("."), constructor.__name__),
|
|
}
|
|
if method is not None:
|
|
out["method"] = method
|
|
if args is not None:
|
|
out["args"] = args
|
|
if kwargs is not None:
|
|
out["kwargs"] = kwargs
|
|
return out
|
|
|
|
def _reviver(self, value: dict[str, Any]) -> Any:
|
|
if (
|
|
value.get("lc", None) == 2
|
|
and value.get("type", None) == "constructor"
|
|
and value.get("id", None) is not None
|
|
):
|
|
id_list = value["id"]
|
|
is_safe = _is_safe_json_type(id_list)
|
|
if self._allowed_json_modules or is_safe:
|
|
try:
|
|
return self._revive_lc2(value)
|
|
except InvalidModuleError as e:
|
|
if not is_safe:
|
|
logger.warning(
|
|
"Object %s is not in the deserialization allowlist.\n%s",
|
|
value["id"],
|
|
e.message,
|
|
)
|
|
|
|
return LC_REVIVER(value)
|
|
|
|
def _revive_lc2(self, value: dict[str, Any]) -> Any:
|
|
self._check_allowed_json_modules(value)
|
|
|
|
[*module, name] = value["id"]
|
|
try:
|
|
mod = importlib.import_module(".".join(module))
|
|
cls = getattr(mod, name)
|
|
method = value.get("method")
|
|
if isinstance(method, str):
|
|
methods = [getattr(cls, method)]
|
|
elif isinstance(method, list):
|
|
methods = [cls if m is None else getattr(cls, m) for m in method]
|
|
else:
|
|
methods = [cls]
|
|
args = value.get("args")
|
|
kwargs = value.get("kwargs")
|
|
for method in methods:
|
|
try:
|
|
if isclass(method) and issubclass(method, BaseException):
|
|
return None
|
|
if args and kwargs:
|
|
return method(*args, **kwargs)
|
|
elif args:
|
|
return method(*args)
|
|
elif kwargs:
|
|
return method(**kwargs)
|
|
else:
|
|
return method()
|
|
except Exception:
|
|
continue
|
|
except Exception:
|
|
return None
|
|
|
|
def _check_allowed_json_modules(self, value: dict[str, Any]) -> None:
|
|
needed = tuple(value["id"])
|
|
method = value.get("method")
|
|
if isinstance(method, list):
|
|
method_display = ",".join(m or "<init>" for m in method)
|
|
elif isinstance(method, str):
|
|
method_display = method
|
|
else:
|
|
method_display = "<init>"
|
|
|
|
dotted = ".".join(needed)
|
|
# Safe types (the same set already allowed for msgpack deserialization) are
|
|
# permitted without an explicit allowlist — they are known-safe LangGraph and
|
|
# LangChain types. This restores backwards-compat for old "json" checkpoints
|
|
# that pre-date the msgpack migration without reopening the broader security gate.
|
|
if _is_safe_json_type(list(needed)):
|
|
return
|
|
|
|
if not self._allowed_json_modules:
|
|
raise InvalidModuleError(
|
|
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
|
"No allowed_json_modules configured.\n\n"
|
|
"Unblock with ONE of:\n"
|
|
f" • JsonPlusSerializer(allowed_json_modules=[{needed!r}, ...])\n"
|
|
" • (DANGEROUS) JsonPlusSerializer(allowed_json_modules=True)\n\n"
|
|
"Note: Prefix allowlists are intentionally unsupported; prefer exact symbols "
|
|
"or plain-JSON representations revived without import-time side effects."
|
|
)
|
|
|
|
if self._allowed_json_modules is True:
|
|
return
|
|
if needed in self._allowed_json_modules:
|
|
return
|
|
|
|
raise InvalidModuleError(
|
|
f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
|
|
"Symbol is not in the deserialization allowlist.\n\n"
|
|
"Add exactly this symbol to unblock:\n"
|
|
f" JsonPlusSerializer(allowed_json_modules=[{needed!r}, ...])\n"
|
|
"Or, as a last resort (DANGEROUS):\n"
|
|
" JsonPlusSerializer(allowed_json_modules=True)"
|
|
)
|
|
|
|
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
|
|
if obj is None:
|
|
return "null", EMPTY_BYTES
|
|
elif isinstance(obj, bytes):
|
|
return "bytes", obj
|
|
elif isinstance(obj, bytearray):
|
|
return "bytearray", obj
|
|
else:
|
|
try:
|
|
return "msgpack", _msgpack_enc(obj)
|
|
except ormsgpack.MsgpackEncodeError as exc:
|
|
if self.pickle_fallback:
|
|
return "pickle", pickle.dumps(obj)
|
|
raise exc
|
|
|
|
def loads_typed(self, data: tuple[str, bytes]) -> Any:
|
|
type_, data_ = data
|
|
if type_ == "null":
|
|
return None
|
|
elif type_ == "bytes":
|
|
return data_
|
|
elif type_ == "bytearray":
|
|
return bytearray(data_)
|
|
elif type_ == "json":
|
|
return json.loads(data_, object_hook=self._reviver)
|
|
elif type_ == "msgpack":
|
|
return ormsgpack.unpackb(
|
|
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
elif self.pickle_fallback and type_ == "pickle":
|
|
return pickle.loads(data_)
|
|
else:
|
|
raise NotImplementedError(f"Unknown serialization type: {type_}")
|
|
|
|
|
|
# --- msgpack ---
|
|
|
|
EXT_CONSTRUCTOR_SINGLE_ARG = 0
|
|
EXT_CONSTRUCTOR_POS_ARGS = 1
|
|
EXT_CONSTRUCTOR_KW_ARGS = 2
|
|
EXT_METHOD_SINGLE_ARG = 3
|
|
EXT_PYDANTIC_V1 = 4
|
|
EXT_PYDANTIC_V2 = 5
|
|
EXT_NUMPY_ARRAY = 6
|
|
EXT_DELTA_SNAPSHOT = 7
|
|
|
|
|
|
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
|
|
if isinstance(obj, _DeltaSnapshot):
|
|
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
|
|
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
|
|
return ormsgpack.Ext(
|
|
EXT_PYDANTIC_V2,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj.model_dump(),
|
|
"model_validate_json",
|
|
),
|
|
),
|
|
)
|
|
elif hasattr(obj, "get_secret_value") and callable(obj.get_secret_value):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj.get_secret_value(),
|
|
),
|
|
),
|
|
)
|
|
elif hasattr(obj, "dict") and callable(obj.dict): # pydantic v1
|
|
return ormsgpack.Ext(
|
|
EXT_PYDANTIC_V1,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj.dict(),
|
|
),
|
|
),
|
|
)
|
|
elif hasattr(obj, "_asdict") and callable(obj._asdict): # namedtuple
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_KW_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj._asdict(),
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, pathlib.Path):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, obj.parts),
|
|
),
|
|
)
|
|
elif isinstance(obj, re.Pattern):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
("re", "compile", (obj.pattern, obj.flags)),
|
|
),
|
|
)
|
|
elif isinstance(obj, UUID):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, obj.hex),
|
|
),
|
|
)
|
|
elif isinstance(obj, decimal.Decimal):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, str(obj)),
|
|
),
|
|
)
|
|
elif isinstance(obj, (set, frozenset, deque)):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, tuple(obj)),
|
|
),
|
|
)
|
|
elif isinstance(obj, (IPv4Address, IPv4Interface, IPv4Network)):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, str(obj)),
|
|
),
|
|
)
|
|
elif isinstance(obj, (IPv6Address, IPv6Interface, IPv6Network)):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, str(obj)),
|
|
),
|
|
)
|
|
elif isinstance(obj, datetime):
|
|
return ormsgpack.Ext(
|
|
EXT_METHOD_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj.isoformat(),
|
|
"fromisoformat",
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, timedelta):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
(obj.days, obj.seconds, obj.microseconds),
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, date):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
(obj.year, obj.month, obj.day),
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, time):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_KW_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
{
|
|
"hour": obj.hour,
|
|
"minute": obj.minute,
|
|
"second": obj.second,
|
|
"microsecond": obj.microsecond,
|
|
"tzinfo": obj.tzinfo,
|
|
"fold": obj.fold,
|
|
},
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, timezone):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
obj.__getinitargs__(), # type: ignore[attr-defined]
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, ZoneInfo):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, obj.key),
|
|
),
|
|
)
|
|
elif isinstance(obj, Enum):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_SINGLE_ARG,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, obj.value),
|
|
),
|
|
)
|
|
elif isinstance(obj, SendProtocol):
|
|
args: tuple[Any, ...] = (obj.node, obj.arg)
|
|
if (timeout := getattr(obj, "timeout", None)) is not None:
|
|
args = (obj.node, obj.arg, timeout)
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_POS_ARGS,
|
|
_msgpack_enc(
|
|
(obj.__class__.__module__, obj.__class__.__name__, args),
|
|
),
|
|
)
|
|
elif dataclasses.is_dataclass(obj):
|
|
# doesn't use dataclasses.asdict to avoid deepcopy and recursion
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_KW_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
{
|
|
field.name: getattr(obj, field.name)
|
|
for field in dataclasses.fields(obj)
|
|
},
|
|
),
|
|
),
|
|
)
|
|
elif isinstance(obj, Item):
|
|
return ormsgpack.Ext(
|
|
EXT_CONSTRUCTOR_KW_ARGS,
|
|
_msgpack_enc(
|
|
(
|
|
obj.__class__.__module__,
|
|
obj.__class__.__name__,
|
|
{k: getattr(obj, k) for k in obj.__slots__},
|
|
),
|
|
),
|
|
)
|
|
elif (np_mod := sys.modules.get("numpy")) is not None and isinstance(
|
|
obj, np_mod.ndarray
|
|
):
|
|
order = "F" if obj.flags.f_contiguous and not obj.flags.c_contiguous else "C"
|
|
if obj.flags.c_contiguous:
|
|
mv = memoryview(obj)
|
|
try:
|
|
meta = (obj.dtype.str, obj.shape, order, mv)
|
|
return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
|
|
finally:
|
|
mv.release()
|
|
else:
|
|
buf = obj.tobytes(order="A")
|
|
meta = (obj.dtype.str, obj.shape, order, buf)
|
|
return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
|
|
|
|
elif isinstance(obj, BaseException):
|
|
return repr(obj)
|
|
else:
|
|
raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable")
|
|
|
|
|
|
def _send_from_args(args: Sequence[Any]) -> Any:
|
|
# ya we have a cyclic import here ¯\_(ツ)_/¯
|
|
from langgraph.types import Send # type: ignore
|
|
|
|
if len(args) == 2:
|
|
return Send(*args)
|
|
return Send(args[0], args[1], timeout=args[2])
|
|
|
|
|
|
def _create_msgpack_ext_hook(
|
|
allowed_modules: set[tuple[str, ...]] | Literal[True] | None,
|
|
) -> Callable[[int, bytes], Any]:
|
|
"""Create msgpack ext hook with allowlist.
|
|
|
|
Args:
|
|
allowed_modules: Set of (module, name) tuples that are allowed to be
|
|
deserialized, or True to allow all with warnings for unregistered types, or None to only allow safe types.
|
|
|
|
Returns:
|
|
An ext_hook function for use with ormsgpack.unpackb.
|
|
"""
|
|
|
|
def _check_allowed(module: str, name: str) -> bool:
|
|
"""Check if type is allowed. Returns True if allowed, False if blocked."""
|
|
key = (module, name)
|
|
|
|
if key in _lg_msgpack.SAFE_MSGPACK_TYPES:
|
|
return True
|
|
|
|
if allowed_modules is True:
|
|
# default is to warn but allow unregistered types
|
|
emit_serde_event(
|
|
{
|
|
"kind": "msgpack_unregistered_allowed",
|
|
"module": module,
|
|
"name": name,
|
|
}
|
|
)
|
|
_warn_once(
|
|
_warned_unregistered_types,
|
|
key,
|
|
"Deserializing unregistered type %s.%s from checkpoint. "
|
|
"This will be blocked in a future version. "
|
|
"Set LANGGRAPH_STRICT_MSGPACK=true to block now, or add "
|
|
"to allowed_msgpack_modules to allow explicitly: [(%r, %r)]",
|
|
module,
|
|
name,
|
|
module,
|
|
name,
|
|
)
|
|
return True
|
|
if allowed_modules is not None:
|
|
if key in allowed_modules:
|
|
return True
|
|
# strict mode blocks unregistered types
|
|
emit_serde_event(
|
|
{
|
|
"kind": "msgpack_blocked",
|
|
"module": module,
|
|
"name": name,
|
|
}
|
|
)
|
|
_warn_once(
|
|
_warned_blocked_types,
|
|
key,
|
|
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
|
|
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
|
|
module,
|
|
name,
|
|
module,
|
|
name,
|
|
)
|
|
return False
|
|
|
|
def _check_allowed_method(module: str, name: str, method: str) -> bool:
|
|
"""Check if a method invocation is allowed."""
|
|
key = (module, name, method)
|
|
if key in _lg_msgpack.SAFE_MSGPACK_METHODS:
|
|
return True
|
|
emit_serde_event(
|
|
{
|
|
"kind": "msgpack_method_blocked",
|
|
"module": module,
|
|
"name": name,
|
|
"method": method,
|
|
}
|
|
)
|
|
logger.warning(
|
|
"Blocked deserialization of method call %s.%s.%s - "
|
|
"not in allowed methods set.",
|
|
module,
|
|
name,
|
|
method,
|
|
)
|
|
return False
|
|
|
|
def ext_hook(code: int, data: bytes) -> Any:
|
|
if code == EXT_DELTA_SNAPSHOT:
|
|
return _DeltaSnapshot(
|
|
ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
)
|
|
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed(tup[0], tup[1]):
|
|
# We default to returning the raw data. If the user
|
|
# is using this in the context of a pydantic state, etc., then
|
|
# it would be validated upon construction.
|
|
return tup[2]
|
|
# module, name, arg
|
|
return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
|
|
except Exception:
|
|
return None
|
|
elif code == EXT_CONSTRUCTOR_POS_ARGS:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed(tup[0], tup[1]):
|
|
return tup[2]
|
|
if tup[0] == "langgraph.types" and tup[1] == "Send":
|
|
return _send_from_args(tup[2])
|
|
# module, name, args
|
|
return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
|
|
except Exception:
|
|
return None
|
|
elif code == EXT_CONSTRUCTOR_KW_ARGS:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed(tup[0], tup[1]):
|
|
return tup[2]
|
|
# module, name, kwargs
|
|
return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
|
|
except Exception:
|
|
return None
|
|
elif code == EXT_METHOD_SINGLE_ARG:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed_method(tup[0], tup[1], tup[3]):
|
|
return tup[2]
|
|
# module, name, arg, method
|
|
return getattr(
|
|
getattr(importlib.import_module(tup[0]), tup[1]), tup[3]
|
|
)(tup[2])
|
|
except Exception:
|
|
return None
|
|
elif code == EXT_PYDANTIC_V1:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed(tup[0], tup[1]):
|
|
return tup[2]
|
|
# module, name, kwargs
|
|
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
|
try:
|
|
return cls(**tup[2])
|
|
except Exception:
|
|
return cls.construct(**tup[2])
|
|
except Exception:
|
|
# for pydantic objects we can't find/reconstruct
|
|
# let's return the kwargs dict instead
|
|
try:
|
|
return tup[2]
|
|
except NameError:
|
|
return None
|
|
elif code == EXT_PYDANTIC_V2:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
if not _check_allowed(tup[0], tup[1]):
|
|
return tup[2]
|
|
# module, name, kwargs, method
|
|
cls = getattr(importlib.import_module(tup[0]), tup[1])
|
|
try:
|
|
return cls(**tup[2])
|
|
except Exception:
|
|
return cls.model_construct(**tup[2])
|
|
except Exception:
|
|
# for pydantic objects we can't find/reconstruct
|
|
# let's return the kwargs dict instead
|
|
try:
|
|
return tup[2]
|
|
except NameError:
|
|
return None
|
|
elif code == EXT_NUMPY_ARRAY:
|
|
try:
|
|
import numpy as _np
|
|
|
|
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
|
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
|
)
|
|
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
|
return arr.reshape(shape, order=order)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
return ext_hook
|
|
|
|
|
|
# Aliasing in case anyone imported it directly
|
|
_msgpack_ext_hook = _create_msgpack_ext_hook(allowed_modules=None)
|
|
|
|
|
|
def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:
|
|
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
if tup[0] == "uuid" and tup[1] == "UUID":
|
|
hex_ = tup[2]
|
|
return (
|
|
f"{hex_[:8]}-{hex_[8:12]}-{hex_[12:16]}-{hex_[16:20]}-{hex_[20:]}"
|
|
)
|
|
# module, name, arg
|
|
return tup[2]
|
|
except Exception:
|
|
return
|
|
elif code == EXT_CONSTRUCTOR_POS_ARGS:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
if tup[0] == "langgraph.types" and tup[1] == "Send":
|
|
return _send_from_args(tup[2])
|
|
# module, name, args
|
|
return tup[2]
|
|
except Exception:
|
|
return
|
|
elif code == EXT_CONSTRUCTOR_KW_ARGS:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
# module, name, args
|
|
return tup[2]
|
|
except Exception:
|
|
return
|
|
elif code == EXT_METHOD_SINGLE_ARG:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
# module, name, arg, method
|
|
return tup[2]
|
|
except Exception:
|
|
return
|
|
elif code == EXT_PYDANTIC_V1:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
# module, name, kwargs
|
|
return tup[2]
|
|
except Exception:
|
|
# for pydantic objects we can't find/reconstruct
|
|
# let's return the kwargs dict instead
|
|
return
|
|
elif code == EXT_PYDANTIC_V2:
|
|
try:
|
|
tup = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
# module, name, kwargs, method
|
|
return tup[2]
|
|
except Exception:
|
|
return
|
|
elif code == EXT_NUMPY_ARRAY:
|
|
try:
|
|
import numpy as _np
|
|
|
|
dtype_str, shape, order, buf = ormsgpack.unpackb(
|
|
data,
|
|
ext_hook=_msgpack_ext_hook_to_json,
|
|
option=ormsgpack.OPT_NON_STR_KEYS,
|
|
)
|
|
arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
|
|
return arr.reshape(shape, order=order).tolist()
|
|
except Exception:
|
|
return
|
|
|
|
|
|
class InvalidModuleError(Exception):
|
|
"""Exception raised when a module is not in the allowlist."""
|
|
|
|
def __init__(self, message: str):
|
|
self.message = message
|
|
|
|
|
|
_option = (
|
|
ormsgpack.OPT_NON_STR_KEYS
|
|
| ormsgpack.OPT_PASSTHROUGH_DATACLASS
|
|
| ormsgpack.OPT_PASSTHROUGH_DATETIME
|
|
| ormsgpack.OPT_PASSTHROUGH_ENUM
|
|
| ormsgpack.OPT_PASSTHROUGH_UUID
|
|
| ormsgpack.OPT_REPLACE_SURROGATES
|
|
)
|
|
|
|
|
|
def _msgpack_enc(data: Any) -> bytes:
|
|
return ormsgpack.packb(data, default=_msgpack_default, option=_option)
|
|
|
|
|
|
def _normalize_allowlist(
|
|
allowlist: AllowedMsgpackModules | Literal[True] | None,
|
|
) -> set[tuple[str, ...]] | Literal[True] | None:
|
|
if allowlist is True:
|
|
return allowlist
|
|
elif allowlist:
|
|
return _normalize_module_keys(allowlist)
|
|
else:
|
|
return None
|
|
|
|
|
|
def _normalize_module_keys(
|
|
modules: AllowedMsgpackModules,
|
|
) -> set[tuple[str, ...]]:
|
|
normalized: set[tuple[str, ...]] = set()
|
|
for module in modules:
|
|
if isclass(module):
|
|
normalized.add((module.__module__, module.__name__))
|
|
else:
|
|
normalized.add(cast(tuple[str, ...], module))
|
|
return normalized
|