Welcome Json - bye String #95
Conversation
Poort attempt to fix redis#69
|
Hi, I’m Jit, a friendly security platform designed to help developers build secure applications from day zero with an MVS (Minimal viable security) mindset. In case there are security findings, they will be communicated to you as a comment inside the PR. Hope you’ll enjoy using Jit. Questions? Comments? Want to learn more? Get in touch with us. |
|
For backward compatibility, adding a hybrid strategy : Lazy migration when an old string type is met -> migrated to JSON -> recheck existing type string, if none, it activates a fast path. On startup, check if at least one string type exists without blocking a long startup in case of many keys to convert |
- Cli for migration (very large number of keys) - Benchmark - Early exit - edge case bench (1M json key - 1 string key)
|
Edge case 1M json key - 1 string key in the end
|
There was a problem hiding this comment.
Pull request overview
This PR migrates working memory storage from Redis STRING format (JSON serialized as strings) to native Redis JSON format using the RedisJSON module. The migration includes backward compatibility through lazy migration and a CLI command for bulk migration.
Key Changes:
- Removed custom JSON datetime handler in favor of pydantic's built-in serialization
- Implemented lazy migration with early-exit optimization for startup performance
- Added CLI command
migrate-working-memoryfor bulk migration of existing keys
Reviewed changes
Copilot reviewed 6 out of 7 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| agent_memory_server/working_memory.py | Core migration logic with global state tracking, lazy migration on read, and native JSON storage for writes |
| agent_memory_server/cli.py | New CLI command for bulk migration with dry-run support and progress reporting |
| agent_memory_server/main.py | Application startup now checks migration status to enable/disable lazy migration |
| tests/test_working_memory_strategies.py | Updated mocks to use redis.json() API instead of string get/set |
| tests/test_working_memory.py | Added comprehensive migration tests including backward compatibility scenarios |
| tests/benchmarks/test_migration_benchmark.py | New benchmark suite to measure migration performance at scale |
| .gitignore | Added pattern to ignore augment markdown files |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| click.echo("\n--dry-run specified, no changes made.") | ||
| return | ||
|
|
||
| # Migrate keys in batches using pipeline | ||
| click.echo(f"\nMigrating {len(string_keys)} keys...") | ||
| migrate_start = time.time() | ||
| migrated = 0 | ||
| errors = 0 | ||
|
|
||
| # Process in batches | ||
| for batch_start in range(0, len(string_keys), batch_size): | ||
| batch_keys = string_keys[batch_start : batch_start + batch_size] | ||
|
|
||
| # First, read all string data in a pipeline | ||
| read_pipe = redis.pipeline() | ||
| for key in batch_keys: |
There was a problem hiding this comment.
The migration command does not preserve TTL when migrating keys from string to JSON format. Before deleting the old key (line 193), get and preserve its TTL:
# Get TTL before deleting
ttl = await redis.ttl(key)
await redis.delete(key)
await redis.json().set(key, "$", data)
# Restore TTL if it was set (ttl > 0)
if ttl > 0:
await redis.expire(key, ttl)This ensures keys with expiration don't become permanent after migration.
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
| # Flag to track if all string keys have been migrated to JSON | ||
| # When True, we skip the type() check and go straight to json().get() | ||
| _string_keys_migrated: bool = False | ||
|
|
||
| def json_datetime_handler(obj): | ||
| """JSON serializer for datetime objects.""" | ||
| if isinstance(obj, datetime): | ||
| return obj.isoformat() | ||
| raise TypeError(f"Object of type {type(obj)} is not JSON serializable") | ||
| # Counter for remaining string keys (avoids re-scanning after each migration) | ||
| _remaining_string_keys: int = 0 |
There was a problem hiding this comment.
The global variables _string_keys_migrated and _remaining_string_keys are not thread-safe and won't be shared across multiple uvicorn workers. In a multi-worker deployment, each worker process will have its own copy of these variables, causing:
- Redundant migration checks on startup (each worker scans independently)
- Race conditions when decrementing the counter during lazy migration
- Inconsistent migration state across workers
Consider using Redis itself to store the migration status (e.g., a Redis key like "migration:working_memory:complete") or use proper inter-process synchronization mechanisms.
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
| _string_keys_migrated = True | ||
| _remaining_string_keys = 0 | ||
| return True | ||
|
|
||
|
|
||
| def _decrement_string_key_count() -> None: | ||
| """ | ||
| Decrement the string key counter after a successful migration. | ||
| Note: When _remaining_string_keys is -1, we don't know the exact count | ||
| (early exit mode). In this case, lazy migration stays enabled until | ||
| the migration script is run. | ||
| """ | ||
| global _string_keys_migrated, _remaining_string_keys | ||
|
|
||
| # If we don't know the count (-1), we can't track completion | ||
| # The migration script will set the flag when done | ||
| if _remaining_string_keys == -1: | ||
| return | ||
|
|
There was a problem hiding this comment.
The function _decrement_string_key_count() modifies global state (_remaining_string_keys, _string_keys_migrated) without any synchronization mechanism. In a concurrent environment, multiple async tasks could call this function simultaneously, leading to:
- Lost decrements (two tasks read the same value, decrement, and write back)
- Race conditions when checking if migration is complete
Consider using an asyncio.Lock or storing the counter in Redis with atomic operations (e.g., DECR command).
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
There was a problem hiding this comment.
actually a good idea...
| # Delete the old string key and set as JSON | ||
| await redis_client.delete(key) | ||
| await redis_client.json().set(key, "$", data) | ||
|
|
||
| # Preserve TTL if it was set | ||
| # Note: TTL is lost during migration since we deleted the key | ||
| # The next set_working_memory call will restore it if configured |
There was a problem hiding this comment.
TTL is lost during migration from string to JSON format. The comment acknowledges this (line 164-165 in original file), but this means that working memory entries with TTL will have their expiration removed during migration, potentially causing them to persist indefinitely until explicitly deleted.
Consider preserving the TTL during migration by:
- Getting the TTL before deleting the old key:
ttl = await redis_client.ttl(key) - After setting the JSON key, restore the TTL if it was set:
if ttl > 0: await redis_client.expire(key, ttl)
| # Delete the old string key and set as JSON | |
| await redis_client.delete(key) | |
| await redis_client.json().set(key, "$", data) | |
| # Preserve TTL if it was set | |
| # Note: TTL is lost during migration since we deleted the key | |
| # The next set_working_memory call will restore it if configured | |
| # Preserve TTL if it was set | |
| ttl = await redis_client.ttl(key) | |
| # Delete the old string key and set as JSON | |
| await redis_client.delete(key) | |
| await redis_client.json().set(key, "$", data) | |
| # Restore TTL if it was set (ttl > 0) | |
| if ttl > 0: | |
| await redis_client.expire(key, ttl) |
|
|
||
| # Scan for working_memory:* keys | ||
| cursor = 0 | ||
| json_keys_found = 0 |
There was a problem hiding this comment.
Variable string_keys_found is not used.
There was a problem hiding this comment.
Well, i disagree ;-)
|
Ok one last touch for those who go with the cli to migrate and no longer need any backward compatibilty logic : introducing yet another env variable |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…nt-memory-server into working_memory_json
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
|
@fcenedes I'm having trouble pushing a fix to your branch, but basically, you want to run the pre-commit formatter: |
|
Feel free to add a |
|
done |
both cli and lazy migration test
well... light refactor isn't it ? ;-)
|
Confirmed that tests pass locally (API keys issue with PRs). |
2 participants
Poor attempt to fix #69
Changes