langgraph-api 0.4.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

langgraph_api/config/schemas.py

@@ -0,0 +1,431 @@
+import os
+import re
+from typing import Annotated, Literal
+
+from pydantic.functional_validators import AfterValidator
+from typing_extensions import TypedDict
+
+__all__ = [
+    "CheckpointerConfig",
+    "ConfigurableHeaders",
+    "CorsConfig",
+    "EncryptionConfig",
+    "HttpConfig",
+    "IndexConfig",
+    "MiddlewareOrders",
+    "SecurityConfig",
+    "SerdeConfig",
+    "StoreConfig",
+    "TTLConfig",
+    "ThreadTTLConfig",
+]
+
+
+class CorsConfig(TypedDict, total=False):
+    allow_origins: list[str]
+    """List of origins allowed to access the API (e.g., ["https://app.com"]).
+
+    Use ["*"] to allow any origin. Combined with `allow_origin_regex` when both are set.
+    """
+    allow_methods: list[str]
+    """HTTP methods permitted for cross-origin requests (e.g., ["GET", "POST"])."""
+    allow_headers: list[str]
+    """Request headers clients may send in cross-origin requests (e.g., ["authorization"])."""
+    allow_credentials: bool
+    """Whether browsers may include credentials (cookies, Authorization) in cross-origin requests."""
+    allow_origin_regex: str
+    """Regular expression that matches allowed origins; evaluated against the request Origin header."""
+    expose_headers: list[str]
+    """Response headers that browsers are allowed to read from CORS responses."""
+    max_age: int
+    """Number of seconds browsers may cache the CORS preflight (OPTIONS) response."""
+
+
+class ConfigurableHeaders(TypedDict, total=False):
+    includes: list[str] | None
+    """Header name patterns to include.
+
+    Patterns support literal names and the "*" wildcard (e.g., "x-*", "user-agent").
+    Matching is performed against lower-cased header names.
+    """
+    excludes: list[str] | None
+    """Header name patterns to exclude after inclusion rules are applied."""
+    include: list[str] | None
+    """Alias of `includes` for convenience/backwards compatibility."""
+    exclude: list[str] | None
+    """Alias of `excludes` for convenience/backwards compatibility."""
+
+
+MiddlewareOrders = Literal["auth_first", "middleware_first"]
+
+
+class HttpConfig(TypedDict, total=False):
+    app: str
+    """Import path for a custom Starlette/FastAPI app to mount"""
+    disable_assistants: bool
+    """Disable /assistants routes"""
+    disable_threads: bool
+    """Disable /threads routes"""
+    disable_runs: bool
+    """Disable /runs routes"""
+    disable_store: bool
+    """Disable /store routes"""
+    disable_meta: bool
+    """Disable /ok, /info, /metrics, and /docs routes"""
+    disable_webhooks: bool
+    """Disable webhooks calls on run completion in all routes"""
+    cors: CorsConfig | None
+    """CORS configuration for all mounted routes; see CorsConfig for details."""
+    disable_ui: bool
+    """Disable /ui routes"""
+    disable_mcp: bool
+    """Disable /mcp routes"""
+    disable_a2a: bool
+    """Disable /a2a routes"""
+    mount_prefix: str
+    """Prefix for mounted routes. E.g., "/my-deployment/api"."""
+    configurable_headers: ConfigurableHeaders | None
+    """Controls which inbound request headers are surfaced to runs as `config.configurable`.
+
+    Only headers that match `includes` and do not match `excludes` are exposed, except
+    for tracing headers like `langsmith-trace` and select `baggage` keys which are
+    always forwarded. Patterns are matched case-insensitively against lower-cased names.
+    """
+    logging_headers: ConfigurableHeaders | None
+    """Controls which inbound request headers may appear in access/application logs.
+
+    Use restrictive patterns (e.g., include "user-agent", exclude "authorization") to
+    avoid logging sensitive information.
+    """
+    enable_custom_route_auth: bool
+    """If true, apply the configured authentication middleware to user-supplied routes."""
+    middleware_order: MiddlewareOrders | None
+    """Ordering for authentication vs custom middleware on user routes.
+
+    - "auth_first": run auth middleware before user middleware
+    - "middleware_first": run user middleware before auth (default behavior)
+    """
+
+
+class ThreadTTLConfig(TypedDict, total=False):
+    strategy: Literal["delete", "keep_latest"]
+    """Action taken when a thread exceeds its TTL.
+
+    - "delete": Remove the thread and all its data entirely.
+    - "keep_latest": Prune old checkpoints but keep the thread and its latest state.
+      Requires core API (FF_USE_CORE_API=true).
+    """
+    default_ttl: float | None
+    """Default thread TTL in minutes; threads past this age are subject to the `strategy`."""
+    sweep_interval_minutes: int | None
+    """How often to scan for expired threads, in minutes."""
+    sweep_limit: int | None
+    """Maximum number of threads to process per sweep iteration. Defaults to 1000."""
+
+
+class IndexConfig(TypedDict, total=False):
+    """Configuration for indexing documents for semantic search in the store."""
+
+    dims: int
+    """Number of dimensions in the embedding vectors.
+    
+    Common embedding models have the following dimensions:
+        - OpenAI text-embedding-3-large: 256, 1024, or 3072
+        - OpenAI text-embedding-3-small: 512 or 1536
+        - OpenAI text-embedding-ada-002: 1536
+        - Cohere embed-english-v3.0: 1024
+        - Cohere embed-english-light-v3.0: 384
+        - Cohere embed-multilingual-v3.0: 1024
+        - Cohere embed-multilingual-light-v3.0: 384
+    """
+
+    embed: str
+    """Either a path to an embedding model (./path/to/file.py:embedding_model)
+    or a name of an embedding model (openai:text-embedding-3-small)
+    
+    Note: LangChain is required to use the model format specification.
+    """
+
+    fields: list[str] | None
+    """Fields to extract text from for embedding generation.
+    
+    Defaults to the root ["$"], which embeds the json object as a whole.
+    """
+
+
+class TTLConfig(TypedDict, total=False):
+    """Configuration for TTL (time-to-live) behavior in the store."""
+
+    refresh_on_read: bool
+    """Default behavior for refreshing TTLs on read operations (GET and SEARCH).
+    
+    If True, TTLs will be refreshed on read operations (get/search) by default.
+    This can be overridden per-operation by explicitly setting refresh_ttl.
+    Defaults to True if not configured.
+    """
+    default_ttl: float | None
+    """Default TTL (time-to-live) in minutes for new items.
+    
+    If provided, new items will expire after this many minutes after their last access.
+    The expiration timer refreshes on both read and write operations.
+    Defaults to None (no expiration).
+    """
+    sweep_interval_minutes: int | None
+    """Interval in minutes between TTL sweep operations.
+    
+    If provided, the store will periodically delete expired items based on TTL.
+    Defaults to None (no sweeping).
+    """
+
+
+class StoreConfig(TypedDict, total=False):
+    path: str
+    """Import path to a custom `BaseStore` or a callable returning one.
+
+    Examples:
+    - "./my_store.py:create_store"
+    - "my_package.store:Store"
+
+    When provided, this replaces the default Postgres-backed store.
+    """
+    index: IndexConfig
+    """Vector index settings for the built-in store (ignored by custom stores)."""
+    ttl: TTLConfig
+    """TTL behavior for stored items in the built-in store (custom stores may not support TTL)."""
+
+
+class SerdeConfig(TypedDict, total=False):
+    """Configuration for the built-in serde, which handles checkpointing of state.
+
+    If omitted, no serde is set up (the object store will still be present, however)."""
+
+    allowed_json_modules: list[list[str]] | Literal[True] | None
+    """Optional. List of allowed python modules to de-serialize custom objects from.
+    
+    If provided, only the specified modules will be allowed to be deserialized.
+    If omitted, no modules are allowed, and the object returned will simply be a json object OR
+    a deserialized langchain object.
+    
+    Example:
+    {...
+        "serde": {
+            "allowed_json_modules": [
+                ["my_agent", "my_file", "SomeType"],
+            ]
+        }
+    }
+
+    If you set this to True, any module will be allowed to be deserialized.
+
+    Example:
+    {...
+        "serde": {
+            "allowed_json_modules": true
+        }
+    }
+    
+    """
+    pickle_fallback: bool
+    """Optional. Whether to allow pickling as a fallback for deserialization.
+    
+    If True, pickling will be allowed as a fallback for deserialization.
+    If False, pickling will not be allowed as a fallback for deserialization.
+    Defaults to True if not configured."""
+
+
+class CheckpointerConfig(TypedDict, total=False):
+    """Configuration for the built-in checkpointer, which handles checkpointing of state.
+
+    If omitted, no checkpointer is set up (the object store will still be present, however).
+    """
+
+    ttl: ThreadTTLConfig | None
+    """Optional. Defines the TTL (time-to-live) behavior configuration.
+    
+    If provided, the checkpointer will apply TTL settings according to the configuration.
+    If omitted, no TTL behavior is configured.
+    """
+    serde: SerdeConfig | None
+    """Optional. Defines the configuration for how checkpoints are serialized."""
+
+
+class SecurityConfig(TypedDict, total=False):
+    securitySchemes: dict
+    """OpenAPI `components.securitySchemes` definition to merge into the spec."""
+    security: list
+    """Default security requirements applied to all operations (OpenAPI `security`)."""
+    # path => {method => security}
+    paths: dict[str, dict[str, list]]
+    """Per-route overrides of security requirements.
+
+    Mapping of path -> method -> OpenAPI `security` array.
+    """
+
+
+class CacheConfig(TypedDict, total=False):
+    cache_keys: list[str]
+    """Header names used to build the cache key for auth decisions."""
+    ttl_seconds: int
+    """How long to cache successful auth decisions, in seconds."""
+    max_size: int
+    """Maximum number of distinct auth cache entries to retain."""
+
+
+class AuthConfig(TypedDict, total=False):
+    path: str
+    """Path to the authentication function in a Python file."""
+    disable_studio_auth: bool
+    """Whether to disable auth when connecting from the LangSmith Studio."""
+    openapi: SecurityConfig
+    """The schema to use for updating the openapi spec.
+
+    Example:
+        {
+            "securitySchemes": {
+                "OAuth2": {
+                    "type": "oauth2",
+                    "flows": {
+                        "password": {
+                            "tokenUrl": "/token",
+                            "scopes": {
+                                "me": "Read information about the current user",
+                                "items": "Access to create and manage items"
+                            }
+                        }
+                    }
+                }
+            },
+            "security": [
+                {"OAuth2": ["me"]}  # Default security requirement for all endpoints
+            ]
+        }
+    """
+    cache: CacheConfig | None
+    """Optional cache settings for the custom auth backend to reduce repeated lookups."""
+
+
+class WebhookUrlPolicy(TypedDict, total=False):
+    require_https: bool
+    """Enforce HTTPS scheme for absolute URLs; reject `http://` when true."""
+    allowed_domains: list[str]
+    """Hostname allowlist. Supports exact hosts and wildcard subdomains.
+
+    Use entries like "hooks.example.com" or "*.mycorp.com". The wildcard only
+    matches subdomains ("foo.mycorp.com"), not the apex ("mycorp.com"). When
+    empty or omitted, any public host is allowed (subject to SSRF IP checks).
+    """
+    allowed_ports: list[int]
+    """Explicit port allowlist for absolute URLs.
+
+    If set, requests must use one of these ports. Defaults are respected when
+    a port is not present in the URL (443 for https, 80 for http).
+    """
+    max_url_length: int
+    """Maximum permitted URL length in characters; longer inputs are rejected early."""
+    disable_loopback: bool
+    """Disallow relative URLs (internal loopback calls) when true."""
+
+
+# Matches things like "${{ env.LG_WEBHOOK_FOO_BAR }}"
+_WEBHOOK_TEMPLATE_RE = re.compile(r"\$\{\{\s*([^}]+?)\s*\}\}")
+
+
+def _validate_url_policy(
+    policy: "WebhookUrlPolicy | None",
+) -> "WebhookUrlPolicy | None":
+    if not policy:
+        return policy
+    if "allowed_domains" in policy:
+        doms = policy["allowed_domains"]
+        if not isinstance(doms, list) or not all(
+            isinstance(d, str) and d for d in doms
+        ):
+            raise ValueError(
+                f"webhooks.url.allowed_domains must be a list of non-empty strings. Got: {doms}"
+            )
+        for d in doms:
+            if "*" in d and not d.startswith("*."):
+                raise ValueError(
+                    f"webhooks.url.allowed_domains wildcard can only be used as a prefix, in the form '*.domain'. Got: {doms}"
+                )
+    if "require_https" not in policy:
+        policy["require_https"] = True
+    if "allowed_domains" not in policy:
+        policy["allowed_domains"] = []
+    if "allowed_ports" not in policy:
+        policy["allowed_ports"] = []
+    if "max_url_length" not in policy:
+        policy["max_url_length"] = 2048
+    if "disable_loopback" not in policy:
+        policy["disable_loopback"] = False
+    return policy
+
+
+class WebhooksConfig(TypedDict, total=False):
+    env_prefix: str
+    """Required prefix for environment variables referenced in header templates.
+
+    Acts as an allowlist boundary to prevent leaking arbitrary environment
+    variables. Defaults to "LG_WEBHOOK_" when omitted.
+    """
+    url: Annotated[WebhookUrlPolicy, AfterValidator(_validate_url_policy)]
+    """URL validation policy for user-supplied webhook endpoints."""
+    headers: dict[str, str]
+    """Static headers to include with webhook requests.
+
+    Values may contain templates of the form "${{ env.VAR }}". On startup, these
+    are resolved via the process environment after verifying `VAR` starts with
+    `env_prefix`. Mixed literals and multiple templates are allowed.
+    """
+
+
+def webhooks_validator(cfg: "WebhooksConfig") -> "WebhooksConfig":
+    # Enforce env prefix & actual env presence at the aggregate level
+    headers = cfg.get("headers") or {}
+    if headers:
+        env_prefix = cfg.get("env_prefix", "LG_WEBHOOK_")
+        if env_prefix is None:
+            raise ValueError("webhook headers: Invalid null env_prefix")
+
+        def _replace(m: re.Match[str]) -> str:
+            expr = m.group(1).strip()
+            # Only variables we support are of the form "env.FOO_BAR" right now.
+            if not expr.startswith("env."):
+                raise ValueError(
+                    f"webhook headers: only env.VAR references are allowed (e.g. {{{{ env.{env_prefix}FOO_BAR }}}})"
+                )
+            var = expr[len("env.") :]
+            if not var or "." in var:
+                raise ValueError(
+                    f"webhook headers: invalid env reference '{var}'. Use env.VAR with no dots (e.g. {{{{ env.{env_prefix}FOO_BAR }}}})"
+                )
+            if env_prefix and not var.startswith(env_prefix):
+                raise ValueError(
+                    f"webhook headers: environment variable name '{var}' must start with the configured env_prefix '{env_prefix}'"
+                )
+            val = os.getenv(var)
+            if val is None:
+                raise ValueError(
+                    f"webhook headers: missing required environment variable '{var}'"
+                )
+            return val
+
+        rendered_headers = {}
+        for k, v in headers.items():
+            if not isinstance(v, str):
+                raise ValueError(f"Webhook header values must be strings. Got: {v}")
+            rendered = _WEBHOOK_TEMPLATE_RE.sub(_replace, v)
+            rendered_headers[k] = rendered
+        cfg["headers"] = rendered_headers
+    return cfg
+
+
+class EncryptionConfig(TypedDict, total=False):
+    path: str
+    """Path to the encryption module in a Python file.
+
+    Example: "./encryption.py:my_encryption"
+
+    The module should export an Encrypt instance with registered
+    encryption and decryption handlers for blobs and metadata.
+    """

