Python: preserve functional workflow state on response-only HITL resume

python/packages/core/agent_framework/_workflows/_functional.py

@@ -739,6 +739,7 @@ def __init__(
         self._last_message: Any = None
         self._last_step_cache: dict[tuple[str, int], Any] = {}
         self._last_step_cache_auto_request_info_counts: dict[tuple[str, int], int] = {}
+        self._last_state: dict[str, Any] = {}
         self._last_pending_request_ids: set[str] = set()
 
         # Signature arity is validated once at decoration time.
@@ -752,6 +753,30 @@ def __init__(
 
         functools.update_wrapper(self, func)  # type: ignore[arg-type]
 
+    def _capture_replay_state(self, ctx: RunContext, message: Any | None = None) -> None:
+        """Capture the state needed to continue a response-only HITL replay."""
+        if message is not None:
+            self._last_message = message
+        self._last_step_cache = dict(ctx._step_cache)
+        self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
+        self._last_state = dict(ctx._state)
+        self._last_pending_request_ids = set(ctx._pending_requests)
+
+    def _restore_replay_state(self, ctx: RunContext) -> Any:
+        """Restore cached execution state and return the message used by the replay."""
+        ctx._step_cache = dict(self._last_step_cache)
+        ctx._step_cache_auto_request_info_counts = dict(self._last_step_cache_auto_request_info_counts)
+        ctx._state = dict(self._last_state)
+        return self._last_message
+
+    def _clear_replay_state(self) -> None:
+        """Clear all state retained for a response-only replay."""
+        self._last_message = None
+        self._last_step_cache = {}
+        self._last_step_cache_auto_request_info_counts = {}
+        self._last_state = {}
+        self._last_pending_request_ids = set()
+
     @staticmethod
     def _classify_signature(func: Callable[..., Any]) -> list[str]:
         """Return the names of non-ctx parameters, validating arity.
@@ -1011,10 +1036,9 @@ async def _run_core(
 
         # For response-only replay (no checkpoint), restore cached state
         if checkpoint_id is None and responses:
+            replay_message = self._restore_replay_state(ctx)
             if message is None:
-                message = self._last_message
-            ctx._step_cache = dict(self._last_step_cache)
-            ctx._step_cache_auto_request_info_counts = dict(self._last_step_cache_auto_request_info_counts)
+                message = replay_message
 
         # Store message for future replays
         if message is not None:
@@ -1062,8 +1086,7 @@ async def _on_step_completed() -> None:
                     )
 
                 # Persist step cache for response-only replay
-                self._last_step_cache = dict(ctx._step_cache)
-                self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
+                self._capture_replay_state(ctx, message)
 
             # Yield collected events.
             # NOTE: Events are buffered during _execute() and yielded after
@@ -1084,23 +1107,17 @@ async def _on_step_completed() -> None:
 
             # Final status
             if saw_request:
-                self._last_pending_request_ids = set(ctx._pending_requests)
                 yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
             else:
                 # Clean completion — drop cross-run replay state.
-                self._last_message = None
-                self._last_step_cache = {}
-                self._last_step_cache_auto_request_info_counts = {}
-                self._last_pending_request_ids = set()
+                self._clear_replay_state()
                 yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE)
 
             span.add_event(OtelAttr.WORKFLOW_COMPLETED)
 
         except WorkflowInterrupted:
             # Persist step cache for response-only replay
-            self._last_step_cache = dict(ctx._step_cache)
-            self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts)
-            self._last_pending_request_ids = set(ctx._pending_requests)
+            self._capture_replay_state(ctx, message)
 
             # HITL interruption — yield events collected so far
             for event in ctx._get_events():

python/packages/core/tests/workflow/test_functional_workflow.py

@@ -1470,6 +1470,68 @@ def emit(self, record: logging.LogRecord) -> None:
 class TestHITLInStepWithCaching:
     """Regression tests: request_info inside @step combined with caching and bypass."""
 
+    def test_replay_state_helpers_restore_and_clear_the_full_bundle(self):
+        """Replay helpers keep message, caches, state, and pending IDs together."""
+
+        @built_workflow
+        async def wf(data: str) -> str:
+            return data
+
+        source_ctx = _RunContext("wf")
+        source_ctx._step_cache = {("seed_state", 0): "seeded"}
+        source_ctx._step_cache_auto_request_info_counts = {("seed_state", 0): 1}
+        source_ctx._state = {"marker": "ok"}
+        source_ctx._pending_requests = {
+            "r1": WorkflowEvent.request_info(
+                request_id="r1",
+                source_executor_id="wf",
+                request_data="question",
+                response_type=str,
+            )
+        }
+
+        wf._capture_replay_state(source_ctx, "input")
+
+        restored_ctx = _RunContext("wf")
+        assert wf._restore_replay_state(restored_ctx) == "input"
+        assert restored_ctx._step_cache == source_ctx._step_cache
+        assert restored_ctx._step_cache_auto_request_info_counts == source_ctx._step_cache_auto_request_info_counts
+        assert restored_ctx._state == source_ctx._state
+        assert wf._last_pending_request_ids == {"r1"}
+
+        wf._clear_replay_state()
+
+        assert wf._last_message is None
+        assert wf._last_step_cache == {}
+        assert wf._last_step_cache_auto_request_info_counts == {}
+        assert wf._last_state == {}
+        assert wf._last_pending_request_ids == set()
+
+    async def test_response_only_resume_restores_state_from_cached_step(self):
+        """Response-only HITL resumes must preserve state written before a cached step."""
+        seed_calls = 0
+
+        @step
+        async def seed_state(ctx: RunContext) -> str:
+            nonlocal seed_calls
+            seed_calls += 1
+            ctx.set_state("marker", "ok")
+            return "seeded"
+
+        @built_workflow
+        async def wf(data: str, ctx: RunContext) -> str:
+            value = await seed_state()
+            answer = await ctx.request_info("question", response_type=str, request_id="r1")
+            return f"{ctx.get_state('marker', 'MISSING')}:{value}:{answer}"
+
+        result1 = await wf.run("input")
+        assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+
+        result2 = await wf.run(responses={"r1": "ok"})
+
+        assert seed_calls == 1
+        assert result2.get_outputs() == ["ok:seeded:ok"]
+
     async def test_preceding_step_bypassed_on_hitl_resume(self):
         """When a step after a completed step calls request_info and interrupts,
         resuming should bypass the first step (cached) and re-execute the HITL step."""