fix(core): start at the default entrypoint when resume_at_next_action is False

burr/core/application.py

@@ -2707,8 +2707,9 @@ def _init_state_from_persister(
                     self.start = last_position
                     self.reset_to_entrypoint()
             else:
-                # self.start is already set to the default. We don't need to do anything.
-                pass
+                # self.start is already set to the default, but the loaded state still has the
+                # prior step, which would otherwise decide the next action. Clear it.
+                self.reset_to_entrypoint()
 
     def _load_from_sync_persister(self):
         """Loads from the set sync persister and into this current object.

tests/core/test_application.py

@@ -67,6 +67,7 @@
     BaseStateLoader,
     BaseStatePersister,
     DevNullPersister,
+    InMemoryPersister,
     PersistedStateData,
     PersisterHookAsync,
     SQLLitePersister,
@@ -4454,3 +4455,76 @@ def noop(state: State) -> State:
     app = builder.build()
 
     assert app.state["x"] == 100
+
+
+@pytest.mark.parametrize("last_step_failed", [False, True])
+@pytest.mark.parametrize(
+    "resume_at_next_action,expected_next_action",
+    [(True, "result"), (False, "counter")],
+)
+def test_initialize_from_resume_at_next_action(
+    resume_at_next_action, expected_next_action, last_step_failed
+):
+    """resume_at_next_action=False should start at the default entrypoint, keeping the loaded state,
+    whether the last saved step completed or failed"""
+    persister = InMemoryPersister()
+    result = base_broken_action if last_step_failed else Result("count")
+
+    def build(**initialize_kwargs):
+        return (
+            ApplicationBuilder()
+            .with_actions(counter=base_counter_action, result=result)
+            .with_transitions(("counter", "result", default))
+            .initialize_from(
+                persister, default_state={}, default_entrypoint="counter", **initialize_kwargs
+            )
+            .with_identifiers(app_id="app", partition_key="pk")
+            .with_state_persister(persister)
+            .build()
+        )
+
+    build(resume_at_next_action=True).step()
+    if last_step_failed:
+        with pytest.raises(BrokenStepException):
+            build(resume_at_next_action=True).step()
+    app = build(resume_at_next_action=resume_at_next_action)
+    assert app.get_next_action().name == expected_next_action
+    assert app.state["count"] == 1
+    if not resume_at_next_action:
+        assert PRIOR_STEP not in app.state.get_all()
+
+
+@pytest.mark.parametrize("last_step_failed", [False, True])
+@pytest.mark.parametrize(
+    "resume_at_next_action,expected_next_action",
+    [(True, "result"), (False, "counter")],
+)
+async def test_initialize_from_resume_at_next_action_async(
+    resume_at_next_action, expected_next_action, last_step_failed
+):
+    """Same as above, for the async builder"""
+    persister = AsyncInMemoryPersister()
+    result = base_broken_action_async if last_step_failed else Result("count")
+
+    async def abuild(**initialize_kwargs):
+        return await (
+            ApplicationBuilder()
+            .with_actions(counter=base_counter_action_async, result=result)
+            .with_transitions(("counter", "result", default))
+            .initialize_from(
+                persister, default_state={}, default_entrypoint="counter", **initialize_kwargs
+            )
+            .with_identifiers(app_id="app", partition_key="pk")
+            .with_state_persister(persister)
+            .abuild()
+        )
+
+    await (await abuild(resume_at_next_action=True)).astep()
+    if last_step_failed:
+        with pytest.raises(BrokenStepException):
+            await (await abuild(resume_at_next_action=True)).astep()
+    app = await abuild(resume_at_next_action=resume_at_next_action)
+    assert app.get_next_action().name == expected_next_action
+    assert app.state["count"] == 1
+    if not resume_at_next_action:
+        assert PRIOR_STEP not in app.state.get_all()