From 9dde905dcb9e8071f6501a5db1184b8793ba05c8 Mon Sep 17 00:00:00 2001 From: Connor Braa Date: Tue, 11 Aug 2026 11:16:57 -0700 Subject: [PATCH] fix(sdk): preserve decrypt handler types --- .../langgraph_sdk/encryption/__init__.py | 7 ++++-- libs/sdk-py/tests/test_encryption.py | 25 ++++++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/libs/sdk-py/langgraph_sdk/encryption/__init__.py b/libs/sdk-py/langgraph_sdk/encryption/__init__.py index 54e269274..c2b3f8a03 100644 --- a/libs/sdk-py/langgraph_sdk/encryption/__init__.py +++ b/libs/sdk-py/langgraph_sdk/encryption/__init__.py @@ -18,6 +18,9 @@ import warnings from langgraph_sdk.encryption import types +_BlobDecryptorT = typing.TypeVar("_BlobDecryptorT", bound=types.BlobDecryptor) +_JsonDecryptorT = typing.TypeVar("_JsonDecryptorT", bound=types.JsonDecryptor) + class LangGraphBetaWarning(UserWarning): """Warning for beta features in LangGraph SDK.""" @@ -141,7 +144,7 @@ class _DecryptDecorators: def __init__(self, parent: Encryption): self._parent = parent - def blob(self, fn: types.BlobDecryptor) -> types.BlobDecryptor: + def blob(self, fn: _BlobDecryptorT) -> _BlobDecryptorT: """Register a blob decryption handler. The handler will be called to decrypt opaque data like checkpoint blobs. @@ -172,7 +175,7 @@ class _DecryptDecorators: self._parent._blob_decryptor = fn return fn - def json(self, fn: types.JsonDecryptor) -> types.JsonDecryptor: + def json(self, fn: _JsonDecryptorT) -> _JsonDecryptorT: """Register the JSON decryption handler. Example: diff --git a/libs/sdk-py/tests/test_encryption.py b/libs/sdk-py/tests/test_encryption.py index 1d4cd8e54..5fabbe952 100644 --- a/libs/sdk-py/tests/test_encryption.py +++ b/libs/sdk-py/tests/test_encryption.py @@ -1,6 +1,8 @@ +from collections.abc import Awaitable, Callable + import pytest -from langgraph_sdk import DecryptResult +from langgraph_sdk import DecryptResult, EncryptionContext from langgraph_sdk.encryption import DuplicateHandlerError, Encryption @@ -12,6 +14,27 @@ def test_decrypt_result(): assert DecryptResult(plaintext={"plain": True}).replacement is None +def test_decrypt_decorators_preserve_return_types(): + encryption = Encryption() + + @encryption.decrypt.blob + async def blob_dec(_ctx: EncryptionContext, data: bytes) -> bytes: + return data + + @encryption.decrypt.json + async def json_dec( + _ctx: EncryptionContext, data: dict[str, object] + ) -> dict[str, object]: + return data + + blob_handler: Callable[[EncryptionContext, bytes], Awaitable[bytes]] = blob_dec + json_handler: Callable[ + [EncryptionContext, dict[str, object]], Awaitable[dict[str, object]] + ] = json_dec + assert blob_handler is blob_dec + assert json_handler is json_dec + + class TestHandlerValidation: """Test duplicate handler and signature validation."""