initial commit

This commit is contained in:
2026-05-11 12:36:20 +05:30
commit 384cbe8019
15377 changed files with 2360544 additions and 0 deletions

View File

@@ -0,0 +1,466 @@
"""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 _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
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
def json(self, fn: types.JsonEncryptor) -> types.JsonEncryptor:
"""Register the JSON encryption handler.
Example:
```python
@encryption.encrypt.json
async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:
# Encrypt the data
return encrypt_data(data)
```
Args:
fn: The encryption handler function
Returns:
The registered handler function
Raises:
DuplicateHandlerError: If JSON encryptor already registered
TypeError: If handler has invalid signature
"""
if self._parent._json_encryptor is not None:
raise DuplicateHandlerError("JSON encryptor already registered")
_validate_handler(fn, "JSON encryptor")
self._parent._json_encryptor = fn
return fn
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
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
def json(self, fn: types.JsonDecryptor) -> types.JsonDecryptor:
"""Register the JSON decryption handler.
Example:
```python
@encryption.decrypt.json
async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:
# Decrypt the data
return decrypt_data(data)
```
Args:
fn: The decryption handler function
Returns:
The registered handler function
Raises:
DuplicateHandlerError: If JSON decryptor already registered
TypeError: If handler has invalid signature
"""
if self._parent._json_decryptor is not None:
raise DuplicateHandlerError("JSON decryptor already registered")
_validate_handler(fn, "JSON decryptor")
self._parent._json_decryptor = fn
return fn
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.
!!! warning "JSON Encryptors Must Preserve Keys"
JSON encryptors **must not add or remove keys** from the input dict.
Only values may be transformed. This constraint is **enforced at runtime
by the server** and exists because SQL JSONB merge operations (used for
partial updates) work at the key level.
**Correct (per-key encryption):**
```python
# Input: {"secret": "value", "plain": "x"}
# Output: {"secret": "<encrypted>", "plain": "x"} ✓ Keys preserved
```
**Incorrect (key consolidation):**
```python
# Input: {"secret": "value", "plain": "x"}
# Output: {"__encrypted__": "<blob>", "plain": "x"} ✗ Key changed
```
If your encryptor needs to store auxiliary data (DEK, IV, etc.), embed it
within the encrypted value itself, not as separate keys.
???+ example "Basic Usage"
```python
from langgraph_sdk import Encryption, EncryptionContext
my_encryption = Encryption()
SKIP_FIELDS = {"tenant_id", "owner", "thread_id", "assistant_id"}
ENCRYPTED_PREFIX = "encrypted:"
@my_encryption.encrypt.blob
async def encrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:
return your_encrypt_bytes(blob)
@my_encryption.decrypt.blob
async def decrypt_blob(ctx: EncryptionContext, blob: bytes) -> bytes:
return your_decrypt_bytes(blob)
@my_encryption.encrypt.json
async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:
result = {}
for k, v in data.items():
if k in SKIP_FIELDS or v is None:
result[k] = v
else:
result[k] = ENCRYPTED_PREFIX + your_encrypt_string(v)
return result
@my_encryption.decrypt.json
async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:
result = {}
for k, v in data.items():
if isinstance(v, str) and v.startswith(ENCRYPTED_PREFIX):
result[k] = your_decrypt_string(v[len(ENCRYPTED_PREFIX):])
else:
result[k] = v
return result
```
???+ example "Field-Specific Logic"
The `ctx.model` and `ctx.field` attributes tell you which model type and
specific field is being encrypted, allowing different logic:
```python
@my_encryption.encrypt.json
async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:
if ctx.field == "metadata":
# 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)
```
!!! warning "Model/Field May Differ Between Encrypt and Decrypt"
Data encrypted with one `(model, field)` pair is **not guaranteed**
to be decrypted with the same pair. The server performs SQL JSONB
merges that can move encrypted values between models (e.g., cron
metadata → run metadata). Your decryption logic must handle data
regardless of the `ctx.model` or `ctx.field` values at decrypt time.
**Safe:** Use `ctx.model`/`ctx.field` for logging or metrics only.
**Safe:** Encrypt different keys based on `ctx.field`, but use a
single decrypt handler that decrypts any value with the encrypted
prefix (and passes through plaintext unchanged):
```python
ENCRYPTED_PREFIX = "enc:"
@my_encryption.encrypt.json
async def encrypt_json(ctx: EncryptionContext, data: dict) -> dict:
# Encrypt different keys depending on the field
if ctx.field == "context":
keys_to_encrypt = {"api_key", "secret_token"}
else:
keys_to_encrypt = {"email", "ssn"}
return {
k: ENCRYPTED_PREFIX + encrypt(v) if k in keys_to_encrypt else v
for k, v in data.items()
}
@my_encryption.decrypt.json
async def decrypt_json(ctx: EncryptionContext, data: dict) -> dict:
# Decrypt ANY value with the prefix, regardless of model/field
return {
k: decrypt(v[len(ENCRYPTED_PREFIX):])
if isinstance(v, str) and v.startswith(ENCRYPTED_PREFIX)
else v
for k, v in data.items()
}
```
**Unsafe:** Using different encryption keys or algorithms based on
`ctx.model`/`ctx.field` will cause decryption failures.
"""
__slots__ = (
"_blob_decryptor",
"_blob_encryptor",
"_context_handler",
"_json_decryptor",
"_json_encryptor",
"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._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, # kept for langgraph-api compat
) -> types.JsonEncryptor | None:
"""Get the JSON encryptor.
Args:
_model: Ignored. Kept for backwards compatibility with langgraph-api
which passes model_type to this method.
Returns:
The JSON encryptor, or None if not registered.
"""
return self._json_encryptor
def get_json_decryptor(
self,
_model: str | None = None, # kept for langgraph-api compat
) -> types.JsonDecryptor | None:
"""Get the JSON decryptor.
Args:
_model: Ignored. Kept for backwards compatibility with langgraph-api
which passes model_type to this method.
Returns:
The JSON decryptor, or None if not registered.
"""
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._context_handler:
handlers.append("context_handler")
return f"Encryption(handlers=[{', '.join(handlers)}])"

