"""Custom encryption support for LangGraph. .. warning:: This API is in beta and may change in future versions. This module provides a framework for implementing custom at-rest encryption in LangGraph applications. Similar to the Auth system, it allows developers to define custom encryption and decryption handlers that are executed server-side. """ from __future__ import annotations import functools import inspect import typing import warnings from langgraph_sdk.encryption import types class LangGraphBetaWarning(UserWarning): """Warning for beta features in LangGraph SDK.""" @functools.lru_cache(maxsize=1) def _warn_encryption_beta() -> None: warnings.warn( "The Encryption API is in beta and may change in future versions.", LangGraphBetaWarning, stacklevel=4, ) class DuplicateHandlerError(Exception): """Raised when attempting to register a duplicate encryption/decryption handler.""" pass def _validate_handler(fn: typing.Callable, handler_type: str) -> None: """Validate that a handler function has the correct signature. Args: fn: The handler function to validate handler_type: Description of the handler for error messages Raises: TypeError: If the handler is not an async function or has wrong parameter count """ if not inspect.iscoroutinefunction(fn): raise TypeError(f"{handler_type} must be an async function, got {type(fn)}") sig = inspect.signature(fn) params = [ p for p in sig.parameters.values() if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) ] if len(params) != 2: raise TypeError( f"{handler_type} must accept exactly 2 parameters " f"(ctx, data), got {len(params)}" ) class _JsonEncryptDecorators: """Dynamic decorator factory for JSON encryption handlers. Supports both default and model-specific handlers: - @encrypt.json - default handler for all models - @encrypt.json.thread - handler for thread model """ def __init__(self, parent: Encryption): self._parent = parent def __call__(self, fn: types.JsonEncryptor) -> types.JsonEncryptor: """Register the default JSON encryption handler. Args: fn: The handler function Returns: The registered handler function Raises: DuplicateHandlerError: If handler already registered TypeError: If handler has invalid signature """ if self._parent._json_encryptor is not None: raise DuplicateHandlerError("Default JSON encryptor already registered") _validate_handler(fn, "Default JSON encryptor") self._parent._json_encryptor = fn return fn def __getattr__( self, model: str ) -> typing.Callable[[types.JsonEncryptor], types.JsonEncryptor]: """Dynamic attribute access for model-specific handlers. Allows @encryption.encrypt.json.thread, @encryption.encrypt.json.assistant, etc. Raises: DuplicateHandlerError: If handler already registered for this model TypeError: If handler has invalid signature """ def decorator(fn: types.JsonEncryptor) -> types.JsonEncryptor: if model in self._parent._json_encryptors: raise DuplicateHandlerError( f"JSON encryptor for model '{model}' already registered" ) _validate_handler(fn, f"JSON encryptor for model '{model}'") self._parent._json_encryptors[model] = fn return fn return decorator class _JsonDecryptDecorators: """Dynamic decorator factory for JSON decryption handlers. Supports both default and model-specific handlers: - @encryption.decrypt.json - default handler for all models - @encryption.decrypt.json.thread - handler for thread model """ def __init__(self, parent: Encryption): self._parent = parent def __call__(self, fn: types.JsonDecryptor) -> types.JsonDecryptor: """Register the default JSON decryption handler. Args: fn: The handler function Returns: The registered handler function Raises: DuplicateHandlerError: If handler already registered TypeError: If handler has invalid signature """ if self._parent._json_decryptor is not None: raise DuplicateHandlerError("Default JSON decryptor already registered") _validate_handler(fn, "Default JSON decryptor") self._parent._json_decryptor = fn return fn def __getattr__( self, model: str ) -> typing.Callable[[types.JsonDecryptor], types.JsonDecryptor]: """Dynamic attribute access for model-specific handlers. Allows @encryption.decrypt.json.thread, @encryption.decrypt.json.assistant, etc. Raises: DuplicateHandlerError: If handler already registered for this model TypeError: If handler has invalid signature """ def decorator(fn: types.JsonDecryptor) -> types.JsonDecryptor: if model in self._parent._json_decryptors: raise DuplicateHandlerError( f"JSON decryptor for model '{model}' already registered" ) _validate_handler(fn, f"JSON decryptor for model '{model}'") self._parent._json_decryptors[model] = fn return fn return decorator class _EncryptDecorators: """Decorators for encryption handlers. Provides @encryption.encrypt.blob and @encryption.encrypt.json decorators for registering encryption functions. """ def __init__(self, parent: Encryption): self._parent = parent self._json = _JsonEncryptDecorators(parent) def blob(self, fn: types.BlobEncryptor) -> types.BlobEncryptor: """Register a blob encryption handler. The handler will be called to encrypt opaque data like checkpoint blobs. Example: ```python @encryption.encrypt.blob async def encrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes: # Encrypt the blob using your encryption service return encrypted_blob ``` Args: fn: The encryption handler function Returns: The registered handler function Raises: DuplicateHandlerError: If blob encryptor already registered TypeError: If handler has invalid signature """ if self._parent._blob_encryptor is not None: raise DuplicateHandlerError("Blob encryptor already registered") _validate_handler(fn, "Blob encryptor") self._parent._blob_encryptor = fn return fn @property def json(self) -> _JsonEncryptDecorators: """Access JSON encryption decorators. Supports model-specific handlers: - @encryption.encrypt.json - default handler for all models - @encryption.encrypt.json.thread - handler for thread model only - @encryption.encrypt.json.assistant - handler for assistant model only Example: ```python @encryption.encrypt.json async def default_encrypt(ctx: EncryptionContext, data: dict) -> dict: # Default encryption for all models return encrypt_data(data) @encryption.encrypt.json.thread async def encrypt_thread(ctx: EncryptionContext, data: dict) -> dict: # Special encryption for thread model only return encrypt_thread_data(data) ``` """ return self._json class _DecryptDecorators: """Decorators for decryption handlers. Provides @encryption.decrypt.blob and @encryption.decrypt.json decorators for registering decryption functions. """ def __init__(self, parent: Encryption): self._parent = parent self._json = _JsonDecryptDecorators(parent) def blob(self, fn: types.BlobDecryptor) -> types.BlobDecryptor: """Register a blob decryption handler. The handler will be called to decrypt opaque data like checkpoint blobs. Example: ```python @encryption.decrypt.blob async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes: # Decrypt the blob using your encryption service return decrypted_blob ``` Args: fn: The decryption handler function Returns: The registered handler function Raises: DuplicateHandlerError: If blob decryptor already registered TypeError: If handler has invalid signature """ if self._parent._blob_decryptor is not None: raise DuplicateHandlerError("Blob decryptor already registered") _validate_handler(fn, "Blob decryptor") self._parent._blob_decryptor = fn return fn @property def json(self) -> _JsonDecryptDecorators: """Access JSON decryption decorators. Supports model-specific handlers: - @encryption.decrypt.json - default handler for all models - @encryption.decrypt.json.thread - handler for thread model only - @encryption.decrypt.json.assistant - handler for assistant model only Example: ```python @encryption.decrypt.json async def default_decrypt(ctx: EncryptionContext, data: dict) -> dict: # Default decryption for all models return decrypt_data(data) @encryption.decrypt.json.thread async def decrypt_thread(ctx: EncryptionContext, data: dict) -> dict: # Special decryption for thread model only return decrypt_thread_data(data) ``` """ return self._json class Encryption: """Add custom at-rest encryption to your LangGraph application. .. warning:: This API is in beta and may change in future versions. The Encryption class provides a system for implementing custom encryption of data at rest in LangGraph applications. It supports encryption of both opaque blobs (like checkpoints) and structured JSON data (like metadata, context, kwargs, values, etc.). To use, create a separate Python file and add the path to the file to your LangGraph API configuration file (`langgraph.json`). Within that file, create an instance of the Encryption class and register encryption and decryption handlers as needed. Example `langgraph.json` file: ```json { "dependencies": ["."], "graphs": { "agent": "./my_agent/agent.py:graph" }, "env": ".env", "encryption": { "path": "./encryption.py:my_encryption" } } ``` Then the LangGraph server will load your encryption file and use it to encrypt/decrypt data at rest. ???+ example "Basic Usage" ```python from langgraph_sdk import Encryption, EncryptionContext my_encryption = Encryption() @my_encryption.encrypt.blob async def encrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes: # Call your encryption service return encrypted_blob @my_encryption.decrypt.blob async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes: # Call your decryption service return decrypted_blob @my_encryption.encrypt.json async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict: # Practical encryption strategy: # - "owner" field: unencrypted (for search/filtering) # - "my.customer.org/" prefixed fields: encrypt VALUES only # - All other fields: pass through unencrypted encrypted = {} for key, value in data.items(): if key.startswith("my.customer.org/"): # Encrypt VALUE for sensitive customer data encrypted[key] = encrypt_value(value) else: # Pass through (including "owner" for search) encrypted[key] = value return encrypted @my_encryption.decrypt.json async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict: # Decrypt VALUES for "my.customer.org/" prefixed fields decrypted = {} for key, value in data.items(): if key.startswith("my.customer.org/"): decrypted[key] = decrypt_value(value) else: decrypted[key] = value return decrypted ``` ???+ example "Model-Specific Handlers" You can register different encryption handlers for different model types (thread, assistant, run, cron, checkpoint, etc.): ```python from langgraph_sdk import Encryption, EncryptionContext my_encryption = Encryption() # Default handler for models without specific handlers @my_encryption.encrypt.json async def default_encrypt(ctx: EncryptionContext, data: dict) -> dict: return standard_encrypt(data) # Thread-specific handler (uses different KMS key) @my_encryption.encrypt.json.thread async def encrypt_thread(ctx: EncryptionContext, data: dict) -> dict: return encrypt_with_thread_key(data) # Assistant-specific handler @my_encryption.encrypt.json.assistant async def encrypt_assistant(ctx: EncryptionContext, data: dict) -> dict: return encrypt_with_assistant_key(data) # Same pattern for decryption @my_encryption.decrypt.json async def default_decrypt(ctx: EncryptionContext, data: dict) -> dict: return standard_decrypt(data) @my_encryption.decrypt.json.thread async def decrypt_thread(ctx: EncryptionContext, data: dict) -> dict: return decrypt_with_thread_key(data) ``` ???+ example "Field-Specific Logic" The `ctx.field` attribute tells you which specific field is being encrypted, allowing different logic within the same model: ```python @my_encryption.encrypt.json.thread async def encrypt_thread(ctx: EncryptionContext, data: dict) -> dict: if ctx.field == "metadata": # Thread metadata - standard encryption return encrypt_standard(data) elif ctx.field == "values": # Thread values - more sensitive, use stronger encryption return encrypt_sensitive(data) else: return encrypt_standard(data) ``` """ __slots__ = ( "_blob_decryptor", "_blob_encryptor", "_context_handler", "_json_decryptor", "_json_decryptors", "_json_encryptor", "_json_encryptors", "decrypt", "encrypt", ) types = types """Reference to encryption type definitions. Provides access to all type definitions used in the encryption system, including EncryptionContext, BlobEncryptor, BlobDecryptor, JsonEncryptor, and JsonDecryptor. """ def __init__(self) -> None: """Initialize the Encryption instance.""" _warn_encryption_beta() self.encrypt = _EncryptDecorators(self) self.decrypt = _DecryptDecorators(self) self._blob_encryptor: types.BlobEncryptor | None = None self._blob_decryptor: types.BlobDecryptor | None = None self._json_encryptor: types.JsonEncryptor | None = None self._json_decryptor: types.JsonDecryptor | None = None self._json_encryptors: dict[str, types.JsonEncryptor] = {} self._json_decryptors: dict[str, types.JsonDecryptor] = {} self._context_handler: types.ContextHandler | None = None def context(self, fn: types.ContextHandler) -> types.ContextHandler: """Register a context handler to derive encryption context from auth. The handler receives the authenticated user and current EncryptionContext, and returns a dict that becomes ctx.metadata for encrypt/decrypt handlers. This allows encryption context to be derived from JWT claims or other auth-derived data instead of requiring a separate X-Encryption-Context header. Note: The context handler is called once per request in middleware, so ctx.model and ctx.field will be None in the handler. Example: ```python from langgraph_sdk import Encryption, EncryptionContext from starlette.authentication import BaseUser encryption = Encryption() @encryption.context async def get_context(user: BaseUser, ctx: EncryptionContext) -> dict: # Derive encryption context from authenticated user return { **ctx.metadata, # preserve X-Encryption-Context header if present "tenant_id": user.tenant_id, } ``` Args: fn: The context handler function Returns: The registered handler function """ self._context_handler = fn return fn def get_json_encryptor( self, model: str | None = None ) -> types.JsonEncryptor | None: """Get the JSON encryptor for a specific model. Args: model: The model type (e.g., "thread", "assistant"). If None, returns default. Returns: Model-specific encryptor if registered, otherwise default encryptor, or None. """ if model and model in self._json_encryptors: return self._json_encryptors[model] return self._json_encryptor def get_json_decryptor( self, model: str | None = None ) -> types.JsonDecryptor | None: """Get the JSON decryptor for a specific model. Args: model: The model type (e.g., "thread", "assistant"). If None, returns default. Returns: Model-specific decryptor if registered, otherwise default decryptor, or None. """ if model and model in self._json_decryptors: return self._json_decryptors[model] return self._json_decryptor def __repr__(self) -> str: handlers = [] if self._blob_encryptor: handlers.append("blob_encryptor") if self._blob_decryptor: handlers.append("blob_decryptor") if self._json_encryptor: handlers.append("json_encryptor") if self._json_decryptor: handlers.append("json_decryptor") if self._json_encryptors: handlers.append(f"json_encryptors({list(self._json_encryptors.keys())})") if self._json_decryptors: handlers.append(f"json_decryptors({list(self._json_decryptors.keys())})") if self._context_handler: handlers.append("context_handler") return f"Encryption(handlers=[{', '.join(handlers)}])"