diff --git a/burr/core/application.py b/burr/core/application.py index 91ef8b1b6..aa76d14f1 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -958,17 +958,26 @@ def _context_factory(self, action: Action, sequence_id: int) -> ApplicationConte ) def _step( - self, inputs: Optional[Dict[str, Any]], _run_hooks: bool = True + self, + inputs: Optional[Dict[str, Any]], + _run_hooks: bool = True, + _processed_inputs: Optional[Dict[str, Any]] = None, ) -> Optional[Tuple[Action, dict, State]]: """Internal-facing version of step. This is the same as step, but with an additional - parameter to hide hook execution so async can leverage it.""" + parameter to hide hook execution so async can leverage it. Processed inputs can also be + supplied so callers that already processed them (async `astep`) reuse the same injected + dependencies for the hooks and the action instead of invoking the factories twice.""" with self.context: next_action = self.get_next_action() if next_action is None: return None if inputs is None: inputs = {} - action_inputs = self._process_inputs(inputs, next_action) + action_inputs = ( + self._process_inputs(inputs, next_action) + if _processed_inputs is None + else _processed_inputs + ) if _run_hooks: self._adapter_set.call_all_lifecycle_hooks_sync( "pre_run_step", @@ -1105,12 +1114,13 @@ async def _astep(self, inputs: Optional[Dict[str, Any]], _run_hooks: bool = True return None if inputs is None: inputs = {} + action_inputs = self._process_inputs(inputs, next_action) if _run_hooks: await self._adapter_set.call_all_lifecycle_hooks_sync_and_async( "pre_run_step", action=next_action, state=self._state, - inputs=inputs, + inputs=action_inputs, sequence_id=self.sequence_id, app_id=self._uid, partition_key=self._partition_key, @@ -1124,16 +1134,17 @@ 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). - # In this case we allow the self._step to do input processing # 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) + output = self._step( + inputs=inputs, + _run_hooks=False, + _processed_inputs=action_inputs, + ) 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 else: - 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 diff --git a/tests/core/test_application.py b/tests/core/test_application.py index efb7d205b..b7400aa24 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -3401,6 +3401,155 @@ async def test_application_run_step_runs_hooks(): assert len(hooks[0].pre_run_execute_calls) == 1 +async def test_astep_processes_inputs_before_pre_run_step_hook(): + class InputCaptureHook(PreRunStepHook): + def __init__(self): + self.captured_inputs = [] + + def pre_run_step(self, *, inputs, **kwargs): + self.captured_inputs.append(inputs) + + @action(reads=[], writes=['val']) + def sync_noop(state: State) -> State: + return state.update(val=1) + + @action(reads=[], writes=['val']) + async def async_noop(state: State) -> State: + return state.update(val=2) + + # 1. Test sync action executed via astep() + sync_hook = InputCaptureHook() + app_sync = ( + ApplicationBuilder() + .with_actions(sync_noop) + .with_entrypoint('sync_noop') + .with_hooks(sync_hook) + .build() + ) + await app_sync.astep(inputs={'unused': 123}) + assert sync_hook.captured_inputs == [{}] + + # 2. Test async action executed via astep() + async_hook = InputCaptureHook() + app_async = ( + ApplicationBuilder() + .with_actions(async_noop) + .with_entrypoint('async_noop') + .with_hooks(async_hook) + .build() + ) + await app_async.astep(inputs={'unused': 123}) + assert async_hook.captured_inputs == [{}] + + +async def test_astep_validates_inputs_before_pre_run_step_hook(): + class InputCaptureHook(PreRunStepHook): + def __init__(self): + self.call_count = 0 + + def pre_run_step(self, *, inputs, **kwargs): + self.call_count += 1 + + @action(reads=[], writes=['val']) + def sync_requires_param(state: State, required_param: int) -> State: + return state.update(val=required_param) + + @action(reads=[], writes=['val']) + async def async_requires_param(state: State, required_param: int) -> State: + return state.update(val=required_param) + + # 1. Sync action under astep: missing required input should raise ValueError before pre_run_step + hook_sync = InputCaptureHook() + app_sync = ( + ApplicationBuilder() + .with_actions(sync_requires_param) + .with_entrypoint('sync_requires_param') + .with_hooks(hook_sync) + .build() + ) + with pytest.raises(ValueError): + await app_sync.astep(inputs={}) + assert hook_sync.call_count == 0 + + # 2. Async action under astep: missing required input should raise ValueError before pre_run_step + hook_async = InputCaptureHook() + app_async = ( + ApplicationBuilder() + .with_actions(async_requires_param) + .with_entrypoint('async_requires_param') + .with_hooks(hook_async) + .build() + ) + with pytest.raises(ValueError): + await app_async.astep(inputs={}) + assert hook_async.call_count == 0 + + +@pytest.mark.parametrize( + 'execute', ['step', 'astep_with_sync_action', 'astep_with_async_action'] +) +async def test_step_and_astep_share_one_processed_inputs_dict(execute: str): + """The pre_run_step hook and the action must receive the same processed inputs. + + Regression for the sync branch of ``astep``: it used to delegate with + ``_step(inputs=inputs)``, which ran ``_process_inputs`` a second time. An action + that declares ``__context`` therefore got a second object from the dependency + factory, so the hook and the action did not share the same instance (#979). + """ + created_contexts = [] + hook_contexts = [] + action_contexts = [] + + def counting_context_factory(action, sequence_id): + context = object() + created_contexts.append(context) + return context + + @action(reads=[], writes=['val']) + def sync_with_context(state: State, __context) -> State: + action_contexts.append(__context) + return state.update(val=1) + + @action(reads=[], writes=['val']) + async def async_with_context(state: State, __context) -> State: + action_contexts.append(__context) + return state.update(val=2) + + class ContextCaptureHook(PreRunStepHook): + def pre_run_step(self, *, inputs, **kwargs): + hook_contexts.append(inputs['__context']) + + if execute == 'astep_with_async_action': + entrypoint = 'async_with_context' + target_action = async_with_context + else: + entrypoint = 'sync_with_context' + target_action = sync_with_context + + app = ( + ApplicationBuilder() + .with_actions(target_action) + .with_entrypoint(entrypoint) + .with_hooks(ContextCaptureHook()) + .build() + ) + # Swap in a factory that returns a fresh object on every call, so a second run of + # _process_inputs is visible both as a call count and as a distinct instance. + app._dependency_factory['__context'] = counting_context_factory + + if execute == 'step': + assert app.step() is not None + else: + await app.astep() + + assert len(created_contexts) == 1, ( + f'__context factory ran {len(created_contexts)} times, expected once' + ) + assert hook_contexts == created_contexts + assert action_contexts == created_contexts + assert hook_contexts[0] is action_contexts[0] + + def test_application_post_application_create_hook(): class PostApplicationCreateTracker(PostApplicationCreateHook): def __init__(self):