codekingpro/portable-devtools

1"""Custom encryption support for LangGraph.2 3.. warning::4    This API is in beta and may change in future versions.5 6This module provides a framework for implementing custom at-rest encryption7in LangGraph applications. Similar to the Auth system, it allows developers8to define custom encryption and decryption handlers that are executed9server-side.10"""11 12from __future__ import annotations13 14import functools15import inspect16import typing17import warnings18 19from langgraph_sdk.encryption import types20 21 22class LangGraphBetaWarning(UserWarning):23    """Warning for beta features in LangGraph SDK."""24 25 26@functools.lru_cache(maxsize=1)27def _warn_encryption_beta() -> None:28    warnings.warn(29        "The Encryption API is in beta and may change in future versions.",30        LangGraphBetaWarning,31        stacklevel=4,32    )33 34 35class DuplicateHandlerError(Exception):36    """Raised when attempting to register a duplicate encryption/decryption handler."""37 38    pass39 40 41def _validate_handler(fn: typing.Callable, handler_type: str) -> None:42    """Validate that a handler function has the correct signature.43 44    Args:45        fn: The handler function to validate46        handler_type: Description of the handler for error messages47 48    Raises:49        TypeError: If the handler is not an async function or has wrong parameter count50    """51    if not inspect.iscoroutinefunction(fn):52        raise TypeError(f"{handler_type} must be an async function, got {type(fn)}")53 54    sig = inspect.signature(fn)55    params = [56        p57        for p in sig.parameters.values()58        if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)59    ]60    if len(params) != 2:61        raise TypeError(62            f"{handler_type} must accept exactly 2 parameters "63            f"(ctx, data), got {len(params)}"64        )65 66 67class _EncryptDecorators:68    """Decorators for encryption handlers.69 70    Provides @encryption.encrypt.blob and @encryption.encrypt.json decorators for71    registering encryption functions.72    """73 74    def __init__(self, parent: Encryption):75        self._parent = parent76 77    def blob(self, fn: types.BlobEncryptor) -> types.BlobEncryptor:78        """Register a blob encryption handler.79 80        The handler will be called to encrypt opaque data like checkpoint blobs.81 82        Example:83            ```python84            @encryption.encrypt.blob85            async def encrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:86                # Encrypt the blob using your encryption service87                return encrypted_blob88            ```89 90        Args:91            fn: The encryption handler function92 93        Returns:94            The registered handler function95 96        Raises:97            DuplicateHandlerError: If blob encryptor already registered98            TypeError: If handler has invalid signature99        """100        if self._parent._blob_encryptor is not None:101            raise DuplicateHandlerError("Blob encryptor already registered")102        _validate_handler(fn, "Blob encryptor")103        self._parent._blob_encryptor = fn104        return fn105 106    def json(self, fn: types.JsonEncryptor) -> types.JsonEncryptor:107        """Register the JSON encryption handler.108 109        Example:110            ```python111            @encryption.encrypt.json112            async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:113                # Encrypt the data114                return encrypt_data(data)115            ```116 117        Args:118            fn: The encryption handler function119 120        Returns:121            The registered handler function122 123        Raises:124            DuplicateHandlerError: If JSON encryptor already registered125            TypeError: If handler has invalid signature126        """127        if self._parent._json_encryptor is not None:128            raise DuplicateHandlerError("JSON encryptor already registered")129        _validate_handler(fn, "JSON encryptor")130        self._parent._json_encryptor = fn131        return fn132 133 134class _DecryptDecorators:135    """Decorators for decryption handlers.136 137    Provides @encryption.decrypt.blob and @encryption.decrypt.json decorators for138    registering decryption functions.139    """140 141    def __init__(self, parent: Encryption):142        self._parent = parent143 144    def blob(self, fn: types.BlobDecryptor) -> types.BlobDecryptor:145        """Register a blob decryption handler.146 147        The handler will be called to decrypt opaque data like checkpoint blobs.148 149        Example:150            ```python151            @encryption.decrypt.blob152            async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:153                # Decrypt the blob using your encryption service154                return decrypted_blob155            ```156 157        Args:158            fn: The decryption handler function159 160        Returns:161            The registered handler function162 163        Raises:164            DuplicateHandlerError: If blob decryptor already registered165            TypeError: If handler has invalid signature166        """167        if self._parent._blob_decryptor is not None:168            raise DuplicateHandlerError("Blob decryptor already registered")169        _validate_handler(fn, "Blob decryptor")170        self._parent._blob_decryptor = fn171        return fn172 173    def json(self, fn: types.JsonDecryptor) -> types.JsonDecryptor:174        """Register the JSON decryption handler.175 176        Example:177            ```python178            @encryption.decrypt.json179            async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:180                # Decrypt the data181                return decrypt_data(data)182            ```183 184        Args:185            fn: The decryption handler function186 187        Returns:188            The registered handler function189 190        Raises:191            DuplicateHandlerError: If JSON decryptor already registered192            TypeError: If handler has invalid signature193        """194        if self._parent._json_decryptor is not None:195            raise DuplicateHandlerError("JSON decryptor already registered")196        _validate_handler(fn, "JSON decryptor")197        self._parent._json_decryptor = fn198        return fn199 200 201class Encryption:202    """Add custom at-rest encryption to your LangGraph application.203 204    .. warning::205        This API is in beta and may change in future versions.206 207    The Encryption class provides a system for implementing custom encryption208    of data at rest in LangGraph applications. It supports encryption of209    both opaque blobs (like checkpoints) and structured JSON data (like210    metadata, context, kwargs, values, etc.).211 212    To use, create a separate Python file and add the path to the file to your213    LangGraph API configuration file (`langgraph.json`). Within that file, create214    an instance of the Encryption class and register encryption and decryption215    handlers as needed.216 217    Example `langgraph.json` file:218 219    ```json220    {221      "dependencies": ["."],222      "graphs": {223        "agent": "./my_agent/agent.py:graph"224      },225      "env": ".env",226      "encryption": {227        "path": "./encryption.py:my_encryption"228      }229    }230    ```231 232    Then the LangGraph server will load your encryption file and use it to233    encrypt/decrypt data at rest.234 235    !!! warning "JSON Encryptors Must Preserve Keys"236 237        JSON encryptors **must not add or remove keys** from the input dict.238        Only values may be transformed. This constraint is **enforced at runtime239        by the server** and exists because SQL JSONB merge operations (used for240        partial updates) work at the key level.241 242        **Correct (per-key encryption):**243        ```python244        # Input:  {"secret": "value", "plain": "x"}245        # Output: {"secret": "<encrypted>", "plain": "x"}  ✓ Keys preserved246        ```247 248        **Incorrect (key consolidation):**249        ```python250        # Input:  {"secret": "value", "plain": "x"}251        # Output: {"__encrypted__": "<blob>", "plain": "x"}  ✗ Key changed252        ```253 254        If your encryptor needs to store auxiliary data (DEK, IV, etc.), embed it255        within the encrypted value itself, not as separate keys.256 257    ???+ example "Basic Usage"258 259        ```python260        from langgraph_sdk import Encryption, EncryptionContext261 262        my_encryption = Encryption()263 264        SKIP_FIELDS = {"tenant_id", "owner", "thread_id", "assistant_id"}265        ENCRYPTED_PREFIX = "encrypted:"266 267        @my_encryption.encrypt.blob268        async def encrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:269            return your_encrypt_bytes(blob)270 271        @my_encryption.decrypt.blob272        async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:273            return your_decrypt_bytes(blob)274 275        @my_encryption.encrypt.json276        async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:277            result = {}278            for k, v in data.items():279                if k in SKIP_FIELDS or v is None:280                    result[k] = v281                else:282                    result[k] = ENCRYPTED_PREFIX + your_encrypt_string(v)283            return result284 285        @my_encryption.decrypt.json286        async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:287            result = {}288            for k, v in data.items():289                if isinstance(v, str) and v.startswith(ENCRYPTED_PREFIX):290                    result[k] = your_decrypt_string(v[len(ENCRYPTED_PREFIX):])291                else:292                    result[k] = v293            return result294        ```295 296    ???+ example "Field-Specific Logic"297 298        The `ctx.model` and `ctx.field` attributes tell you which model type and299        specific field is being encrypted, allowing different logic:300 301        ```python302        @my_encryption.encrypt.json303        async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:304            if ctx.field == "metadata":305                # Metadata - standard encryption306                return encrypt_standard(data)307            elif ctx.field == "values":308                # Thread values - more sensitive, use stronger encryption309                return encrypt_sensitive(data)310            else:311                return encrypt_standard(data)312        ```313 314        !!! warning "Model/Field May Differ Between Encrypt and Decrypt"315 316            Data encrypted with one `(model, field)` pair is **not guaranteed**317            to be decrypted with the same pair. The server performs SQL JSONB318            merges that can move encrypted values between models (e.g., cron319            metadata → run metadata). Your decryption logic must handle data320            regardless of the `ctx.model` or `ctx.field` values at decrypt time.321 322            **Safe:** Use `ctx.model`/`ctx.field` for logging or metrics only.323 324            **Safe:** Encrypt different keys based on `ctx.field`, but use a325            single decrypt handler that decrypts any value with the encrypted326            prefix (and passes through plaintext unchanged):327 328            ```python329            ENCRYPTED_PREFIX = "enc:"330 331            @my_encryption.encrypt.json332            async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:333                # Encrypt different keys depending on the field334                if ctx.field == "context":335                    keys_to_encrypt = {"api_key", "secret_token"}336                else:337                    keys_to_encrypt = {"email", "ssn"}338                return {339                    k: ENCRYPTED_PREFIX + encrypt(v) if k in keys_to_encrypt else v340                    for k, v in data.items()341                }342 343            @my_encryption.decrypt.json344            async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:345                # Decrypt ANY value with the prefix, regardless of model/field346                return {347                    k: decrypt(v[len(ENCRYPTED_PREFIX):])348                       if isinstance(v, str) and v.startswith(ENCRYPTED_PREFIX)349                       else v350                    for k, v in data.items()351                }352            ```353 354            **Unsafe:** Using different encryption keys or algorithms based on355            `ctx.model`/`ctx.field` will cause decryption failures.356    """357 358    __slots__ = (359        "_blob_decryptor",360        "_blob_encryptor",361        "_context_handler",362        "_json_decryptor",363        "_json_encryptor",364        "decrypt",365        "encrypt",366    )367 368    types = types369    """Reference to encryption type definitions.370 371    Provides access to all type definitions used in the encryption system,372    including EncryptionContext, BlobEncryptor, BlobDecryptor,373    JsonEncryptor, and JsonDecryptor.374    """375 376    def __init__(self) -> None:377        """Initialize the Encryption instance."""378        _warn_encryption_beta()379        self.encrypt = _EncryptDecorators(self)380        self.decrypt = _DecryptDecorators(self)381        self._blob_encryptor: types.BlobEncryptor | None = None382        self._blob_decryptor: types.BlobDecryptor | None = None383        self._json_encryptor: types.JsonEncryptor | None = None384        self._json_decryptor: types.JsonDecryptor | None = None385        self._context_handler: types.ContextHandler | None = None386 387    def context(self, fn: types.ContextHandler) -> types.ContextHandler:388        """Register a context handler to derive encryption context from auth.389 390        The handler receives the authenticated user and current EncryptionContext,391        and returns a dict that becomes ctx.metadata for encrypt/decrypt handlers.392 393        This allows encryption context to be derived from JWT claims or other394        auth-derived data instead of requiring a separate X-Encryption-Context header.395 396        Note: The context handler is called once per request in middleware,397        so ctx.model and ctx.field will be None in the handler.398 399        Example:400            ```python401            from langgraph_sdk import Encryption, EncryptionContext402            from starlette.authentication import BaseUser403 404            encryption = Encryption()405 406            @encryption.context407            async def get_context(user: BaseUser, ctx: EncryptionContext) -> dict:408                # Derive encryption context from authenticated user409                return {410                    **ctx.metadata,  # preserve X-Encryption-Context header if present411                    "tenant_id": user.tenant_id,412                }413            ```414 415        Args:416            fn: The context handler function417 418        Returns:419            The registered handler function420        """421        self._context_handler = fn422        return fn423 424    def get_json_encryptor(425        self,426        _model: str | None = None,  # kept for langgraph-api compat427    ) -> types.JsonEncryptor | None:428        """Get the JSON encryptor.429 430        Args:431            _model: Ignored. Kept for backwards compatibility with langgraph-api432                which passes model_type to this method.433 434        Returns:435            The JSON encryptor, or None if not registered.436        """437        return self._json_encryptor438 439    def get_json_decryptor(440        self,441        _model: str | None = None,  # kept for langgraph-api compat442    ) -> types.JsonDecryptor | None:443        """Get the JSON decryptor.444 445        Args:446            _model: Ignored. Kept for backwards compatibility with langgraph-api447                which passes model_type to this method.448 449        Returns:450            The JSON decryptor, or None if not registered.451        """452        return self._json_decryptor453 454    def __repr__(self) -> str:455        handlers = []456        if self._blob_encryptor:457            handlers.append("blob_encryptor")458        if self._blob_decryptor:459            handlers.append("blob_decryptor")460        if self._json_encryptor:461            handlers.append("json_encryptor")462        if self._json_decryptor:463            handlers.append("json_decryptor")464        if self._context_handler:465            handlers.append("context_handler")466        return f"Encryption(handlers=[{', '.join(handlers)}])"467