View File

@@ -0,0 +1,147 @@
"""Encryption and decryption types for LangGraph.
This module defines the core types used for custom at-rest encryption
in LangGraph. It includes context types and typed dictionaries for
encryption operations.
"""
from __future__ import annotations
import typing
from collections.abc import Awaitable, Callable
Json = dict[str, typing.Any]
"""JSON-serializable dictionary type for structured data encryption."""
class EncryptionContext:
"""Context passed to encryption/decryption handlers.
Contains arbitrary non-secret key-values that will be stored on encrypt.
These key-values are intended to be sent to an external service that
manages keys and handles the actual encryption and decryption of data.
Attributes:
model: The model type being encrypted (e.g., "assistant", "thread", "run", "checkpoint")
field: The specific field being encrypted (e.g., "metadata", "context", "kwargs", "values")
metadata: Additional context metadata that can be used for encryption decisions
"""
__slots__ = ("field", "metadata", "model")
def __init__(
self,
model: str | None = None,
metadata: dict[str, typing.Any] | None = None,
field: str | None = None,
):
self.model = model
self.field = field
self.metadata = metadata or {}
def __repr__(self) -> str:
return f"EncryptionContext(model={self.model!r}, field={self.field!r}, metadata={self.metadata!r})"
BlobEncryptor = Callable[[EncryptionContext, bytes], Awaitable[bytes]]
"""Handler for encrypting opaque blob data like checkpoints.
Note: Must be an async function. Encryption typically involves I/O operations
(calling external KMS services), which should be async.
Args:
ctx: Encryption context with model type and metadata
blob: The raw bytes to encrypt
Returns:
Awaitable that resolves to encrypted bytes
"""
BlobDecryptor = Callable[[EncryptionContext, bytes], Awaitable[bytes]]
"""Handler for decrypting opaque blob data like checkpoints.
Note: Must be an async function. Decryption typically involves I/O operations
(calling external KMS services), which should be async.
Args:
ctx: Encryption context with model type and metadata
blob: The encrypted bytes to decrypt
Returns:
Awaitable that resolves to decrypted bytes
"""
JsonEncryptor = Callable[[EncryptionContext, Json], Awaitable[Json]]
"""Handler for encrypting structured JSON data.
Note: Must be an async function. Encryption typically involves I/O operations
(calling external KMS services), which should be async.
Used for encrypting structured data like metadata, context, kwargs, values,
and other JSON-serializable fields across different model types.
Maps plaintext fields to encrypted fields. A practical approach:
- Keep "owner" field unencrypted for search/filtering
- Encrypt VALUES (not keys) for fields with specific prefix (e.g., "my.customer.org/")
- Pass through all other fields unencrypted
Example:
Input: {"owner": "user123", "my.customer.org/email": "john@example.com", "tenant_id": "t-456"}
Output: {"owner": "user123", "my.customer.org/email": "ENCRYPTED", "tenant_id": "t-456"}
Note: Encrypted field VALUES cannot be reliably searched, as most real-world
encryption implementations use nonces (non-deterministic encryption).
Only unencrypted fields can be used in search queries.
Args:
ctx: Encryption context with model type, field name, and metadata
data: The plaintext JSON dictionary
Returns:
Awaitable that resolves to encrypted JSON dictionary
"""
JsonDecryptor = Callable[[EncryptionContext, Json], Awaitable[Json]]
"""Handler for decrypting structured JSON data.
Note: Must be an async function. Decryption typically involves I/O operations
(calling external KMS services), which should be async.
Inverse of JsonEncryptor. Must be able to decrypt data that
was encrypted by the corresponding encryptor.
Args:
ctx: Encryption context with model type, field name, and metadata
data: The encrypted JSON dictionary
Returns:
Awaitable that resolves to decrypted JSON dictionary
"""
if typing.TYPE_CHECKING:
from starlette.authentication import BaseUser
ContextHandler = Callable[
["BaseUser", EncryptionContext], Awaitable[dict[str, typing.Any]]
]
"""Handler for deriving encryption context from authenticated user info.
Note: Must be an async function as it may involve I/O operations.
The context handler is called once per request in middleware (after auth),
allowing encryption context to be derived from JWT claims, user properties,
or other auth-derived data instead of requiring a separate X-Encryption-Context header.
The return value becomes ctx.metadata for subsequent encrypt/decrypt operations
and is persisted with encrypted data for later decryption.
Note: ctx.model and ctx.field will be None in context handlers since
the handler runs once per request before any specific model/field is known.
Args:
user: The authenticated user (from Starlette's AuthenticationMiddleware)
ctx: Current encryption context with metadata from X-Encryption-Context header
Returns:
Awaitable that resolves to dict that becomes the new ctx.metadata
"""