diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..2c6b20802 100644 --- a/burr/core/application.py +++ b/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)) diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 9313cefc9..14fa3e3a9 100644 --- a/tests/core/test_application.py +++ b/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())