fix(checkpoint): clear in-memory store after successful persistence (#751)

* fix(checkpoint): clear in-memory store after successful persistence

* test(checkpoint): add unit test for memory leak check

* Update tests/unit/checkpoint/test_memory_leak.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Willem Jiang
2025-12-09 23:32:13 +08:00
committed by GitHub
parent fde7a69562
commit 84c449cf79
2 changed files with 60 additions and 2 deletions

View File

@@ -210,14 +210,26 @@ class ChatStreamManager:
return False
# Choose persistence method based on available connection
success = False
if self.mongo_db is not None:
return self._persist_to_mongodb(thread_id, messages)
success = self._persist_to_mongodb(thread_id, messages)
elif self.postgres_conn is not None:
return self._persist_to_postgresql(thread_id, messages)
success = self._persist_to_postgresql(thread_id, messages)
else:
self.logger.warning("No database connection available")
return False
if success:
try:
for item in memories:
self.store.delete(store_namespace, item.key)
except Exception as e:
self.logger.error(
f"Error cleaning up memory store for thread {thread_id}: {e}"
)
return success
except Exception as e:
self.logger.error(
f"Error persisting conversation for thread {thread_id}: {e}"