diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 7a365ec122..297cffba75 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -81,6 +81,9 @@ def __init__( self._iteration = 0 self._max_iterations = max_iterations self._state = state + # When True, Workflow.run must reject successors even if the ResponseStream + # weakref is already gone (cleanup still in progress after a drop/cancel). + self._blocking_reuse_until_cleanup = False # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None @@ -128,51 +131,85 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: # Run iteration concurrently with live event streaming: we poll # for new events while the iteration coroutine progresses. iteration_task = asyncio.create_task(self._run_iteration()) + # Track commit so cancel/abort cleanup spans the whole superstep + # (polling → await iteration → drain → commit), not only the poll loop (#7859). + committed = False + # Defer failure events until after discard so dropping the ResponseStream + # cannot race a successor commit of stale pending writes (#7859). + deferred_failure_events: list[WorkflowEvent] = [] try: + self._blocking_reuse_until_cleanup = True while not iteration_task.done(): try: # Wait briefly for any new event; timeout allows progress checks event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05) - yield event + if event.type == "executor_failed": + deferred_failure_events.append(event) + else: + yield event except asyncio.TimeoutError: # Periodically continue to let iteration advance continue - except asyncio.CancelledError: - # Propagate cancellation to the iteration task to avoid orphaned work - iteration_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await iteration_task - raise - # Propagate errors from iteration, but first surface any pending events - try: - await iteration_task - except Exception: - # Make sure failure-related events (like ExecutorFailedEvent) are surfaced - if await self._ctx.has_events(): - for event in await self._ctx.drain_events(): + # Propagate errors from iteration, but first surface any pending events + try: + await iteration_task + except Exception: + # Discard staged writes immediately — before any await/yield — so a + # streaming consumer that stops after executor_failed cannot leave + # pending state for a later run to commit (#7859). + self._state.discard() + for event in deferred_failure_events: yield event - raise - self._iteration += 1 + deferred_failure_events.clear() + if await self._ctx.has_events(): + for event in await self._ctx.drain_events(): + yield event + raise - # Drain any straggler events emitted at tail end - if await self._ctx.has_events(): - for event in await self._ctx.drain_events(): + for event in deferred_failure_events: yield event + deferred_failure_events.clear() - logger.info(f"Completed superstep {self._iteration}") - - # Commit pending state changes at superstep boundary - self._state.commit() + self._iteration += 1 - # Create checkpoint after each superstep iteration - await self.create_checkpoint_if_enabled() - - yield WorkflowEvent.superstep_completed(iteration=self._iteration) + # Drain any straggler events emitted at tail end + if await self._ctx.has_events(): + for event in await self._ctx.drain_events(): + yield event - # Check for convergence: no more messages to process - if not await self._ctx.has_messages(): - break + logger.info(f"Completed superstep {self._iteration}") + + # Commit pending state changes at superstep boundary + self._state.commit() + + # Create checkpoint after each superstep iteration. Keep + # ``committed`` false until this returns so residual pending + # staged during checkpoint prep is discarded on cancel/failure. + await self.create_checkpoint_if_enabled() + committed = True + + yield WorkflowEvent.superstep_completed(iteration=self._iteration) + + # Check for convergence: no more messages to process + if not await self._ctx.has_messages(): + break + except BaseException: + # Cancel during poll/drain, or an iteration task that ends cancelled, + # must still abandon staged writes before the commit boundary (#7859). + # Await cleanup with broad suppression so a raising executor ``finally`` + # cannot skip discard below. + if not iteration_task.done(): + iteration_task.cancel() + with contextlib.suppress(BaseException): + await iteration_task + raise + finally: + # Always discard on abort — including when iteration-task await above + # raised from executor cleanup — so staged writes cannot leak (#7859). + if not committed: + self._state.discard() + self._blocking_reuse_until_cleanup = False logger.info(f"Workflow completed after {self._iteration} supersteps") @@ -236,10 +273,18 @@ async def _prepare_checkpoint_state(self) -> None: This is used by checkpoint capture paths that need a complete, restorable state payload without necessarily writing to a checkpoint storage backend. + + If staging executor/edge state fails or is cancelled before ``commit()``, + discard residual pending writes so a later run cannot commit a partial + checkpoint payload (#7859). """ - await self._save_executor_states() - self._save_edge_runner_states() - self._state.commit() + try: + await self._save_executor_states() + self._save_edge_runner_states() + self._state.commit() + except BaseException: + self._state.discard() + raise async def create_checkpoint_if_enabled(self) -> None: """Create a checkpoint and save the checkpoint to the configured storage if one is configured. @@ -250,9 +295,11 @@ async def create_checkpoint_if_enabled(self) -> None: if not self._ctx.has_checkpointing(): return + prepared = False try: # Save executor states into committed state before creating the checkpoint. await self._prepare_checkpoint_state() + prepared = True checkpoint_id = await self._ctx.create_checkpoint( self._workflow_name, @@ -270,6 +317,10 @@ async def create_checkpoint_if_enabled(self) -> None: ) self._previous_checkpoint_id = checkpoint_id except Exception as e: + # ``_prepare_checkpoint_state`` discards on its own failure; if we + # never reached that commit, clear any residual pending here too. + if not prepared: + self._state.discard() logger.warning( "Failed to create checkpoint at iteration %d: %s. " "Note that this does not fail the workflow run. " diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index afbf45af11..c6d0b0cc1e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -409,6 +409,9 @@ def __init__( # ever iterating, the weakref dereferences to ``None`` once Python collects it, # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # Strong hold for non-streaming exclusive entry points (e.g. cancel_pending_requests) + # that mutate the same runner/state without installing a ResponseStream weakref. + self._exclusive_run_hold: bool = False @property def status(self) -> WorkflowRunState: @@ -1323,6 +1326,14 @@ async def cancel_pending_requests( if not all(isinstance(request_id, str) and request_id for request_id in selected_ids): raise ValueError("Pending workflow request IDs must be non-empty strings.") + # Share the same active/cleanup lock as ``run()`` so a cancellation continuation + # cannot start during ResponseStream cleanup (or while another run holds the lock) + # and commit pending State before the abandoned run discards it (#7859). + if self._is_run_active(): + raise WorkflowException( + "Workflow is already running; concurrent runs are not allowed on the same instance." + ) + async def apply_cancellations() -> None: cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) for request_id, request_event in cancelled_events.items(): @@ -1337,12 +1348,16 @@ async def apply_cancellations() -> None: ) await executor._cancel_pending_request(request_id, context) # pyright: ignore[reportPrivateUsage] - if checkpoint_storage is not None: - self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) - runtime_tools = normalize_tools(tools) if tools is not None else None - self._runner.context.set_runtime_tools(runtime_tools) + # Acquire the hold first, but keep setup inside try/finally so a failure in + # normalize_tools / checkpoint restore cannot leave the workflow permanently locked. + self._exclusive_run_hold = True events: list[WorkflowEvent[Any]] = [] + runtime_tools = None try: + if checkpoint_storage is not None: + self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + runtime_tools = normalize_tools(tools) if tools is not None else None + self._runner.context.set_runtime_tools(runtime_tools) if checkpoint_id is not None: await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) async for event in self._run_workflow_with_tracing( @@ -1357,6 +1372,7 @@ async def apply_cancellations() -> None: continue events.append(event) finally: + self._exclusive_run_hold = False if checkpoint_storage is not None: self._runner.context.clear_runtime_checkpoint_storage() self._runner.context.clear_runtime_tools() @@ -1411,5 +1427,11 @@ def _is_run_active(self) -> bool: Returns: True if a run is active, False otherwise. """ + if self._exclusive_run_hold: + return True + # Runner cleanup can outlive a dropped ResponseStream (weakref cleared on GC). + # Keep the instance reserved until pending State is discarded (#7859). + if getattr(self._runner, "_blocking_reuse_until_cleanup", False): + return True existing_stream = self._active_run() if self._active_run is not None else None return existing_stream is not None diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index c5219a6744..31120a91d0 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -1496,3 +1496,104 @@ async def test_runner_drains_straggler_events_at_iteration_end(): output_events = [e for e in events if e.type == "output"] # We should have output events from both executors assert len(output_events) >= 2 + +@pytest.mark.asyncio +async def test_failed_superstep_discards_pending_state_before_next_run() -> None: + """Pending State writes from a failed superstep must not leak into a later run (#7859).""" + from agent_framework import WorkflowBuilder + + @dataclass + class Msg: + fail: bool + + class FlakyExecutor(Executor): + @handler + async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None: + if message.fail: + ctx.set_state("secret", "leaked-from-failed-run") + raise RuntimeError("simulated transient failure") + await ctx.yield_output("ok") + + workflow = WorkflowBuilder(start_executor=FlakyExecutor(id="flaky")).build() + + with pytest.raises(RuntimeError, match="simulated transient failure"): + async for _ in workflow.run(Msg(fail=True), stream=True): + pass + + async for _ in workflow.run(Msg(fail=False), stream=True): + pass + + committed = workflow._runner.state.export_state() # pyright: ignore[reportPrivateUsage] + assert "secret" not in committed + + +@pytest.mark.asyncio +async def test_cancelled_superstep_discards_pending_state_before_next_run() -> None: + """Pending State writes from a cancelled superstep must not leak into a later run (#7859).""" + from agent_framework import WorkflowBuilder + + @dataclass + class Msg: + cancel: bool + + started = asyncio.Event() + + class StagingThenBlockingExecutor(Executor): + @handler + async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None: + if message.cancel: + ctx.set_state("secret", "leaked-from-cancelled-run") + started.set() + await asyncio.sleep(3600) + await ctx.yield_output("ok") + + workflow = WorkflowBuilder(start_executor=StagingThenBlockingExecutor(id="blocker")).build() + + async def run_and_cancel() -> None: + async for _ in workflow.run(Msg(cancel=True), stream=True): + pass + + task = asyncio.create_task(run_and_cancel()) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + async for _ in workflow.run(Msg(cancel=False), stream=True): + pass + + committed = workflow._runner.state.export_state() # pyright: ignore[reportPrivateUsage] + assert "secret" not in committed + + +@pytest.mark.asyncio +async def test_failed_superstep_discards_even_if_executor_cleanup_raises() -> None: + """Executor cleanup that raises after set_state must not skip State.discard (#7859).""" + from agent_framework import WorkflowBuilder + + @dataclass + class Msg: + fail: bool + + class CleanupRaisesExecutor(Executor): + @handler + async def run(self, message: Msg, ctx: WorkflowContext[Msg, str]) -> None: + if message.fail: + ctx.set_state("secret", "leaked-from-cleanup-raise") + try: + raise RuntimeError("primary failure") + finally: + raise RuntimeError("cleanup failure") # noqa: B012 + await ctx.yield_output("ok") + + workflow = WorkflowBuilder(start_executor=CleanupRaisesExecutor(id="cleanup")).build() + + with pytest.raises(RuntimeError): + async for _ in workflow.run(Msg(fail=True), stream=True): + pass + + async for _ in workflow.run(Msg(fail=False), stream=True): + pass + + committed = workflow._runner.state.export_state() # pyright: ignore[reportPrivateUsage] + assert "secret" not in committed diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py index 7b346f12a4..d6828f09a2 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py @@ -401,6 +401,14 @@ def set_state_data(self, data: DeclarativeStateData) -> None: """Set the full state data dict in state.""" self._state.set(DECLARATIVE_STATE_KEY, data) + def commit(self) -> None: + """Commit pending runner State writes. + + Used when an action publishes diagnostic state and then fails the + superstep. The runner discards uncommitted writes on failure (#7859). + """ + self._state.commit() + def get(self, path: str, default: Any = None) -> Any: """Get a value from the state using a dot-notated path. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_http.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_http.py index 0c67b579ab..7d43cd05d0 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_http.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_http.py @@ -206,6 +206,9 @@ async def handle_action( # Non-success path: still publish headers diagnostically, then raise. self._assign_response_headers(state, result) + # Runner discards pending State when the superstep fails (#7859 / #8306). + # Commit first so diagnostic headers remain readable after the error. + state.commit() raise DeclarativeActionError(f"HTTP request to '{url}' failed with status code {result.status_code}.") # ----- Field resolution ----------------------------------------------------