(burr) branch main updated: fix: keep missing in-memory state loads out of application lists (#958)

[eb...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

elijahbenizzy pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/burr.git


The following commit(s) were added to refs/heads/main by this push:
     new 9088aad8 fix: keep missing in-memory state loads out of application lists (#958)
9088aad8 is described below

commit 9088aad8cb6de62b9923816b89eaa562c1162c65
Author: MIRS571 <25...@users.noreply.github.com>
AuthorDate: Sun Oct 18 01:17:11 2026 +0800

    fix: keep missing in-memory state loads out of application lists (#958)
    
    Generated-by: OpenAI Codex
    
    Co-authored-by: Elijah Ben Izzy <el...@users.noreply.github.com>
---
 burr/core/persistence.py       |  4 +--
 tests/core/test_persistence.py | 71 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/burr/core/persistence.py b/burr/core/persistence.py
index ef08f6b7..d0775b55 100644
--- a/burr/core/persistence.py
+++ b/burr/core/persistence.py
@@ -627,7 +627,7 @@ class InMemoryPersister(BaseStatePersister):
         if app_id is None:
             return None
 
-        if not (states := self._storage[partition_key][app_id]):
+        if not (states := self._storage.get(partition_key, {}).get(app_id)):
             return None
 
         if sequence_id is None:
@@ -680,7 +680,7 @@ class AsyncInMemoryPersister(AsyncBaseStatePersister):
         if app_id is None:
             return None
 
-        if not (states := self._storage[partition_key][app_id]):
+        if not (states := self._storage.get(partition_key, {}).get(app_id)):
             return None
 
         if sequence_id is None:
diff --git a/tests/core/test_persistence.py b/tests/core/test_persistence.py
index 0d462083..1944bbb1 100644
--- a/tests/core/test_persistence.py
+++ b/tests/core/test_persistence.py
@@ -96,6 +96,38 @@ def test_sqlite_persister_list_app_ids_is_unique_and_newest_first():
     assert persister.list_app_ids("pk") == ["a-app", "z-app"]
 
 
+@pytest.mark.parametrize("partition_key", [None, "partition_key"])
+@pytest.mark.parametrize("has_saved_app", [False, True])
+def test_memory_persistence_missing_load_preserves_app_ids(partition_key, has_saved_app):
+    persistence = InMemoryPersister()
+    if has_saved_app:
+        for sequence_id in (1, 2):
+            persistence.save(
+                partition_key,
+                "saved",
+                sequence_id,
+                "position",
+                State({"step": sequence_id}),
+                "status",
+            )
+    expected_app_ids = ["saved"] if has_saved_app else []
+
+    assert persistence.load(partition_key, None) is None
+    for sequence_id in (None, 99):
+        assert persistence.load(partition_key, "missing", sequence_id) is None
+        assert persistence.list_app_ids(partition_key) == expected_app_ids
+
+    if has_saved_app:
+        assert persistence.load(partition_key, "saved", 99) is None
+        assert persistence.load(partition_key, "saved")["state"] == State({"step": 2})
+        assert persistence.load(partition_key, "saved", 1)["state"] == State({"step": 1})
+        assert persistence.list_app_ids(partition_key) == expected_app_ids
+
+    persistence.save(partition_key, "missing", 1, "position", State({"step": 3}), "status")
+    assert set(persistence.list_app_ids(partition_key)) == set(expected_app_ids + ["missing"])
+    assert persistence.load(partition_key, "missing")["state"] == State({"step": 3})
+
+
 def test_persistence_is_initialized_false(initializing_persistence):
     assert not initializing_persistence.is_initialized()
 
@@ -298,6 +330,45 @@ async def test_async_persistence_lists_app_ids(async_persistence):
     assert set(app_ids) == set(["app_id1", "app_id2"])
 
 
+@pytest.mark.parametrize("partition_key", [None, "partition_key"])
+@pytest.mark.parametrize("has_saved_app", [False, True])
+async def test_async_persistence_missing_load_preserves_app_ids(
+    async_persistence, partition_key, has_saved_app
+):
+    if has_saved_app:
+        for sequence_id in (1, 2):
+            await async_persistence.save(
+                partition_key,
+                "saved",
+                sequence_id,
+                "position",
+                State({"step": sequence_id}),
+                "status",
+            )
+    expected_app_ids = ["saved"] if has_saved_app else []
+
+    assert await async_persistence.load(partition_key, None) is None
+    for sequence_id in (None, 99):
+        assert await async_persistence.load(partition_key, "missing", sequence_id) is None
+        assert await async_persistence.list_app_ids(partition_key) == expected_app_ids
+
+    if has_saved_app:
+        assert await async_persistence.load(partition_key, "saved", 99) is None
+        assert (await async_persistence.load(partition_key, "saved"))["state"] == State({"step": 2})
+        assert (await async_persistence.load(partition_key, "saved", 1))["state"] == State(
+            {"step": 1}
+        )
+        assert await async_persistence.list_app_ids(partition_key) == expected_app_ids
+
+    await async_persistence.save(
+        partition_key, "missing", 1, "position", State({"step": 3}), "status"
+    )
+    assert set(await async_persistence.list_app_ids(partition_key)) == set(
+        expected_app_ids + ["missing"]
+    )
+    assert (await async_persistence.load(partition_key, "missing"))["state"] == State({"step": 3})
+
+
 @pytest.mark.parametrize(
     "method_name,kwargs",
     [