-
Notifications
You must be signed in to change notification settings - Fork 26
fix: always refresh TTL when refresh_on_read is enabled #95
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
+589
−17
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||
|---|---|---|---|---|
| @@ -0,0 +1,305 @@ | ||||
| """Test async TTL synchronization behavior for AsyncRedisSaver.""" | ||||
|
|
||||
| import asyncio | ||||
| import time | ||||
| from contextlib import asynccontextmanager | ||||
| from typing import AsyncGenerator | ||||
| from uuid import uuid4 | ||||
|
|
||||
| import pytest | ||||
| from langgraph.checkpoint.base import Checkpoint | ||||
|
|
||||
| from langgraph.checkpoint.redis.aio import AsyncRedisSaver | ||||
|
|
||||
|
|
||||
| @asynccontextmanager | ||||
| async def _saver( | ||||
| redis_url: str, ttl_config: dict | ||||
| ) -> AsyncGenerator[AsyncRedisSaver, None]: | ||||
| """Create an AsyncRedisSaver with the given TTL configuration.""" | ||||
| async with AsyncRedisSaver.from_conn_string(redis_url, ttl=ttl_config) as saver: | ||||
| await saver.setup() | ||||
| yield saver | ||||
|
|
||||
|
|
||||
| @pytest.mark.asyncio | ||||
| async def test_async_ttl_refresh_on_read(redis_url: str) -> None: | ||||
| """Test that TTL is always refreshed when refresh_on_read is enabled (async).""" | ||||
|
|
||||
| # Configure with TTL refresh on read | ||||
| ttl_config = { | ||||
| "default_ttl": 2, # 2 minutes = 120 seconds | ||||
| "refresh_on_read": True, | ||||
| } | ||||
|
|
||||
| async with _saver(redis_url, ttl_config) as saver: | ||||
| thread_id = str(uuid4()) | ||||
| checkpoint_ns = "" | ||||
| checkpoint_id = str(uuid4()) | ||||
|
|
||||
| # Create a checkpoint | ||||
| checkpoint = Checkpoint( | ||||
| v=1, | ||||
| id=checkpoint_id, | ||||
| ts="2024-01-01T00:00:00+00:00", | ||||
| channel_values={"test": "value"}, | ||||
| channel_versions={}, | ||||
| versions_seen={}, | ||||
| ) | ||||
|
|
||||
| config = { | ||||
| "configurable": { | ||||
| "thread_id": thread_id, | ||||
| "checkpoint_ns": checkpoint_ns, | ||||
| } | ||||
| } | ||||
|
|
||||
| # Save the checkpoint | ||||
| saved_config = await saver.aput(config, checkpoint, {"test": "metadata"}, {}) | ||||
|
|
||||
| # Get the checkpoint key | ||||
| from langgraph.checkpoint.redis.base import BaseRedisSaver | ||||
| from langgraph.checkpoint.redis.util import ( | ||||
| to_storage_safe_id, | ||||
| to_storage_safe_str, | ||||
| ) | ||||
|
|
||||
| checkpoint_key = BaseRedisSaver._make_redis_checkpoint_key( | ||||
| to_storage_safe_id(thread_id), | ||||
| to_storage_safe_str(checkpoint_ns), | ||||
| to_storage_safe_id(saved_config["configurable"]["checkpoint_id"]), | ||||
| ) | ||||
|
|
||||
| # Check initial TTL (should be around 120 seconds) | ||||
| initial_ttl = await saver._redis.ttl(checkpoint_key) | ||||
| assert ( | ||||
| 115 <= initial_ttl <= 120 | ||||
| ), f"Initial TTL should be ~120s, got {initial_ttl}" | ||||
|
|
||||
| # Wait a bit (simulate time passing) | ||||
| await asyncio.sleep(2) | ||||
|
|
||||
| # Read the checkpoint - this should refresh TTL to full value | ||||
| result = await saver.aget_tuple(saved_config) | ||||
| assert result is not None | ||||
|
|
||||
| # Check TTL after read - should be refreshed to full value | ||||
| refreshed_ttl = await saver._redis.ttl(checkpoint_key) | ||||
| assert ( | ||||
| 115 <= refreshed_ttl <= 120 | ||||
| ), f"TTL should be refreshed to ~120s, got {refreshed_ttl}" | ||||
|
|
||||
|
|
||||
| @pytest.mark.asyncio | ||||
| async def test_async_ttl_no_refresh_when_disabled(redis_url: str) -> None: | ||||
| """Test that TTL is not refreshed when refresh_on_read is disabled (async).""" | ||||
|
|
||||
| # Configure without TTL refresh on read | ||||
| ttl_config = { | ||||
| "default_ttl": 2, # 2 minutes = 120 seconds | ||||
| "refresh_on_read": False, # Don't refresh TTL on read | ||||
| } | ||||
|
|
||||
| async with _saver(redis_url, ttl_config) as saver: | ||||
| thread_id = str(uuid4()) | ||||
| checkpoint_ns = "" | ||||
| checkpoint_id = str(uuid4()) | ||||
|
|
||||
| # Create a checkpoint | ||||
| checkpoint = Checkpoint( | ||||
| v=1, | ||||
| id=checkpoint_id, | ||||
| ts="2024-01-01T00:00:00+00:00", | ||||
| channel_values={"test": "value"}, | ||||
| channel_versions={}, | ||||
| versions_seen={}, | ||||
| ) | ||||
|
|
||||
| config = { | ||||
| "configurable": { | ||||
| "thread_id": thread_id, | ||||
| "checkpoint_ns": checkpoint_ns, | ||||
| } | ||||
| } | ||||
|
|
||||
| # Save the checkpoint | ||||
| saved_config = await saver.aput(config, checkpoint, {"test": "metadata"}, {}) | ||||
|
|
||||
| # Get the checkpoint key | ||||
| from langgraph.checkpoint.redis.base import BaseRedisSaver | ||||
| from langgraph.checkpoint.redis.util import ( | ||||
| to_storage_safe_id, | ||||
| to_storage_safe_str, | ||||
| ) | ||||
|
||||
| ) |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[nitpick] The import statements for checkpoint key generation are repeated in multiple test functions. Consider moving these imports to the top of the file or creating a helper function to reduce code duplication.