From 4460b05fc916d28446825c92d53cdf15d35b5b2e Mon Sep 17 00:00:00 2001 From: atty57 <99388680+atty57@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:46:58 -0400 Subject: [PATCH 1/2] Python: Reject unsafe functional step replay (#9385) --- .../agent_framework/_workflows/_functional.py | 47 ++++-- .../workflow/test_functional_workflow.py | 147 ++++++++++++++++++ 2 files changed, 180 insertions(+), 14 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 0408e021ca1..ce247628270 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -777,7 +777,8 @@ def _export_step_cache(self) -> dict[str, Any]: if self._non_replayable_step_cache_keys: raise ValueError( "Cannot checkpoint @step calls without a durable replay identity and safely copyable result. " - "Use replay_key for captured state, return replayable values, or run without checkpointing/HITL." + "Use replay_key for opaque arguments or captured state, return replayable values, " + "or run without checkpointing/HITL." ) return dict(self._step_cache) @@ -828,10 +829,10 @@ class StepWrapper(Generic[R]): * **Caching** — results use a stable identity derived from bound, canonical arguments so HITL replay and checkpoint restore skip the matching logical invocation even if concurrent scheduling changes. - Use ``replay_key`` for opaque arguments or concurrent calls whose - identical arguments do not imply interchangeable results. On cache - hit a single ``executor_bypassed`` event is emitted instead of the - normal ``executor_invoked`` / ``executor_completed`` pair. + Use ``replay_key`` for opaque arguments in checkpoint/HITL replay or + concurrent calls whose identical arguments do not imply interchangeable + results. On cache hit a single ``executor_bypassed`` event is emitted + instead of the normal ``executor_invoked`` / ``executor_completed`` pair. * **Event emission** — ``executor_invoked`` / ``executor_completed`` / ``executor_failed`` events are emitted for observability. * **RunContext injection** — if the step function declares a parameter @@ -850,6 +851,8 @@ class StepWrapper(Generic[R]): string-key mappings containing those values. Concurrent calls with the same automatic identity must produce interchangeable results; use distinct ``replay_key`` values otherwise. + Opaque arguments require ``replay_key`` for checkpoint/HITL replay; + sequential execution without replay still supports them. Functions with captured closure state require ``replay_key`` for concurrent or durable replay; the key is the caller's versioning contract for that captured state. @@ -941,16 +944,19 @@ def _get_replay_identity( if not self._wrapper_identity_is_durable: return None + identity = self._get_argument_replay_identity(args, kwargs) + if identity is None: + return None + return ("auto", identity) + + def _get_argument_replay_identity(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str | None: identity_kwargs = kwargs if self._ctx_param_name is not None and self._ctx_param_name in kwargs: identity_kwargs = dict(kwargs) identity_kwargs.pop(self._ctx_param_name) bound = self._identity_signature.bind(*args, **identity_kwargs) bound.apply_defaults() - identity = _hash_step_identity(bound.arguments) - if identity is None: - return None - return ("auto", identity) + return _hash_step_identity(bound.arguments) async def _return_cached_result(self, ctx: RunContext, cache_key: str, cached: Any) -> R: ctx._advance_auto_request_info_index_for_cached_step(cache_key) @@ -1027,6 +1033,12 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: "Regenerate the checkpoint with the current version and provide replay_key for opaque or " "non-interchangeable concurrent calls." ) + if replay_identity is None and self._get_argument_replay_identity(args, kwargs) is None: + raise ValueError( + f"Cannot safely replay legacy order-based cache entry for @step '{self.name}' " + "with opaque arguments. Start a new run with @step(replay_key=...) " + "to identify each logical invocation." + ) return await self._return_cached_result(ctx, legacy_cache_key, cached) if replay_identity is None and ctx._is_concurrent_execution(): @@ -1035,6 +1047,12 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: "Use @step(replay_key=...) to identify each logical invocation." ) + if replay_identity is None and ctx._on_step_completed is not None: + raise ValueError( + f"Cannot checkpoint @step '{self.name}' without a stable replay identity. " + "Use @step(replay_key=...) to identify each logical invocation." + ) + # Inject RunContext if the step function declares it call_args, call_kwargs = self._build_call_args_with_ctx(ctx, args, kwargs) @@ -1069,7 +1087,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: ctx._auto_request_info_index - auto_request_info_index_before, ) copied, cached_result = _try_snapshot_step_cache_value(result) - identity_is_replayable = self._replay_key is not None or self._wrapper_identity_is_durable + identity_is_replayable = replay_identity is not None ctx._set_cached_result(cache_key, cached_result, replayable=copied and identity_is_replayable) _, event_data = _try_snapshot_step_cache_value(result) await ctx.add_event(WorkflowEvent.executor_completed(self.name, event_data)) @@ -1122,10 +1140,11 @@ def step( Concurrent calls with identical automatic identities must produce interchangeable results. Use distinct ``replay_key`` values when opaque arguments, captured state, generated wrappers, or other state distinguish - the logical invocations. A function with closure state also needs - ``replay_key`` for checkpoint/HITL replay. Legacy order-based checkpoints replay sequential - steps, but concurrent legacy cache hits fail explicitly rather than - returning a potentially mismatched result. + the logical invocations. Calls with opaque arguments or closure state + require ``replay_key`` for checkpoint/HITL replay. Legacy order-based checkpoints + replay sequential steps with supported arguments, but opaque or concurrent + legacy cache hits fail explicitly rather than returning a potentially + mismatched result. Args: func: The async function to decorate (when using the bare diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 1d12399a7e4..737d349547e 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -15,6 +15,7 @@ from typing import Any, cast, overload import pytest +from pydantic import BaseModel from agent_framework import ( AgentResponseUpdate, @@ -41,6 +42,17 @@ _factory_step_calls: list[str] = [] + +class _ReplayJob(BaseModel): + label: str + + +@step +async def _opaque_echo(job: _ReplayJob) -> str: + _factory_step_calls.append(job.label) + return job.label.upper() + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -937,6 +949,141 @@ async def wf(x: int) -> int: class TestStepReplayIdentity: + @pytest.mark.parametrize("checkpoint_in_build", [False, True]) + async def test_checkpointed_nested_opaque_arguments_require_replay_key(self, checkpoint_in_build: bool) -> None: + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + @step + async def wrapper(label: str) -> str: + return await _opaque_echo(_ReplayJob(label=label)) + + @workflow + async def wf(_: str) -> list[str]: + return [await wrapper("a"), await wrapper("b")] + + built = wf.build(checkpoint_storage=storage if checkpoint_in_build else None) + with pytest.raises(ValueError, match="stable replay identity.*replay_key"): + await built.run("input", checkpoint_storage=None if checkpoint_in_build else storage) + + assert _factory_step_calls == [] + assert await storage.list_checkpoints(workflow_name="wf") == [] + + async def test_checkpoint_restore_rejects_opaque_legacy_cache_hit(self) -> None: + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + @built_workflow(checkpoint_storage=storage) + async def wf(_: str) -> str: + return await _opaque_echo(_ReplayJob(label="b")) + + checkpoint = WorkflowCheckpoint( + workflow_name="wf", + graph_signature_hash=wf.graph_signature_hash, + state={"_step_cache": {"_opaque_echo::0": "A"}, "_original_message": "input"}, + ) + checkpoint_id = await storage.save(checkpoint) + + with pytest.raises(ValueError, match="opaque arguments.*replay_key"): + await wf.run(checkpoint_id=checkpoint_id) + + assert _factory_step_calls == [] + + async def test_nested_opaque_arguments_without_replay_succeed(self) -> None: + _factory_step_calls.clear() + + @step + async def wrapper(label: str) -> str: + return await _opaque_echo(_ReplayJob(label=label)) + + @built_workflow + async def wf(_: str) -> list[str]: + return [await wrapper("a"), await wrapper("b")] + + result = await wf.run("input") + + assert result.get_outputs() == [["A", "B"]] + assert _factory_step_calls == ["a", "b"] + + async def test_nested_opaque_arguments_reject_response_only_replay(self) -> None: + _factory_step_calls.clear() + + @step + async def wrapper(label: str) -> str: + return await _opaque_echo(_ReplayJob(label=label)) + + @built_workflow + async def wf(_: str, ctx: RunContext) -> str: + result = await wrapper("a") + await ctx.request_info("continue", response_type=str, request_id="continue") + return result + + with pytest.raises(ValueError, match="replay_key"): + await wf.run("input") + + assert _factory_step_calls == ["a"] + + async def test_checkpoint_restore_nested_opaque_arguments_with_replay_keys(self) -> None: + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + crash_at = "b" + + @step(replay_key=lambda job: job.label) + async def tick(job: _ReplayJob) -> str: + calls.append(job.label) + if job.label == crash_at: + raise RuntimeError("killed") + return job.label.upper() + + @step(replay_key=lambda job: job.label) + async def wrapper(job: _ReplayJob) -> str: + return await tick(job) + + @workflow(name="nested_repro") + async def wf(_: str) -> list[str]: + return [await wrapper(_ReplayJob(label="a")), await wrapper(_ReplayJob(label="b"))] + + with pytest.raises(RuntimeError, match="killed"): + await wf.build(checkpoint_storage=storage).run("input") + + calls.clear() + crash_at = "" + latest = await storage.get_latest(workflow_name="nested_repro") + assert latest is not None + resumed = await wf.build(checkpoint_storage=storage).run(checkpoint_id=latest.checkpoint_id) + + assert resumed.get_outputs() == [["A", "B"]] + assert calls == ["b"] + + async def test_nested_opaque_arguments_with_replay_keys_resume_response_only(self) -> None: + calls: list[str] = [] + + @step(replay_key=lambda job: job.label) + async def tick(job: _ReplayJob, ctx: RunContext) -> str: + calls.append(job.label) + if job.label == "b": + await ctx.request_info("continue", response_type=str, request_id="continue") + return job.label.upper() + + @step(replay_key=lambda job: job.label) + async def wrapper(job: _ReplayJob) -> str: + return await tick(job) + + @built_workflow + async def wf(_: str) -> list[str]: + return [await wrapper(_ReplayJob(label="a")), await wrapper(_ReplayJob(label="b"))] + + interrupted = await wf.run("input") + assert interrupted.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + + calls.clear() + resumed = await wf.run(responses={"continue": "yes"}) + + assert resumed.get_outputs() == [["A", "B"]] + assert calls == ["b"] + async def test_checkpoint_restore_preserves_concurrent_step_invocations(self): storage = InMemoryCheckpointStorage() predecessor_started = asyncio.Event() From 9bb14b350fb68d9c4c7432b4889b145505f47920 Mon Sep 17 00:00:00 2001 From: atty57 <99388680+atty57@users.noreply.github.com> Date: Wed, 14 Oct 2026 13:16:26 -0400 Subject: [PATCH 2/2] Python: Reject unverified legacy checkpoint replay --- .../agent_framework/_workflows/_functional.py | 60 +++---- .../workflow/test_functional_workflow.py | 164 ++++++++++++++++-- 2 files changed, 171 insertions(+), 53 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index ce247628270..a89d6e7f1b5 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -517,7 +517,8 @@ def __init__( self._step_cache_auto_request_info_counts: dict[str, int] = {} # Wrapper identities reachable from the current workflow definition. self._expected_step_wrapper_identities: dict[str, set[str]] = {} - # Legacy per-step call counters retained for old checkpoint compatibility. + # Per-step counters detect old order-only cache entries and support + # sequential execution without a durable replay identity. self._step_call_counters: dict[str, int] = {} # Per-identity counters keep repeated automatic invocations distinct. self._step_identity_counters: dict[tuple[str, str, str], int] = {} @@ -856,8 +857,10 @@ class StepWrapper(Generic[R]): Functions with captured closure state require ``replay_key`` for concurrent or durable replay; the key is the caller's versioning contract for that captured state. - Legacy order-based checkpoints remain replayable sequentially, but a - concurrent legacy cache hit raises instead of risking a mismatched result. + Legacy order-based cache entries lack invocation identity and cannot be + verified, even for top-level sequential calls. Attempting to replay one + raises instead of risking a mismatched result. Start a new run to regenerate + checkpoints with stable replay identities. Args: func: The async function to wrap. @@ -868,9 +871,8 @@ class StepWrapper(Generic[R]): Raises: TypeError: If *func* is not an async (coroutine) function. - ValueError: If replay identity is unavailable or invalid for a - concurrent invocation, or a legacy checkpoint cannot be replayed - safely under concurrency. + ValueError: If replay identity is unavailable or invalid for + checkpointing or concurrency, or a legacy cache hit cannot be verified. """ def __init__( @@ -944,19 +946,16 @@ def _get_replay_identity( if not self._wrapper_identity_is_durable: return None - identity = self._get_argument_replay_identity(args, kwargs) - if identity is None: - return None - return ("auto", identity) - - def _get_argument_replay_identity(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str | None: identity_kwargs = kwargs if self._ctx_param_name is not None and self._ctx_param_name in kwargs: identity_kwargs = dict(kwargs) identity_kwargs.pop(self._ctx_param_name) bound = self._identity_signature.bind(*args, **identity_kwargs) bound.apply_defaults() - return _hash_step_identity(bound.arguments) + identity = _hash_step_identity(bound.arguments) + if identity is None: + return None + return ("auto", identity) async def _return_cached_result(self, ctx: RunContext, cache_key: str, cached: Any) -> R: ctx._advance_auto_request_info_index_for_cached_step(cache_key) @@ -1025,21 +1024,14 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: "Start a new run or restore a checkpoint created by the current workflow version." ) - found, cached = ctx._get_cached_result(legacy_cache_key) - if found: - if ctx._is_concurrent_execution(): - raise ValueError( - f"Cannot safely replay legacy order-based cache entry for concurrent @step '{self.name}'. " - "Regenerate the checkpoint with the current version and provide replay_key for opaque or " - "non-interchangeable concurrent calls." - ) - if replay_identity is None and self._get_argument_replay_identity(args, kwargs) is None: - raise ValueError( - f"Cannot safely replay legacy order-based cache entry for @step '{self.name}' " - "with opaque arguments. Start a new run with @step(replay_key=...) " - "to identify each logical invocation." - ) - return await self._return_cached_result(ctx, legacy_cache_key, cached) + # A skipped enclosing step can shift these counters even when the current + # call is top-level. Order alone cannot verify the cached invocation. + if legacy_cache_key in ctx._step_cache: + raise ValueError( + f"Cannot safely replay legacy order-based cache entry for @step '{self.name}': " + "its invocation identity cannot be verified. Start a new run to regenerate the checkpoint " + "and use @step(replay_key=...) for opaque arguments or captured state." + ) if replay_identity is None and ctx._is_concurrent_execution(): raise ValueError( @@ -1141,10 +1133,9 @@ def step( interchangeable results. Use distinct ``replay_key`` values when opaque arguments, captured state, generated wrappers, or other state distinguish the logical invocations. Calls with opaque arguments or closure state - require ``replay_key`` for checkpoint/HITL replay. Legacy order-based checkpoints - replay sequential steps with supported arguments, but opaque or concurrent - legacy cache hits fail explicitly rather than returning a potentially - mismatched result. + require ``replay_key`` for checkpoint/HITL replay. Legacy order-based cache + entries cannot be verified and are rejected, including top-level sequential + calls. Start a new run to regenerate checkpoints with stable replay identities. Args: func: The async function to decorate (when using the bare @@ -1162,9 +1153,8 @@ def step( Raises: TypeError: If the decorated function is not async. - ValueError: If replay identity is unavailable or invalid for a - concurrent invocation, or a legacy checkpoint cannot be replayed - safely under concurrency. + ValueError: If replay identity is unavailable or invalid for + checkpointing or concurrency, or a legacy cache hit cannot be verified. Examples: diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 737d349547e..276ee4d62f4 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -53,6 +53,19 @@ async def _opaque_echo(job: _ReplayJob) -> str: return job.label.upper() +# Module-level (durable, fingerprintable-argument) nested steps used to reproduce the +# nested legacy-replay corruption: a cached _nested_wrapper hit skips _nested_inner. +@step +async def _nested_inner(label: str) -> str: + _factory_step_calls.append(label) + return label.upper() + + +@step +async def _nested_wrapper(label: str) -> str: + return await _nested_inner(label) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -986,11 +999,131 @@ async def wf(_: str) -> str: ) checkpoint_id = await storage.save(checkpoint) - with pytest.raises(ValueError, match="opaque arguments.*replay_key"): + with pytest.raises(ValueError, match="invocation identity cannot be verified"): + await wf.run(checkpoint_id=checkpoint_id) + + assert _factory_step_calls == [] + + async def test_checkpoint_restore_rejects_nested_legacy_cache_hit_auto_identity(self) -> None: + """A cached enclosing step skips its inner call, so an order-based legacy inner + entry must not silently satisfy a later inner invocation (would return ['A', 'A']).""" + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + @built_workflow(name="nested_legacy_auto", checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + return [await _nested_wrapper("a"), await _nested_wrapper("b")] + + checkpoint = WorkflowCheckpoint( + workflow_name="nested_legacy_auto", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": {"_nested_wrapper::0": "A", "_nested_inner::0": "A"}, + "_step_cache_auto_request_info_counts": {"_nested_wrapper::0": 0, "_nested_inner::0": 0}, + "_original_message": "input", + }, + ) + checkpoint_id = await storage.save(checkpoint) + + with pytest.raises(ValueError, match="legacy order-based cache entry.*invocation identity"): await wf.run(checkpoint_id=checkpoint_id) assert _factory_step_calls == [] + async def test_checkpoint_restore_rejects_nested_legacy_cache_hit_explicit_key(self) -> None: + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + + @step(replay_key=lambda label: label) + async def inner(label: str) -> str: + calls.append(label) + return label.upper() + + @step(replay_key=lambda label: label) + async def wrapper(label: str) -> str: + return await inner(label) + + @built_workflow(name="nested_legacy_explicit", checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + return [await wrapper("a"), await wrapper("b")] + + checkpoint = WorkflowCheckpoint( + workflow_name="nested_legacy_explicit", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": {"wrapper::0": "A", "inner::0": "A"}, + "_step_cache_auto_request_info_counts": {"wrapper::0": 0, "inner::0": 0}, + "_original_message": "input", + }, + ) + checkpoint_id = await storage.save(checkpoint) + + with pytest.raises(ValueError, match="legacy order-based cache entry.*invocation identity"): + await wf.run(checkpoint_id=checkpoint_id) + + assert calls == [] + + @pytest.mark.parametrize("explicit_keys", [False, True]) + @pytest.mark.parametrize("direct_inner", [False, True]) + async def test_versioned_outer_rejects_unverified_legacy_inner( + self, explicit_keys: bool, direct_inner: bool + ) -> None: + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + _factory_step_calls.clear() + + @step(replay_key=lambda label: label) + async def explicit_inner(label: str) -> str: + calls.append(label) + return label.upper() + + @step(replay_key=lambda label: label) + async def explicit_wrapper(label: str) -> str: + return await explicit_inner(label) + + inner = explicit_inner if explicit_keys else _nested_inner + wrapper = explicit_wrapper if explicit_keys else _nested_wrapper + + @built_workflow(checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + first = await wrapper("a") + second = await inner("b") if direct_inner else await wrapper("b") + return [first, second] + + replay_identity = wrapper._get_replay_identity(("a",), {}) # pyright: ignore[reportPrivateUsage] + assert replay_identity is not None + identity_kind, identity = replay_identity + key_ctx = _RunContext("wf") + if identity_kind == "explicit": + outer_key = key_ctx._get_explicit_step_cache_key( # pyright: ignore[reportPrivateUsage] + wrapper.name, wrapper._wrapper_source_identity, identity + ) + else: + outer_key = key_ctx._get_automatic_step_cache_key( # pyright: ignore[reportPrivateUsage] + wrapper.name, wrapper._wrapper_identity, identity + ) + checkpoint = WorkflowCheckpoint( + workflow_name="wf", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": {outer_key: "A", f"{inner.name}::0": "A"}, + "_original_message": "input", + }, + ) + checkpoint_id = await storage.save(checkpoint) + + with pytest.raises(ValueError, match="legacy order-based cache entry"): + await wf.run(checkpoint_id=checkpoint_id) + + assert calls == [] + assert _factory_step_calls == [] + async def test_nested_opaque_arguments_without_replay_succeed(self) -> None: _factory_step_calls.clear() @@ -1837,7 +1970,7 @@ async def wf(_: str) -> threading.Lock: with pytest.raises(ValueError, match="Cannot checkpoint"): await wf.run("input") - async def test_legacy_checkpoint_replays_sequential_step(self): + async def test_legacy_checkpoint_rejects_unverifiable_sequential_step(self): from agent_framework import WorkflowCheckpoint storage = InMemoryCheckpointStorage() @@ -1864,9 +1997,9 @@ async def wf(value: int) -> int: ) checkpoint_id = await storage.save(checkpoint) - replayed = await wf.run(checkpoint_id=checkpoint_id) + with pytest.raises(ValueError, match="invocation identity cannot be verified"): + await wf.run(checkpoint_id=checkpoint_id) - assert replayed.get_outputs() == [6] assert call_count == 0 async def test_legacy_checkpoint_rejects_concurrent_cache_hit(self): @@ -1935,9 +2068,7 @@ async def wf(_: str) -> str: with pytest.raises(ValueError, match="legacy order-based cache entry"): await wf.run(checkpoint_id=checkpoint_id) - async def test_legacy_sequential_replay_ignores_unrelated_caller_task(self): - from agent_framework import WorkflowCheckpoint - + async def test_versioned_sequential_replay_ignores_unrelated_caller_task(self): storage = InMemoryCheckpointStorage() workflow_waiting = asyncio.Event() release_workflow = asyncio.Event() @@ -1947,22 +2078,19 @@ async def test_legacy_sequential_replay_ignores_unrelated_caller_task(self): async def compute(value: int) -> int: return value + 1 - @built_workflow(name="legacy_sequential", checkpoint_storage=storage) + @built_workflow(name="versioned_sequential", checkpoint_storage=storage) async def wf(value: int) -> int: workflow_waiting.set() await release_workflow.wait() return await compute(value) - checkpoint = WorkflowCheckpoint( - workflow_name="legacy_sequential", - graph_signature_hash=wf.graph_signature_hash, - state={ - "_step_cache": {"compute::0": 6}, - "_step_cache_auto_request_info_counts": {"compute::0": 0}, - "_original_message": 5, - }, - ) - checkpoint_id = await storage.save(checkpoint) + release_workflow.set() + await wf.run(5) + checkpoint = await storage.get_latest(workflow_name="versioned_sequential") + assert checkpoint is not None + checkpoint_id = checkpoint.checkpoint_id + workflow_waiting.clear() + release_workflow.clear() loop = asyncio.get_running_loop() original_factory = loop.get_task_factory()