diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..415e0984e 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -1113,27 +1113,30 @@ async def _astep(self, inputs: Optional[Dict[str, Any]], _run_hooks: bool = True # but that's safer than assuming its OK to launch a thread # TODO -- add an option/configuration to launch a thread (yikes, not super safe, but for a pure function # 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( - inputs=inputs, _run_hooks=False - ) # Skip hooks as we already ran all of them/will run all of them in this function's finally + # Execute the synchronous step without hooks here. Async hooks still run + # afterward so async persisters can save the completed state of synchronous steps. + output = self._step(inputs=inputs, _run_hooks=False) + 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..5a252f778 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -63,10 +63,12 @@ from burr.core.graph import Graph, GraphBuilder, Transition from burr.core.persistence import ( AsyncDevNullPersister, + AsyncInMemoryPersister, BaseStateLoader, BaseStatePersister, DevNullPersister, PersistedStateData, + PersisterHookAsync, SQLLitePersister, ) from burr.core.typing import TypingSystem @@ -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())