langgraph_api/cron_scheduler.py

@@ -1,9 +1,13 @@
 import asyncio
 from random import random
+from typing import cast
 
 import structlog
 
+from langgraph_api import config
+from langgraph_api.encryption.middleware import decrypt_response
 from langgraph_api.models.run import create_valid_run
+from langgraph_api.serde import json_loads
 from langgraph_api.utils import next_cron_date
 from langgraph_api.utils.config import run_in_executor
 from langgraph_api.worker import set_auth_ctx_for_run
@@ -13,7 +17,7 @@ from langgraph_runtime.retry import retry_db
 
 logger = structlog.stdlib.get_logger(__name__)
 
-SLEEP_TIME = 5
+SLEEP_TIME = config.CRON_SCHEDULER_SLEEP_TIME
 
 
 @retry_db
@@ -23,7 +27,19 @@ async def cron_scheduler():
         try:
             async with connect() as conn:
                 async for cron in Crons.next(conn):
+                    on_run_completed = cron.get("on_run_completed")
+
                     run_payload = cron["payload"]
+                    if not isinstance(run_payload, dict):
+                        run_payload = json_loads(run_payload)
+                    run_payload = cast("dict", run_payload)
+
+                    run_payload = await decrypt_response(
+                        run_payload, "cron", ["metadata", "context", "input", "config"]
+                    )
+
+                    if on_run_completed == "keep":
+                        run_payload.setdefault("on_completion", "keep")  # type: ignore[union-attr]
 
                     async with set_auth_ctx_for_run(
                         run_payload, user_id=cron["user_id"]

langgraph_api/encryption/__init__.py

@@ -0,0 +1,15 @@
+"""Encryption support for LangGraph API."""
+
+from langgraph_api.encryption.custom import (
+    SUPPORTED_ENCRYPTION_MODELS,
+    ModelType,
+    get_custom_encryption_instance,
+)
+from langgraph_api.encryption.shared import get_encryption
+
+__all__ = [
+    "SUPPORTED_ENCRYPTION_MODELS",
+    "ModelType",
+    "get_custom_encryption_instance",
+    "get_encryption",
+]

langgraph_api/encryption/aes_json.py

@@ -0,0 +1,158 @@
+"""AES encryption for JSON field values.
+
+This module provides opt-in AES encryption for specific JSON keys,
+using the same key and cipher as checkpoint encryption (LANGGRAPH_AES_KEY).
+"""
+
+from __future__ import annotations
+
+import base64
+import functools
+from typing import TYPE_CHECKING, Any
+
+import orjson
+from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
+
+from langgraph_api.encryption.shared import (
+    ENCRYPTION_CONTEXT_KEY,
+)
+
+if TYPE_CHECKING:
+    from langgraph.checkpoint.serde.base import CipherProtocol
+    from langgraph_sdk.encryption.types import JsonDecryptor, JsonEncryptor
+
+    from langgraph_api.encryption.custom import ModelType
+
+AES_ENCRYPTED_PREFIX = "encrypted.aes:"
+
+# Marker key to identify AES encryption in __encryption_context__
+AES_ENCRYPTION_TYPE_KEY = "__langgraph_encryption_type__"
+AES_ENCRYPTION_CONTEXT = {AES_ENCRYPTION_TYPE_KEY: "aes"}
+
+
+def _get_aes_cipher(key: bytes) -> CipherProtocol:
+    """Get AES cipher using the SDK's implementation (same as checkpoint encryption)."""
+    # EncryptedSerializer.from_pycryptodome_aes creates a PycryptodomeAesCipher internally
+    # We extract it to reuse the exact same encryption format as checkpoints
+    return EncryptedSerializer.from_pycryptodome_aes(key=key).cipher
+
+
+def is_aes_encrypted(value: Any) -> bool:
+    """Check if a value is AES-encrypted (has the encryption prefix)."""
+    return isinstance(value, str) and value.startswith(AES_ENCRYPTED_PREFIX)
+
+
+def has_any_aes_encrypted_values(data: dict[str, Any]) -> bool:
+    """Check if any value in the dict is AES-encrypted (top-level only)."""
+    return isinstance(data, dict) and any(is_aes_encrypted(v) for v in data.values())
+
+
+def is_aes_encryption_context(ctx: dict[str, Any] | None) -> bool:
+    """Check if an encryption context indicates AES encryption."""
+    return isinstance(ctx, dict) and ctx.get(AES_ENCRYPTION_TYPE_KEY) == "aes"
+
+
+class EncryptionKeyError(Exception):
+    """Raised when JSON encryptor violates key preservation constraint."""
+
+
+class EncryptionRoutingError(Exception):
+    """Raised when encryption routing fails due to inconsistent markers."""
+
+
+class DecryptorMissingError(Exception):
+    """Raised when data has encryption marker but no decryptor is configured."""
+
+
+class AesEncryptionInstance:
+    """Built-in AES encryption for JSON field values.
+
+    Uses the same AES cipher as checkpoint encryption (via SDK's CipherProtocol).
+    Duck-types the SDK's Encryption interface for use in the middleware.
+    """
+
+    def __init__(self, key: bytes, allowlist: frozenset[str]) -> None:
+        self._cipher = _get_aes_cipher(key)
+        self._allowlist = allowlist
+
+    def encrypt_value(self, value: Any) -> str:
+        """Encrypt a JSON-serializable value, returning prefixed ciphertext."""
+        plaintext = orjson.dumps(value)
+        _, encrypted_blob = self._cipher.encrypt(plaintext)
+        encoded = base64.b64encode(encrypted_blob).decode("ascii")
+        return f"{AES_ENCRYPTED_PREFIX}{encoded}"
+
+    def decrypt_value(self, encrypted: str) -> Any:
+        """Decrypt an AES-encrypted value."""
+        if not encrypted.startswith(AES_ENCRYPTED_PREFIX):
+            raise ValueError(f"Expected prefix '{AES_ENCRYPTED_PREFIX}'")
+        encrypted_blob = base64.b64decode(encrypted[len(AES_ENCRYPTED_PREFIX) :])
+        plaintext = self._cipher.decrypt("aes", encrypted_blob)
+        return orjson.loads(plaintext)
+
+    def encrypt_json(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Encrypt allowlisted keys in a JSON dict."""
+        if not data:
+            return data
+        keys_to_encrypt = data.keys() & self._allowlist
+        if not keys_to_encrypt:
+            return data
+        result = {
+            k: self.encrypt_value(v) if k in keys_to_encrypt and v is not None else v
+            for k, v in data.items()
+        }
+        result[ENCRYPTION_CONTEXT_KEY] = AES_ENCRYPTION_CONTEXT.copy()
+        return result
+
+    def decrypt_json(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Decrypt all AES-encrypted values in a JSON dict."""
+        if not data:
+            return data
+        return {
+            k: self.decrypt_value(v) if is_aes_encrypted(v) else v
+            for k, v in data.items()
+            if k != ENCRYPTION_CONTEXT_KEY
+        }
+
+    def get_json_encryptor(self, model_type: ModelType) -> JsonEncryptor:
+        """Return an async encryptor function for the given model type."""
+
+        async def encryptor(ctx: Any, data: dict[str, Any]) -> dict[str, Any]:
+            return self.encrypt_json(data)
+
+        return encryptor
+
+    def get_json_decryptor(self, model_type: ModelType) -> JsonDecryptor:
+        """Return an async decryptor function for the given model type."""
+
+        async def decryptor(ctx: Any, data: dict[str, Any]) -> dict[str, Any]:
+            return self.decrypt_json(data)
+
+        return decryptor
+
+
+@functools.lru_cache(maxsize=1)
+def get_aes_encryption_instance() -> AesEncryptionInstance | None:
+    """Get the AES encryption instance if configured.
+
+    Returns:
+        - AesEncryptionInstance with allowlist if both LANGGRAPH_AES_KEY and
+          LANGGRAPH_AES_JSON_KEYS are configured (encrypts and decrypts)
+        - AesEncryptionInstance with empty allowlist if only LANGGRAPH_AES_KEY
+          is configured (decrypts only, for migration from AES to custom)
+        - None if LANGGRAPH_AES_KEY is not configured
+    """
+    # Import here to avoid circular imports
+    from langgraph_api.config import LANGGRAPH_AES_JSON_KEYS, LANGGRAPH_AES_KEY
+
+    if not LANGGRAPH_AES_KEY:
+        return None
+
+    # If both key and json_keys are set, use full encryption/decryption
+    if LANGGRAPH_AES_JSON_KEYS:
+        return AesEncryptionInstance(LANGGRAPH_AES_KEY, LANGGRAPH_AES_JSON_KEYS)
+
+    # If only key is set (no json_keys), create decryption-only instance.
+    # This supports migration from AES to custom encryption: the server can
+    # decrypt old AES-encrypted data without re-encrypting new data with AES.
+    return AesEncryptionInstance(LANGGRAPH_AES_KEY, frozenset())

langgraph_api/encryption/context.py

@@ -0,0 +1,35 @@
+"""Request-scoped encryption context storage.
+
+This module provides a ContextVar for storing encryption context
+(tenant ID, key identifiers, etc.) that is accessible throughout
+the async request lifecycle, including in checkpoint serialization.
+"""
+
+from contextvars import ContextVar
+from typing import Any
+
+# Request-scoped encryption context
+# Set by API middleware when X-Encryption-Context header is present
+# Accessed by serializers during checkpoint encryption/decryption
+encryption_context: ContextVar[dict[str, Any] | None] = ContextVar(  # type: ignore[assignment]
+    "encryption_context", default=None
+)
+
+
+def get_encryption_context() -> dict[str, Any]:
+    """Get the current request's encryption context.
+
+    Returns:
+        The encryption context dict, or empty dict if not set.
+    """
+    ctx = encryption_context.get()
+    return ctx if ctx is not None else {}
+
+
+def set_encryption_context(context: dict[str, Any]) -> None:
+    """Set the encryption context for the current request.
+
+    Args:
+        context: The encryption context dict (e.g., {"tenant_id": "...", "key_id": "..."})
+    """
+    encryption_context.set(context)