diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 0408e021ca1..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] = {} @@ -777,7 +778,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 +830,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,11 +852,15 @@ 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. - 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. @@ -865,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__( @@ -1019,15 +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." - ) - 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( @@ -1035,6 +1039,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 +1079,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 +1132,10 @@ 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 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 @@ -1143,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 1d12399a7e4..276ee4d62f4 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,30 @@ _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() + + +# 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 # --------------------------------------------------------------------------- @@ -937,6 +962,261 @@ 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="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() + + @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() @@ -1690,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() @@ -1717,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): @@ -1788,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() @@ -1800,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()