Fix async step persistence for synchronous actions

burr/core/application.py

@@ -1115,25 +1115,29 @@ async def _astep(self, inputs: Optional[Dict[str, Any]], _run_hooks: bool = True
                     # which this is supposed to be its OK).
                     # this delegates hooks to the synchronous version, so we'll call all of them as well
                     # In this case we allow the self._step to do input processing
-                    return self._step(
+                    output = self._step(
                         inputs=inputs, _run_hooks=False
                     )  # Skip hooks as we already ran all of them/will run all of them in this function's finally
+                    if output is None:
+                        return None
+                    next_action, result, new_state = output
                 # In this case we want to process inputs because we run the function directly
-                action_inputs = self._process_inputs(inputs, next_action)
-                if next_action.single_step:
-                    result, new_state = await _arun_single_step_action(
-                        next_action, self._state, inputs=action_inputs
-                    )
                 else:
-                    result = await _arun_function(
-                        next_action,
-                        self._state,
-                        inputs=action_inputs,
-                        name=next_action.name,
-                    )
-                    new_state = _run_reducer(next_action, self._state, result, next_action.name)
-                new_state = self._update_internal_state_value(new_state, next_action)
-                self._set_state(new_state)
+                    action_inputs = self._process_inputs(inputs, next_action)
+                    if next_action.single_step:
+                        result, new_state = await _arun_single_step_action(
+                            next_action, self._state, inputs=action_inputs
+                        )
+                    else:
+                        result = await _arun_function(
+                            next_action,
+                            self._state,
+                            inputs=action_inputs,
+                            name=next_action.name,
+                        )
+                        new_state = _run_reducer(next_action, self._state, result, next_action.name)
+                    new_state = self._update_internal_state_value(new_state, next_action)
+                    self._set_state(new_state)
             except Exception as e:
                 exc = e
                 logger.exception(_format_BASE_ERROR_MESSAGE(next_action, self._state, inputs))

tests/core/test_application.py

@@ -63,9 +63,11 @@
 from burr.core.graph import Graph, GraphBuilder, Transition
 from burr.core.persistence import (
     AsyncDevNullPersister,
+    AsyncInMemoryPersister,
     BaseStateLoader,
     BaseStatePersister,
     DevNullPersister,
+    PersisterHookAsync,
     PersistedStateData,
     SQLLitePersister,
 )
@@ -1900,6 +1902,60 @@ async def test_app_astep():
     assert state[PRIOR_STEP] == "counter_async"  # internal contract, not part of the public API
 
 
+async def test_app_astep_sync_action_persists_executed_state():
+    persister = AsyncInMemoryPersister()
+    tracker = ActionTrackerAsync()
+    counter_action = base_counter_action.with_name("counter")
+    app = await (
+        ApplicationBuilder()
+        .with_actions(counter_action)
+        .with_transitions()
+        .with_entrypoint("counter")
+        .with_state(count=0)
+        .with_identifiers(app_id="app", partition_key="pk")
+        .with_hooks(PersisterHookAsync(persister), tracker)
+        .abuild()
+    )
+
+    action, result, state = await app.astep()
+
+    persisted_state = await persister.load("pk", "app")
+    assert action.name == "counter"
+    assert result == {"count": 1}
+    assert state["count"] == 1
+    assert app.state["count"] == 1
+    assert tracker.post_called[0][1]["result"] == {"count": 1}
+    assert tracker.post_called[0][1]["state"]["count"] == 1
+    assert tracker.post_called[0][1]["exception"] is None
+    assert persisted_state["state"]["count"] == 1
+    assert persisted_state["status"] == "completed"
+
+
+async def test_app_astep_sync_single_step_action_persists_executed_state():
+    persister = AsyncInMemoryPersister()
+    counter_action = base_single_step_counter.with_name("counter")
+    app = await (
+        ApplicationBuilder()
+        .with_actions(counter_action)
+        .with_transitions()
+        .with_entrypoint("counter")
+        .with_state(count=0, tracker=[])
+        .with_identifiers(app_id="app", partition_key="pk")
+        .with_hooks(PersisterHookAsync(persister))
+        .abuild()
+    )
+
+    _, result, state = await app.astep()
+
+    persisted_state = await persister.load("pk", "app")
+    assert result == {"count": 1}
+    assert state["count"] == 1
+    assert state["tracker"] == [1]
+    assert persisted_state["state"]["count"] == 1
+    assert persisted_state["state"]["tracker"] == [1]
+    assert persisted_state["status"] == "completed"
+
+
 def test_app_step_context():
     APP_ID = str(uuid.uuid4())
     PARTITION_KEY = str(uuid.uuid4())