From 34b74519537df002de76fa0b6e6f4409ad8087bf Mon Sep 17 00:00:00 2001 From: CoralGarden52 <2193436736@qq.com> Date: Fri, 11 Sep 2026 14:14:14 +0000 Subject: [PATCH 1/2] Python: preserve functional workflow state on response-only resume --- .../agent_framework/_workflows/_functional.py | 5 ++++ .../workflow/test_functional_workflow.py | 25 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index e83fa9e0e5c..900bd3ebd7f 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/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. @@ -1015,6 +1016,7 @@ async def _run_core( 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) + ctx._state = dict(self._last_state) # Store message for future replays if message is not None: @@ -1064,6 +1066,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._last_state = dict(ctx._state) # Yield collected events. # NOTE: Events are buffered during _execute() and yielded after @@ -1091,6 +1094,7 @@ async def _on_step_completed() -> None: 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() yield _framework_event(WorkflowEvent.status, WorkflowRunState.IDLE) @@ -1100,6 +1104,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._last_state = dict(ctx._state) self._last_pending_request_ids = set(ctx._pending_requests) # HITL interruption — yield events collected so far diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 5ed7ae0e8a3..f2d6593efc6 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -1470,6 +1470,31 @@ def emit(self, record: logging.LogRecord) -> None: class TestHITLInStepWithCaching: """Regression tests: request_info inside @step combined with caching and bypass.""" + 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.""" From 1ee0f2bfa33ef433a92571d938ce82413bce173b Mon Sep 17 00:00:00 2001 From: CoralGarden52 <2193436736@qq.com> Date: Mon, 14 Sep 2026 07:00:07 +0000 Subject: [PATCH 2/2] refactor: centralize functional workflow replay state --- .../agent_framework/_workflows/_functional.py | 46 ++++++++++++------- .../workflow/test_functional_workflow.py | 37 +++++++++++++++ 2 files changed, 66 insertions(+), 17 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 900bd3ebd7f..eba3e8aaa34 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -753,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. @@ -1012,11 +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) - ctx._state = dict(self._last_state) + message = replay_message # Store message for future replays if message is not None: @@ -1064,9 +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._last_state = dict(ctx._state) + self._capture_replay_state(ctx, message) # Yield collected events. # NOTE: Events are buffered during _execute() and yielded after @@ -1087,25 +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_state = {} - 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_state = dict(ctx._state) - 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(): diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index f2d6593efc6..207ae0ecc96 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -1470,6 +1470,43 @@ 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