From 02998ba78306c3ec12e23f52e37029c99956d206 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 10:05:05 +0200 Subject: [PATCH 1/7] Python: fix functional workflow replay identity --- .../agent_framework/_workflows/_functional.py | 305 +++++++++++++--- .../workflow/test_functional_workflow.py | 338 +++++++++++++++++- .../functional/parallel_pipeline.py | 4 +- 3 files changed, 602 insertions(+), 45 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index ae94a98692f..6d4e29a0c0b 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -39,12 +39,15 @@ # pyright: reportPrivateUsage=false # Classes in this module (RunContext, StepWrapper, FunctionalWorkflow) form a # cohesive unit and intentionally access each other's underscore-prefixed members. +import asyncio import functools import hashlib import inspect +import json import logging +import math import typing -from collections.abc import AsyncIterable, Awaitable, Callable, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence from contextvars import ContextVar from copy import deepcopy from typing import Any, Generic, Literal, TypeVar, overload @@ -72,11 +75,111 @@ R = TypeVar("R") +_STEP_CACHE_KEY_V2_PREFIX = "v2::" +_UNSUPPORTED_STEP_IDENTITY = object() + # ContextVar holding the active RunContext during workflow execution. # ContextVar is per-asyncio-Task, so concurrent workflows each get their own context. _active_run_ctx: ContextVar[RunContext | None] = ContextVar("_active_run_ctx", default=None) +def _canonicalize_step_identity_value(value: Any) -> Any: + """Return a type-preserving JSON value, or a sentinel for unsupported input.""" + if value is None: + return ["none"] + if type(value) is bool: + return ["bool", value] + if type(value) is int: + return ["int", value] + if type(value) is float: + if not math.isfinite(value): + return _UNSUPPORTED_STEP_IDENTITY + return ["float", value.hex()] + if type(value) is str: + return ["str", value] + if type(value) is list: + list_items: list[Any] = [] + for item in typing.cast(list[Any], value): + canonical = _canonicalize_step_identity_value(item) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + list_items.append(canonical) + return ["list", list_items] + if type(value) is tuple: + tuple_items: list[Any] = [] + for item in typing.cast(tuple[Any, ...], value): + canonical = _canonicalize_step_identity_value(item) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + tuple_items.append(canonical) + return ["tuple", tuple_items] + if isinstance(value, Mapping): + mapping = typing.cast(Mapping[Any, Any], value) + mapping_items: list[Any] = [] + for key, item in mapping.items(): + if type(key) is not str: + return _UNSUPPORTED_STEP_IDENTITY + canonical = _canonicalize_step_identity_value(item) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + mapping_items.append([key, canonical]) + mapping_type = f"{type(mapping).__module__}.{type(mapping).__qualname__}" + return ["mapping", mapping_type, mapping_items] + return _UNSUPPORTED_STEP_IDENTITY + + +def _hash_step_identity(value: Any) -> str | None: + canonical = _canonicalize_step_identity_value(value) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return None + encoded = json.dumps(canonical, ensure_ascii=False, separators=(",", ":"), allow_nan=False) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _encode_step_cache_key( + kind: Literal["auto", "explicit"], + step_name: str, + identity: str, + occurrence: int | None = None, +) -> str: + payload: list[str | int] = [kind, step_name, identity] + if kind == "auto": + if occurrence is None: + raise ValueError("Automatic step cache keys require an occurrence index.") + payload.append(occurrence) + return f"{_STEP_CACHE_KEY_V2_PREFIX}{json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}" + + +def _validate_step_cache_key(key: Any) -> str: + if not isinstance(key, str): + raise TypeError("Step cache keys must be strings.") + + if key.startswith(_STEP_CACHE_KEY_V2_PREFIX): + try: + raw_payload: object = json.loads(key.removeprefix(_STEP_CACHE_KEY_V2_PREFIX)) + except json.JSONDecodeError: + raw_payload = None + if isinstance(raw_payload, list) and raw_payload: + payload = typing.cast(list[Any], raw_payload) + kind = payload[0] + if ( + kind == "auto" + and len(payload) == 4 + and all(isinstance(item, str) for item in payload[1:3]) + and isinstance(payload[3], int) + and payload[3] >= 0 + ): + return key + if kind == "explicit" and len(payload) == 3 and all(isinstance(item, str) for item in payload[1:]): + return key + raise ValueError("Invalid versioned step cache key.") + + name, idx_str = key.rsplit("::", 1) + if not name or int(idx_str) < 0: + raise ValueError("Invalid legacy step cache key.") + return key + + @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) def get_run_context() -> RunContext | None: """Return the active :class:`RunContext`, or ``None`` if not inside a ``@workflow``. @@ -175,12 +278,19 @@ def __init__( # Event accumulator self._events: list[WorkflowEvent[Any]] = [] - # Step result cache: (step_name, call_index) -> result - self._step_cache: dict[tuple[str, int], Any] = {} + # Step result cache. Keys are already checkpoint-safe strings so result + # and metadata maps cannot drift during serialization. + self._step_cache: dict[str, Any] = {} # Cached step metadata used to keep auto-generated request_info IDs in sync on bypass. - self._step_cache_auto_request_info_counts: dict[tuple[str, int], int] = {} - # Per-step call counters for deterministic cache keys + self._step_cache_auto_request_info_counts: dict[str, int] = {} + # Legacy per-step call counters retained for old checkpoint compatibility. self._step_call_counters: dict[str, int] = {} + # Per-identity counters keep repeated automatic invocations distinct. + self._step_identity_counters: dict[tuple[str, str], int] = {} + # Explicit replay keys identify exactly one logical invocation per run. + self._used_explicit_step_cache_keys: set[str] = set() + # The task executing the workflow body; gather/task-group children differ from this task. + self._root_task: asyncio.Task[Any] | None = None # Deterministic call counter for auto-generated request_info IDs self._auto_request_info_index: int = 0 @@ -328,23 +438,43 @@ def is_streaming(self) -> bool: def _get_events(self) -> list[WorkflowEvent[Any]]: return list(self._events) - def _get_step_cache_key(self, step_name: str) -> tuple[str, int]: + def _get_legacy_step_cache_key(self, step_name: str) -> str: idx = self._step_call_counters.get(step_name, 0) self._step_call_counters[step_name] = idx + 1 - return (step_name, idx) + return f"{step_name}::{idx}" + + def _get_automatic_step_cache_key(self, step_name: str, identity: str) -> str: + counter_key = (step_name, identity) + occurrence = self._step_identity_counters.get(counter_key, 0) + self._step_identity_counters[counter_key] = occurrence + 1 + return _encode_step_cache_key("auto", step_name, identity, occurrence) + + def _get_explicit_step_cache_key(self, step_name: str, identity: str) -> str: + key = _encode_step_cache_key("explicit", step_name, identity) + if key in self._used_explicit_step_cache_keys: + raise ValueError( + f"@step '{step_name}' produced duplicate replay_key values in one workflow run. " + "Each explicit replay key must identify exactly one logical invocation." + ) + self._used_explicit_step_cache_keys.add(key) + return key + + def _is_concurrent_task(self) -> bool: + current_task = asyncio.current_task() + return self._root_task is not None and current_task is not None and current_task is not self._root_task - def _get_cached_result(self, key: tuple[str, int]) -> tuple[bool, Any]: + def _get_cached_result(self, key: str) -> tuple[bool, Any]: if key in self._step_cache: return True, self._step_cache[key] return False, None - def _set_cached_result(self, key: tuple[str, int], value: Any) -> None: + def _set_cached_result(self, key: str, value: Any) -> None: self._step_cache[key] = value - def _set_cached_step_auto_request_info_count(self, key: tuple[str, int], count: int) -> None: + def _set_cached_step_auto_request_info_count(self, key: str, count: int) -> None: self._step_cache_auto_request_info_counts[key] = count - def _advance_auto_request_info_index_for_cached_step(self, key: tuple[str, int]) -> None: + def _advance_auto_request_info_index_for_cached_step(self, key: str) -> None: self._auto_request_info_index += self._step_cache_auto_request_info_counts.get(key, 0) def _set_responses(self, responses: dict[str, Any]) -> None: @@ -376,24 +506,20 @@ def _get_response(self, request_id: str) -> tuple[bool, Any]: return False, None def _export_step_cache(self) -> dict[str, Any]: - """Serialize the step cache for checkpointing. - - Converts tuple keys to strings for JSON compatibility. - """ - return {f"{name}::{idx}": val for (name, idx), val in self._step_cache.items()} + """Serialize the step cache for checkpointing.""" + return dict(self._step_cache) def _export_step_cache_auto_request_info_counts(self) -> dict[str, int]: """Serialize per-step auto request_info counts for checkpointing.""" - return {f"{name}::{idx}": count for (name, idx), count in self._step_cache_auto_request_info_counts.items()} + return dict(self._step_cache_auto_request_info_counts) def _import_step_cache(self, data: dict[str, Any]) -> None: """Restore step cache from checkpoint data.""" self._step_cache = {} for k, v in data.items(): try: - name, idx_str = k.rsplit("::", 1) - self._step_cache[name, int(idx_str)] = v - except (ValueError, TypeError) as exc: + self._step_cache[_validate_step_cache_key(k)] = v + except (ValueError, TypeError, json.JSONDecodeError) as exc: raise ValueError( f"Corrupted step cache entry in checkpoint: key={k!r}. " f"The checkpoint may be from an incompatible version or corrupted. " @@ -405,9 +531,8 @@ def _import_step_cache_auto_request_info_counts(self, data: dict[str, Any]) -> N self._step_cache_auto_request_info_counts = {} for k, v in data.items(): try: - name, idx_str = k.rsplit("::", 1) - self._step_cache_auto_request_info_counts[name, int(idx_str)] = int(v) - except (ValueError, TypeError) as exc: + self._step_cache_auto_request_info_counts[_validate_step_cache_key(k)] = int(v) + except (ValueError, TypeError, json.JSONDecodeError) as exc: raise ValueError( f"Corrupted step cache request_info metadata in checkpoint: key={k!r}, value={v!r}. " f"The checkpoint may be from an incompatible version or corrupted. " @@ -427,10 +552,13 @@ class StepWrapper(Generic[R]): When called inside a running ``@workflow`` function, the wrapper intercepts execution to provide: - * **Caching** — results are cached by ``(step_name, call_index)`` so - that HITL replay and checkpoint restore skip already-completed work. - On cache hit a single ``executor_bypassed`` event is emitted instead - of the normal ``executor_invoked`` / ``executor_completed`` pair. + * **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. * **Event emission** — ``executor_invoked`` / ``executor_completed`` / ``executor_failed`` events are emitted for observability. * **RunContext injection** — if the step function declares a parameter @@ -444,15 +572,34 @@ class StepWrapper(Generic[R]): the original function, making decorated functions fully testable in isolation. + Automatic replay identity supports ``None``, booleans, integers, finite + floats, strings, lists, tuples, and string-key mappings containing those + values. Concurrent calls with the same automatic identity must produce + interchangeable results; use distinct ``replay_key`` values otherwise. + Legacy order-based checkpoints remain replayable sequentially, but a + concurrent legacy cache hit raises instead of risking a mismatched result. + Args: func: The async function to wrap. name: Optional display name. Defaults to ``func.__name__``. + replay_key: Optional callback returning a stable, unique, non-empty + string for each logical invocation. The callback receives the + original user arguments and is not called outside a workflow. 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. """ - def __init__(self, func: Callable[..., Awaitable[R]], *, name: str | None = None) -> None: + def __init__( + self, + func: Callable[..., Awaitable[R]], + *, + name: str | None = None, + replay_key: Callable[..., str] | None = None, + ) -> None: if not inspect.iscoroutinefunction(func): raise TypeError( f"@step can only decorate async functions, but '{func.__name__}' is not a coroutine function." @@ -460,6 +607,7 @@ def __init__(self, func: Callable[..., Awaitable[R]], *, name: str | None = None self._func = func self.name: str = name or func.__name__ self._signature = inspect.signature(func) + self._replay_key = replay_key functools.update_wrapper(self, func) # Detect RunContext parameter for auto-injection inside workflows @@ -480,6 +628,34 @@ def __init__(self, func: Callable[..., Awaitable[R]], *, name: str | None = None self._ctx_param_name = param_name break + identity_parameters = [ + param for param in self._signature.parameters.values() if param.name != self._ctx_param_name + ] + self._identity_signature = self._signature.replace(parameters=identity_parameters) + + def _get_replay_identity( + self, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> tuple[Literal["auto", "explicit"], str] | None: + if self._replay_key is not None: + explicit_key = self._replay_key(*args, **kwargs) + if not isinstance(explicit_key, str) or not explicit_key: + raise ValueError(f"@step '{self.name}' replay_key must return a non-empty string.") + return ("explicit", hashlib.sha256(explicit_key.encode("utf-8")).hexdigest()) + + bound = self._identity_signature.bind(*args, **kwargs) + bound.apply_defaults() + 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) + await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, cached)) + return cached + def _build_call_args_with_ctx( self, ctx: RunContext, @@ -521,14 +697,35 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: # Outside a workflow — pass through directly return await self._func(*args, **kwargs) - cache_key = ctx._get_step_cache_key(self.name) - found, cached = ctx._get_cached_result(cache_key) + legacy_cache_key = ctx._get_legacy_step_cache_key(self.name) + cache_key = legacy_cache_key + replay_identity = self._get_replay_identity(args, kwargs) + if replay_identity is not None: + identity_kind, identity = replay_identity + if identity_kind == "explicit": + cache_key = ctx._get_explicit_step_cache_key(self.name, identity) + else: + cache_key = ctx._get_automatic_step_cache_key(self.name, identity) + + found, cached = ctx._get_cached_result(cache_key) + if found: + return await self._return_cached_result(ctx, cache_key, cached) + + found, cached = ctx._get_cached_result(legacy_cache_key) if found: - ctx._advance_auto_request_info_index_for_cached_step(cache_key) - # Dedicated bypass event so consumers can tell cache-hit replays - # apart from fresh executions. - await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, cached)) - return cached + if ctx._is_concurrent_task(): + 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) + + if replay_identity is None and ctx._is_concurrent_task(): + raise ValueError( + f"Cannot derive a stable replay identity for concurrent @step '{self.name}'. " + "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) @@ -574,7 +771,11 @@ def step(func: Callable[..., Awaitable[R]]) -> StepWrapper[R]: ... @overload -def step(*, name: str | None = None) -> Callable[[Callable[..., Awaitable[R]]], StepWrapper[R]]: ... +def step( + *, + name: str | None = None, + replay_key: Callable[..., str] | None = None, +) -> Callable[[Callable[..., Awaitable[R]]], StepWrapper[R]]: ... @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) @@ -582,6 +783,7 @@ def step( func: Callable[..., Awaitable[Any]] | None = None, *, name: str | None = None, + replay_key: Callable[..., str] | None = None, ) -> StepWrapper[Any] | Callable[[Callable[..., Awaitable[Any]]], StepWrapper[Any]]: """Decorator that marks an async function as a tracked workflow step. @@ -599,11 +801,23 @@ def step( inside ``@workflow`` without it; use ``@step`` only when you need caching, checkpointing, or observability for a particular call. + Automatic replay identity is derived from bound canonical arguments. + Concurrent calls with identical automatic identities must produce + interchangeable results. Use distinct ``replay_key`` values when opaque + arguments or other state distinguish the logical invocations. Legacy + order-based checkpoints replay sequential steps, but concurrent legacy + cache hits fail explicitly rather than returning a potentially mismatched + result. + Args: func: The async function to decorate (when using the bare ``@step`` form). name: Optional display name for the step. Defaults to the function's ``__name__``. + replay_key: Optional callback that returns a stable, unique, + non-empty string for each logical invocation. Use this when step + arguments are opaque or when concurrent calls with identical + arguments can produce different results. Returns: A :class:`StepWrapper` (bare form) or a decorator that produces @@ -611,6 +825,9 @@ 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. Examples: @@ -626,16 +843,21 @@ async def transform_data(raw: dict) -> str: return json.dumps(raw) + @step(replay_key=lambda connection, record_id: record_id) + async def load_record(connection: object, record_id: str) -> dict: + return await connection.load(record_id) + + # Step with HITL — RunContext is auto-injected inside a workflow: @step async def review(doc: str, ctx: RunContext) -> str: return await ctx.request_info({"draft": doc}, response_type=str) """ if func is not None: - return StepWrapper(func, name=name) + return StepWrapper(func, name=name, replay_key=replay_key) def _decorator(fn: Callable[..., Awaitable[Any]]) -> StepWrapper[Any]: - return StepWrapper(fn, name=name) + return StepWrapper(fn, name=name, replay_key=replay_key) return _decorator @@ -741,8 +963,8 @@ def __init__( # Replay state: cleared on clean completion so later responses-only # calls can't silently replay with stale data from a prior run. 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_step_cache: dict[str, Any] = {} + self._last_step_cache_auto_request_info_counts: dict[str, int] = {} self._last_state: dict[str, Any] = {} self._last_pending_request_ids: set[str] = set() @@ -1008,6 +1230,7 @@ async def _run_core( # Build context ctx = RunContext(self.name, streaming=streaming, run_kwargs=kwargs if kwargs else None) + ctx._root_task = asyncio.current_task() # Restore from checkpoint if requested prev_checkpoint_id: str | None = None @@ -1253,8 +1476,6 @@ def _compute_signature_hash(self) -> str: "co_code": co_code_hex, "co_names": list(co_names), } - import json - canonical = json.dumps(sig_data, sort_keys=True, separators=(",", ":")) return hashlib.sha256(canonical.encode("utf-8")).hexdigest() diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 43c0ef84e37..6214c12dc37 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -929,6 +929,311 @@ async def wf(x: int) -> int: assert len(checkpoints) == 3 # 2 from first run + 1 final from restore +# --------------------------------------------------------------------------- +# Step replay identity +# --------------------------------------------------------------------------- + + +class TestStepReplayIdentity: + async def test_checkpoint_restore_preserves_concurrent_step_invocations(self): + storage = InMemoryCheckpointStorage() + predecessor_started = asyncio.Event() + release_predecessor = asyncio.Event() + branch_a_shared_completed = asyncio.Event() + calls: list[str] = [] + + @step + async def predecessor(value: str) -> str: + predecessor_started.set() + await release_predecessor.wait() + return value + + @step + async def shared_step(value: str) -> str: + calls.append(value) + if value == "B": + release_predecessor.set() + await branch_a_shared_completed.wait() + else: + branch_a_shared_completed.set() + return f"result:{value}" + + async def branch_a() -> str: + await predecessor("A") + return await shared_step("A") + + async def branch_b() -> str: + await predecessor_started.wait() + return await shared_step("B") + + @built_workflow(checkpoint_storage=storage) + async def parallel_workflow(_: str) -> list[str]: + return list(await asyncio.gather(branch_a(), branch_b())) + + initial = await parallel_workflow.run("input") + checkpoint = await storage.get_latest(workflow_name="parallel_workflow") + assert checkpoint is not None + + replayed = await parallel_workflow.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == [["result:A", "result:B"]] + assert replayed.get_outputs() == initial.get_outputs() + assert calls == ["B", "A"] + + async def test_intermediate_checkpoint_does_not_cross_branches(self): + storage = InMemoryCheckpointStorage() + release_branch_a = asyncio.Event() + branch_b_returned = asyncio.Event() + calls: list[str] = [] + + @step + async def shared_step(value: str) -> str: + calls.append(value) + return f"result:{value}" + + async def branch_a() -> str: + await release_branch_a.wait() + return await shared_step("A") + + async def branch_b() -> str: + result = await shared_step("B") + branch_b_returned.set() + return result + + @built_workflow(checkpoint_storage=storage) + async def parallel_workflow(_: str) -> list[str]: + return list(await asyncio.gather(branch_a(), branch_b())) + + initial_task = asyncio.ensure_future(parallel_workflow.run("input")) + await branch_b_returned.wait() + intermediate = await storage.get_latest(workflow_name="parallel_workflow") + assert intermediate is not None + + release_branch_a.set() + initial = await initial_task + replayed = await parallel_workflow.run(checkpoint_id=intermediate.checkpoint_id) + + assert initial.get_outputs() == [["result:A", "result:B"]] + assert replayed.get_outputs() == initial.get_outputs() + assert calls == ["B", "A", "A"] + + async def test_repeated_automatic_identity_uses_occurrences(self): + storage = InMemoryCheckpointStorage() + call_count = 0 + + @step + async def echo(value: str) -> str: + nonlocal call_count + call_count += 1 + return value + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> list[str]: + return [await echo(value), await echo(value)] + + initial = await wf.run("same") + checkpoint = await storage.get_latest(workflow_name="wf") + assert checkpoint is not None + replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == [["same", "same"]] + assert replayed.get_outputs() == initial.get_outputs() + assert call_count == 2 + + async def test_automatic_identity_normalizes_bound_arguments(self): + @step + async def combine(first: int, second: int = 2) -> int: + return first + second + + positional = combine._get_replay_identity((1,), {}) # pyright: ignore[reportPrivateUsage] + keyword = combine._get_replay_identity((), {"first": 1, "second": 2}) # pyright: ignore[reportPrivateUsage] + + assert positional == keyword + + async def test_concurrent_opaque_arguments_require_replay_key(self): + import threading + + call_count = 0 + + @step + async def use_lock(lock: threading.Lock, value: str) -> str: + nonlocal call_count + call_count += 1 + return value + + @built_workflow + async def wf(_: str) -> list[str]: + return list( + await asyncio.gather( + use_lock(threading.Lock(), "A"), + use_lock(threading.Lock(), "B"), + ) + ) + + with pytest.raises(ValueError, match="replay_key"): + await wf.run("input") + + assert call_count == 0 + + async def test_explicit_replay_key_supports_opaque_concurrent_arguments(self): + import threading + + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + + @step(replay_key=lambda _lock, value: value) + async def use_lock(lock: threading.Lock, value: str) -> str: + calls.append(value) + return value + + @built_workflow(checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + return list( + await asyncio.gather( + use_lock(threading.Lock(), "A"), + use_lock(threading.Lock(), "B"), + ) + ) + + initial = await wf.run("input") + checkpoint = await storage.get_latest(workflow_name="wf") + assert checkpoint is not None + replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == [["A", "B"]] + assert replayed.get_outputs() == initial.get_outputs() + assert calls == ["A", "B"] + + async def test_explicit_replay_key_must_be_non_empty(self): + @step(replay_key=lambda _value: "") + async def invalid_key(value: str) -> str: + return value + + @built_workflow + async def wf(value: str) -> str: + return await invalid_key(value) + + with pytest.raises(ValueError, match="non-empty string"): + await wf.run("input") + + async def test_explicit_replay_key_must_be_unique_per_run(self): + @step(replay_key=lambda _value: "duplicate") + async def duplicated(value: str) -> str: + return value + + @built_workflow + async def wf(value: str) -> list[str]: + return [await duplicated(value), await duplicated(value)] + + with pytest.raises(ValueError, match="duplicate replay_key"): + await wf.run("input") + + async def test_legacy_checkpoint_replays_sequential_step(self): + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + call_count = 0 + + @step + async def compute(value: int) -> int: + nonlocal call_count + call_count += 1 + return value + 1 + + @built_workflow(name="legacy_wf", checkpoint_storage=storage) + async def wf(value: int) -> int: + return await compute(value) + + checkpoint = WorkflowCheckpoint( + workflow_name="legacy_wf", + 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) + + replayed = 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): + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + + @step + async def shared_step(value: str) -> str: + return f"result:{value}" + + @built_workflow(name="legacy_parallel", checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + return list(await asyncio.gather(shared_step("A"), shared_step("B"))) + + checkpoint = WorkflowCheckpoint( + workflow_name="legacy_parallel", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": {"shared_step::0": "result:B"}, + "_step_cache_auto_request_info_counts": {"shared_step::0": 0}, + "_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) + + async def test_versioned_identity_precedes_unrelated_legacy_entry(self): + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + + @step + async def mixed(value: str) -> str: + calls.append(value) + return f"live:{value}" + + @built_workflow(name="mixed_wf", checkpoint_storage=storage) + async def wf(_: str) -> str: + return (await asyncio.gather(mixed("child")))[0] + + replay_identity = mixed._get_replay_identity(("child",), {}) # pyright: ignore[reportPrivateUsage] + assert replay_identity is not None + identity_kind, identity = replay_identity + assert identity_kind == "auto" + key_ctx = _RunContext("mixed_wf") + versioned_key = key_ctx._get_automatic_step_cache_key( # pyright: ignore[reportPrivateUsage] + "mixed", + identity, + ) + checkpoint = WorkflowCheckpoint( + workflow_name="mixed_wf", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": { + "mixed::0": "legacy:other-call", + versioned_key: "cached:child", + }, + "_step_cache_auto_request_info_counts": { + "mixed::0": 0, + versioned_key: 0, + }, + "_original_message": "input", + }, + ) + checkpoint_id = await storage.save(checkpoint) + + replayed = await wf.run(checkpoint_id=checkpoint_id) + + assert replayed.get_outputs() == ["cached:child"] + assert calls == [] + + # --------------------------------------------------------------------------- # Branching / control flow # --------------------------------------------------------------------------- @@ -1301,6 +1606,35 @@ async def test_import_step_cache_non_integer_index(self): with pytest.raises(ValueError, match="Corrupted step cache"): ctx._import_step_cache({"step_name::abc": 42}) # pyright: ignore[reportPrivateUsage] + async def test_step_cache_round_trips_versioned_and_legacy_keys(self): + ctx = _RunContext("test") + automatic_key = ctx._get_automatic_step_cache_key( # pyright: ignore[reportPrivateUsage] + "automatic", + "a" * 64, + ) + explicit_key = ctx._get_explicit_step_cache_key( # pyright: ignore[reportPrivateUsage] + "explicit", + "b" * 64, + ) + ctx._step_cache = {automatic_key: "auto", explicit_key: "explicit", "legacy::0": "legacy"} + ctx._step_cache_auto_request_info_counts = {automatic_key: 1, explicit_key: 2, "legacy::0": 3} + + restored = _RunContext("test") + restored._import_step_cache(ctx._export_step_cache()) # pyright: ignore[reportPrivateUsage] + restored._import_step_cache_auto_request_info_counts( # pyright: ignore[reportPrivateUsage] + ctx._export_step_cache_auto_request_info_counts() + ) + + assert restored._step_cache == ctx._step_cache + assert restored._step_cache_auto_request_info_counts == ctx._step_cache_auto_request_info_counts + + async def test_import_step_cache_rejects_malformed_versioned_key(self): + ctx = _RunContext("test") + with pytest.raises(ValueError, match="Corrupted step cache"): + ctx._import_step_cache( # pyright: ignore[reportPrivateUsage] + {'v2::["auto","step","identity",-1]': 42} + ) + # --------------------------------------------------------------------------- # executor_bypassed event on replay (review comment #3) @@ -1525,8 +1859,8 @@ 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._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( diff --git a/python/samples/03-workflows/functional/parallel_pipeline.py b/python/samples/03-workflows/functional/parallel_pipeline.py index 87a695bf8e5..8424d9764ee 100644 --- a/python/samples/03-workflows/functional/parallel_pipeline.py +++ b/python/samples/03-workflows/functional/parallel_pipeline.py @@ -47,7 +47,9 @@ async def research_pipeline(topic: str) -> str: # # Tip: if any of these were wrapped with @step (e.g. an expensive agent call), # the pattern is identical — @step composes with asyncio.gather, so each - # branch is independently cached on HITL resume or checkpoint restore. + # branch is independently cached by its arguments on HITL resume or checkpoint + # restore. Use replay_key when arguments are opaque or do not uniquely identify + # non-interchangeable concurrent calls to the same step. web, papers, news = await asyncio.gather( research_web(topic), research_papers(topic), From 4363d0277fba7dd6ea9d734da6cf3ab81938bdf5 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 10:25:11 +0200 Subject: [PATCH 2/7] Python: address replay concurrency feedback --- .../agent_framework/_workflows/_functional.py | 81 +++++++++- .../workflow/test_functional_workflow.py | 151 +++++++++++++++++- 2 files changed, 218 insertions(+), 14 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 6d4e29a0c0b..48e905ea001 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -47,8 +47,8 @@ import logging import math import typing -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence -from contextvars import ContextVar +from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Mapping, Sequence +from contextvars import Context, ContextVar from copy import deepcopy from typing import Any, Generic, Literal, TypeVar, overload @@ -81,6 +81,59 @@ # ContextVar holding the active RunContext during workflow execution. # ContextVar is per-asyncio-Task, so concurrent workflows each get their own context. _active_run_ctx: ContextVar[RunContext | None] = ContextVar("_active_run_ctx", default=None) +_workflow_task_factory_states: dict[asyncio.AbstractEventLoop, _WorkflowTaskFactoryState] = {} + + +class _WorkflowTaskFactoryState: + """Delegate a loop task factory while recording tasks created by active workflows.""" + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self.loop = loop + self.original_factory = loop.get_task_factory() + self.active_runs = 0 + self.factory = self._create_task + + def _create_task( + self, + loop: asyncio.AbstractEventLoop, + coro: Coroutine[Any, Any, Any], + **kwargs: Any, + ) -> asyncio.Future[Any]: + if self.original_factory is None: + task = asyncio.Task(coro, loop=loop, **kwargs) + else: + task = self.original_factory(loop, coro, **kwargs) + + task_context = kwargs.get("context") + ctx = task_context.get(_active_run_ctx) if isinstance(task_context, Context) else _active_run_ctx.get() + if ctx is not None: + ctx._workflow_tasks.add(task) + return task + + +def _track_workflow_tasks() -> Callable[[], None]: + loop = asyncio.get_running_loop() + state = _workflow_task_factory_states.get(loop) + if state is None or loop.get_task_factory() is not state.factory: + state = _WorkflowTaskFactoryState(loop) + _workflow_task_factory_states[loop] = state + loop.set_task_factory(state.factory) + state.active_runs += 1 + released = False + + def _release() -> None: + nonlocal released + if released: + return + released = True + state.active_runs -= 1 + if state.active_runs == 0: + if loop.get_task_factory() is state.factory: + loop.set_task_factory(state.original_factory) + if _workflow_task_factory_states.get(loop) is state: + del _workflow_task_factory_states[loop] + + return _release def _canonicalize_step_identity_value(value: Any) -> Any: @@ -289,6 +342,8 @@ def __init__( self._step_identity_counters: dict[tuple[str, str], int] = {} # Explicit replay keys identify exactly one logical invocation per run. self._used_explicit_step_cache_keys: set[str] = set() + # Tasks created while this workflow context is active. + self._workflow_tasks: set[asyncio.Future[Any]] = set() # The task executing the workflow body; gather/task-group children differ from this task. self._root_task: asyncio.Task[Any] | None = None # Deterministic call counter for auto-generated request_info IDs @@ -459,9 +514,13 @@ def _get_explicit_step_cache_key(self, step_name: str, identity: str) -> str: self._used_explicit_step_cache_keys.add(key) return key - def _is_concurrent_task(self) -> bool: + def _is_concurrent_execution(self) -> bool: current_task = asyncio.current_task() - return self._root_task is not None and current_task is not None and current_task is not self._root_task + if self._root_task is None or current_task is None: + return False + if current_task is not self._root_task: + return True + return any(task is not current_task and not task.done() for task in self._workflow_tasks) def _get_cached_result(self, key: str) -> tuple[bool, Any]: if key in self._step_cache: @@ -644,7 +703,11 @@ def _get_replay_identity( raise ValueError(f"@step '{self.name}' replay_key must return a non-empty string.") return ("explicit", hashlib.sha256(explicit_key.encode("utf-8")).hexdigest()) - bound = self._identity_signature.bind(*args, **kwargs) + 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: @@ -713,7 +776,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: found, cached = ctx._get_cached_result(legacy_cache_key) if found: - if ctx._is_concurrent_task(): + 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 " @@ -721,7 +784,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: ) return await self._return_cached_result(ctx, legacy_cache_key, cached) - if replay_identity is None and ctx._is_concurrent_task(): + if replay_identity is None and ctx._is_concurrent_execution(): raise ValueError( f"Cannot derive a stable replay identity for concurrent @step '{self.name}'. " "Use @step(replay_key=...) to identify each logical invocation." @@ -1230,7 +1293,6 @@ async def _run_core( # Build context ctx = RunContext(self.name, streaming=streaming, run_kwargs=kwargs if kwargs else None) - ctx._root_task = asyncio.current_task() # Restore from checkpoint if requested prev_checkpoint_id: str | None = None @@ -1398,6 +1460,8 @@ async def _execute(self, ctx: RunContext, message: Any) -> Any: ) token = _active_run_ctx.set(ctx) + ctx._root_task = asyncio.current_task() + release_task_tracking = _track_workflow_tasks() try: sig = inspect.signature(self._func) params = list(sig.parameters.values()) @@ -1433,6 +1497,7 @@ async def _execute(self, ctx: RunContext, message: Any) -> Any: return await self._func(*call_args) finally: + release_task_tracking() _active_run_ctx.reset(token) # ------------------------------------------------------------------ diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 6214c12dc37..1abc9548869 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -971,8 +971,8 @@ async def parallel_workflow(_: str) -> list[str]: return list(await asyncio.gather(branch_a(), branch_b())) initial = await parallel_workflow.run("input") - checkpoint = await storage.get_latest(workflow_name="parallel_workflow") - assert checkpoint is not None + checkpoints = await storage.list_checkpoints(workflow_name="parallel_workflow") + checkpoint = checkpoints[-1] replayed = await parallel_workflow.run(checkpoint_id=checkpoint.checkpoint_id) @@ -1032,8 +1032,8 @@ async def wf(value: str) -> list[str]: return [await echo(value), await echo(value)] initial = await wf.run("same") - checkpoint = await storage.get_latest(workflow_name="wf") - assert checkpoint is not None + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) assert initial.get_outputs() == [["same", "same"]] @@ -1050,6 +1050,30 @@ async def combine(first: int, second: int = 2) -> int: assert positional == keyword + async def test_automatic_identity_ignores_explicit_context_keyword(self): + storage = InMemoryCheckpointStorage() + call_count = 0 + + @step + async def use_context(value: str, ctx: RunContext) -> str: + nonlocal call_count + call_count += 1 + return f"{value}:{ctx.get_state('marker')}" + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str, ctx: RunContext) -> str: + ctx.set_state("marker", "ok") + return await use_context(value, ctx=ctx) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["input:ok"] + assert replayed.get_outputs() == initial.get_outputs() + assert call_count == 1 + async def test_concurrent_opaque_arguments_require_replay_key(self): import threading @@ -1096,8 +1120,8 @@ async def wf(_: str) -> list[str]: ) initial = await wf.run("input") - checkpoint = await storage.get_latest(workflow_name="wf") - assert checkpoint is not None + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) assert initial.get_outputs() == [["A", "B"]] @@ -1187,6 +1211,121 @@ async def wf(_: str) -> list[str]: with pytest.raises(ValueError, match="legacy order-based cache entry"): await wf.run(checkpoint_id=checkpoint_id) + async def test_legacy_checkpoint_rejects_root_hit_while_child_is_active(self): + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + child_started = asyncio.Event() + release_child = asyncio.Event() + + @step + async def shared_step(value: str) -> str: + return f"result:{value}" + + async def active_child() -> None: + child_started.set() + await release_child.wait() + + @built_workflow(name="legacy_root_parallel", checkpoint_storage=storage) + async def wf(_: str) -> str: + child = asyncio.create_task(active_child()) + await child_started.wait() + try: + return await shared_step("root") + finally: + release_child.set() + await child + + checkpoint = WorkflowCheckpoint( + workflow_name="legacy_root_parallel", + graph_signature_hash=wf.graph_signature_hash, + state={ + "_step_cache": {"shared_step::0": "result:branch"}, + "_step_cache_auto_request_info_counts": {"shared_step::0": 0}, + "_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) + + async def test_legacy_sequential_replay_ignores_unrelated_caller_task(self): + from agent_framework import WorkflowCheckpoint + + storage = InMemoryCheckpointStorage() + workflow_waiting = asyncio.Event() + release_workflow = asyncio.Event() + release_unrelated = asyncio.Event() + + @step + async def compute(value: int) -> int: + return value + 1 + + @built_workflow(name="legacy_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) + loop = asyncio.get_running_loop() + original_factory = loop.get_task_factory() + + replay_task = asyncio.ensure_future(wf.run(checkpoint_id=checkpoint_id)) + await workflow_waiting.wait() + unrelated_task = asyncio.create_task(release_unrelated.wait()) + release_workflow.set() + replayed = await replay_task + + assert replayed.get_outputs() == [6] + assert loop.get_task_factory() is original_factory + + release_unrelated.set() + await unrelated_task + + async def test_task_tracking_preserves_existing_loop_factory(self): + loop = asyncio.get_running_loop() + original_factory = loop.get_task_factory() + created_tasks: list[asyncio.Future[Any]] = [] + + def custom_factory( + task_loop: asyncio.AbstractEventLoop, + coro: Any, + **kwargs: Any, + ) -> asyncio.Future[Any]: + task = asyncio.Task(coro, loop=task_loop, **kwargs) + created_tasks.append(task) + return task + + loop.set_task_factory(custom_factory) + try: + + @step + async def compute(value: int) -> int: + return value + 1 + + @built_workflow + async def wf(value: int) -> int: + return (await asyncio.gather(compute(value)))[0] + + result = await wf.run(1) + + assert result.get_outputs() == [2] + assert created_tasks + assert loop.get_task_factory() is custom_factory + finally: + loop.set_task_factory(original_factory) + async def test_versioned_identity_precedes_unrelated_legacy_entry(self): from agent_framework import WorkflowCheckpoint From cdbad962cb7691d3e002c4d1798287c6d0c4408f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 11:07:14 +0200 Subject: [PATCH 3/7] Python: harden replay identity edge cases --- .../agent_framework/_workflows/_functional.py | 264 +++++++++++--- .../workflow/test_functional_workflow.py | 329 +++++++++++++++++- 2 files changed, 532 insertions(+), 61 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 48e905ea001..3cc9d7594c6 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -50,7 +50,9 @@ from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Mapping, Sequence from contextvars import Context, ContextVar from copy import deepcopy +from types import CodeType from typing import Any, Generic, Literal, TypeVar, overload +from weakref import WeakSet from .._agents import BaseAgent from .._feature_stage import ExperimentalFeature, experimental @@ -77,6 +79,7 @@ _STEP_CACHE_KEY_V2_PREFIX = "v2::" _UNSUPPORTED_STEP_IDENTITY = object() +_step_wrapper_identity_registry: dict[tuple[str, str], WeakSet[Any]] = {} # ContextVar holding the active RunContext during workflow execution. # ContextVar is per-asyncio-Task, so concurrent workflows each get their own context. @@ -136,8 +139,10 @@ def _release() -> None: return _release -def _canonicalize_step_identity_value(value: Any) -> Any: +def _canonicalize_step_identity_value(value: Any, seen: set[int] | None = None) -> Any: """Return a type-preserving JSON value, or a sentinel for unsupported input.""" + if seen is None: + seen = set() if value is None: return ["none"] if type(value) is bool: @@ -148,36 +153,74 @@ def _canonicalize_step_identity_value(value: Any) -> Any: if not math.isfinite(value): return _UNSUPPORTED_STEP_IDENTITY return ["float", value.hex()] + if type(value) is complex: + if not math.isfinite(value.real) or not math.isfinite(value.imag): + return _UNSUPPORTED_STEP_IDENTITY + return ["complex", value.real.hex(), value.imag.hex()] if type(value) is str: return ["str", value] + if type(value) is bytes: + return ["bytes", value.hex()] if type(value) is list: + list_value = typing.cast(list[Any], value) + value_id = id(list_value) + if value_id in seen: + return _UNSUPPORTED_STEP_IDENTITY + seen.add(value_id) list_items: list[Any] = [] - for item in typing.cast(list[Any], value): - canonical = _canonicalize_step_identity_value(item) - if canonical is _UNSUPPORTED_STEP_IDENTITY: - return _UNSUPPORTED_STEP_IDENTITY - list_items.append(canonical) - return ["list", list_items] + try: + for item in list_value: + canonical = _canonicalize_step_identity_value(item, seen) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + list_items.append(canonical) + return ["list", list_items] + finally: + seen.remove(value_id) if type(value) is tuple: + tuple_value = typing.cast(tuple[Any, ...], value) + value_id = id(tuple_value) + if value_id in seen: + return _UNSUPPORTED_STEP_IDENTITY + seen.add(value_id) tuple_items: list[Any] = [] - for item in typing.cast(tuple[Any, ...], value): - canonical = _canonicalize_step_identity_value(item) - if canonical is _UNSUPPORTED_STEP_IDENTITY: - return _UNSUPPORTED_STEP_IDENTITY - tuple_items.append(canonical) - return ["tuple", tuple_items] + try: + for item in tuple_value: + canonical = _canonicalize_step_identity_value(item, seen) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + tuple_items.append(canonical) + return ["tuple", tuple_items] + finally: + seen.remove(value_id) if isinstance(value, Mapping): mapping = typing.cast(Mapping[Any, Any], value) + value_id = id(mapping) + if value_id in seen: + return _UNSUPPORTED_STEP_IDENTITY + seen.add(value_id) mapping_items: list[Any] = [] - for key, item in mapping.items(): - if type(key) is not str: - return _UNSUPPORTED_STEP_IDENTITY - canonical = _canonicalize_step_identity_value(item) + try: + for key, item in mapping.items(): + if type(key) is not str: + return _UNSUPPORTED_STEP_IDENTITY + canonical = _canonicalize_step_identity_value(item, seen) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return _UNSUPPORTED_STEP_IDENTITY + mapping_items.append([key, canonical]) + mapping_type = f"{type(mapping).__module__}.{type(mapping).__qualname__}" + return ["mapping", mapping_type, mapping_items] + finally: + seen.remove(value_id) + if type(value) is frozenset: + frozenset_items: list[Any] = [] + for item in typing.cast(frozenset[Any], value): + canonical = _canonicalize_step_identity_value(item, seen) if canonical is _UNSUPPORTED_STEP_IDENTITY: return _UNSUPPORTED_STEP_IDENTITY - mapping_items.append([key, canonical]) - mapping_type = f"{type(mapping).__module__}.{type(mapping).__qualname__}" - return ["mapping", mapping_type, mapping_items] + frozenset_items.append(canonical) + frozenset_items.sort(key=lambda item: json.dumps(item, sort_keys=True, separators=(",", ":"))) + return ["frozenset", frozenset_items] return _UNSUPPORTED_STEP_IDENTITY @@ -189,13 +232,105 @@ def _hash_step_identity(value: Any) -> str | None: return hashlib.sha256(encoded.encode("utf-8")).hexdigest() +def _canonicalize_code_value(value: Any) -> Any: + if isinstance(value, CodeType): + return ["code", _get_code_identity_payload(value)] + if value is None or type(value) in (bool, int, str): + return [type(value).__name__, value] + if type(value) is float: + return ["float", value.hex()] + if type(value) is complex: + return ["complex", value.real.hex(), value.imag.hex()] + if type(value) is bytes: + return ["bytes", value.hex()] + if type(value) is tuple: + tuple_value = typing.cast(tuple[Any, ...], value) + return ["tuple", [_canonicalize_code_value(item) for item in tuple_value]] + if type(value) is frozenset: + frozenset_value = typing.cast(frozenset[Any], value) + items = [_canonicalize_code_value(item) for item in frozenset_value] + items.sort(key=lambda item: json.dumps(item, sort_keys=True, separators=(",", ":"))) + return ["frozenset", items] + if value is Ellipsis: + return ["ellipsis"] + return ["unsupported", type(value).__module__, type(value).__qualname__] + + +def _get_code_identity_payload(code: CodeType) -> dict[str, Any]: + return { + "argcount": code.co_argcount, + "posonlyargcount": code.co_posonlyargcount, + "kwonlyargcount": code.co_kwonlyargcount, + "flags": code.co_flags, + "code": code.co_code.hex(), + "consts": [_canonicalize_code_value(item) for item in code.co_consts], + "names": list(code.co_names), + "varnames": list(code.co_varnames), + "freevars": list(code.co_freevars), + "cellvars": list(code.co_cellvars), + } + + +def _canonicalize_wrapper_state(value: Any) -> tuple[Any, bool]: + canonical = _canonicalize_step_identity_value(value) + if canonical is _UNSUPPORTED_STEP_IDENTITY: + return (["opaque", type(value).__module__, type(value).__qualname__], False) + return (canonical, _is_immutable_wrapper_state(value)) + + +def _is_immutable_wrapper_state(value: Any) -> bool: + if value is None or type(value) in (bool, int, float, complex, str, bytes): + return True + if type(value) is tuple: + return all(_is_immutable_wrapper_state(item) for item in typing.cast(tuple[Any, ...], value)) + if type(value) is frozenset: + return all(_is_immutable_wrapper_state(item) for item in typing.cast(frozenset[Any], value)) + return False + + +def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str, bool]: + code = getattr(func, "__code__", None) + defaults, defaults_are_durable = _canonicalize_wrapper_state(func.__defaults__ or ()) + raw_kwdefaults = func.__kwdefaults__ or {} + kwdefaults, _ = _canonicalize_wrapper_state(raw_kwdefaults) + kwdefaults_are_durable = all(_is_immutable_wrapper_state(value) for value in raw_kwdefaults.values()) + closure_items: list[Any] = [] + closure_is_durable = func.__closure__ is None + if code is not None and func.__closure__ is not None: + for name, cell in zip(code.co_freevars, func.__closure__): + try: + value = cell.cell_contents + except ValueError: + closure_items.append([name, ["empty"]]) + closure_is_durable = False + continue + canonical, _ = _canonicalize_wrapper_state(value) + closure_items.append([name, canonical]) + + payload = { + "module": func.__module__, + "qualname": func.__qualname__, + "firstlineno": code.co_firstlineno if code is not None else None, + "code": _get_code_identity_payload(code) if code is not None else None, + "defaults": defaults, + "kwdefaults": kwdefaults, + "closure": closure_items, + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return ( + hashlib.sha256(encoded.encode("utf-8")).hexdigest(), + defaults_are_durable and kwdefaults_are_durable and closure_is_durable, + ) + + def _encode_step_cache_key( kind: Literal["auto", "explicit"], step_name: str, + wrapper_identity: str, identity: str, occurrence: int | None = None, ) -> str: - payload: list[str | int] = [kind, step_name, identity] + payload: list[str | int] = [kind, step_name, wrapper_identity, identity] if kind == "auto": if occurrence is None: raise ValueError("Automatic step cache keys require an occurrence index.") @@ -217,13 +352,13 @@ def _validate_step_cache_key(key: Any) -> str: kind = payload[0] if ( kind == "auto" - and len(payload) == 4 - and all(isinstance(item, str) for item in payload[1:3]) - and isinstance(payload[3], int) - and payload[3] >= 0 + and len(payload) == 5 + and all(isinstance(item, str) for item in payload[1:4]) + and isinstance(payload[4], int) + and payload[4] >= 0 ): return key - if kind == "explicit" and len(payload) == 3 and all(isinstance(item, str) for item in payload[1:]): + if kind == "explicit" and len(payload) == 4 and all(isinstance(item, str) for item in payload[1:]): return key raise ValueError("Invalid versioned step cache key.") @@ -339,7 +474,7 @@ def __init__( # Legacy per-step call counters retained for old checkpoint compatibility. self._step_call_counters: dict[str, int] = {} # Per-identity counters keep repeated automatic invocations distinct. - self._step_identity_counters: dict[tuple[str, str], int] = {} + self._step_identity_counters: dict[tuple[str, str, str], int] = {} # Explicit replay keys identify exactly one logical invocation per run. self._used_explicit_step_cache_keys: set[str] = set() # Tasks created while this workflow context is active. @@ -498,14 +633,14 @@ def _get_legacy_step_cache_key(self, step_name: str) -> str: self._step_call_counters[step_name] = idx + 1 return f"{step_name}::{idx}" - def _get_automatic_step_cache_key(self, step_name: str, identity: str) -> str: - counter_key = (step_name, identity) + def _get_automatic_step_cache_key(self, step_name: str, wrapper_identity: str, identity: str) -> str: + counter_key = (step_name, wrapper_identity, identity) occurrence = self._step_identity_counters.get(counter_key, 0) self._step_identity_counters[counter_key] = occurrence + 1 - return _encode_step_cache_key("auto", step_name, identity, occurrence) + return _encode_step_cache_key("auto", step_name, wrapper_identity, identity, occurrence) - def _get_explicit_step_cache_key(self, step_name: str, identity: str) -> str: - key = _encode_step_cache_key("explicit", step_name, identity) + def _get_explicit_step_cache_key(self, step_name: str, wrapper_identity: str, identity: str) -> str: + key = _encode_step_cache_key("explicit", step_name, wrapper_identity, identity) if key in self._used_explicit_step_cache_keys: raise ValueError( f"@step '{step_name}' produced duplicate replay_key values in one workflow run. " @@ -632,9 +767,13 @@ class StepWrapper(Generic[R]): isolation. Automatic replay identity supports ``None``, booleans, integers, finite - floats, strings, lists, tuples, and string-key mappings containing those - values. Concurrent calls with the same automatic identity must produce - interchangeable results; use distinct ``replay_key`` values otherwise. + floats/complex numbers, strings, bytes, lists, tuples, frozensets, and + string-key mappings containing those values. Concurrent calls with the + same automatic identity must produce interchangeable results; use + distinct ``replay_key`` values otherwise. + Opaque captured defaults/closure state, and multiple wrappers whose + canonical source and captured state are identical, require ``replay_key`` + for concurrent use. Legacy order-based checkpoints remain replayable sequentially, but a concurrent legacy cache hit raises instead of risking a mismatched result. @@ -666,6 +805,15 @@ def __init__( self._func = func self.name: str = name or func.__name__ self._signature = inspect.signature(func) + self._wrapper_identity, self._wrapper_identity_is_durable = _get_step_wrapper_identity(func) + if _active_run_ctx.get() is not None: + self._wrapper_identity_is_durable = False + self._wrapper_registry_key: tuple[str, str] | None = None + if self._wrapper_identity_is_durable: + registry_key = (self.name, self._wrapper_identity) + matching_wrappers = _step_wrapper_identity_registry.setdefault(registry_key, WeakSet()) + matching_wrappers.add(self) + self._wrapper_registry_key = registry_key self._replay_key = replay_key functools.update_wrapper(self, func) @@ -697,12 +845,33 @@ def _get_replay_identity( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> tuple[Literal["auto", "explicit"], str] | None: + if self._wrapper_identity_is_durable: + current_identity, current_is_durable = _get_step_wrapper_identity(self._func) + if not current_is_durable or current_identity != self._wrapper_identity: + raise ValueError( + f"@step '{self.name}' defaults or captured state changed after decoration. " + "Create a new workflow definition or provide a versioned replay_key." + ) + if self._replay_key is not None: explicit_key = self._replay_key(*args, **kwargs) if not isinstance(explicit_key, str) or not explicit_key: raise ValueError(f"@step '{self.name}' replay_key must return a non-empty string.") return ("explicit", hashlib.sha256(explicit_key.encode("utf-8")).hexdigest()) + if self._wrapper_registry_key is not None: + matching_wrappers = _step_wrapper_identity_registry.get(self._wrapper_registry_key) + else: + matching_wrappers = None + if matching_wrappers is not None and len(matching_wrappers) > 1: + raise ValueError( + f"@step '{self.name}' was generated from a function definition used by multiple wrappers. " + "Provide replay_key so each logical invocation has an explicit durable identity." + ) + + if not self._wrapper_identity_is_durable: + return None + identity_kwargs = kwargs if self._ctx_param_name is not None and self._ctx_param_name in kwargs: identity_kwargs = dict(kwargs) @@ -766,9 +935,9 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: if replay_identity is not None: identity_kind, identity = replay_identity if identity_kind == "explicit": - cache_key = ctx._get_explicit_step_cache_key(self.name, identity) + cache_key = ctx._get_explicit_step_cache_key(self.name, self._wrapper_identity, identity) else: - cache_key = ctx._get_automatic_step_cache_key(self.name, identity) + cache_key = ctx._get_automatic_step_cache_key(self.name, self._wrapper_identity, identity) found, cached = ctx._get_cached_result(cache_key) if found: @@ -867,10 +1036,10 @@ def step( Automatic replay identity is derived from bound canonical arguments. Concurrent calls with identical automatic identities must produce interchangeable results. Use distinct ``replay_key`` values when opaque - arguments or other state distinguish the logical invocations. Legacy - order-based checkpoints replay sequential steps, but concurrent legacy - cache hits fail explicitly rather than returning a potentially mismatched - result. + arguments, captured state, generated wrappers, or other state distinguish + the logical invocations. Legacy order-based checkpoints replay sequential + steps, but concurrent legacy cache hits fail explicitly rather than + returning a potentially mismatched result. Args: func: The async function to decorate (when using the bare @@ -1042,10 +1211,17 @@ def __init__( functools.update_wrapper(self, func) # type: ignore[arg-type] + @staticmethod + def _snapshot_replay_message(message: Any) -> Any: + try: + return deepcopy(message) + except Exception: + return message + 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 + if message is not None and self._last_message is None: + self._last_message = self._snapshot_replay_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) @@ -1331,7 +1507,7 @@ async def _run_core( # Store message for future replays if message is not None: - self._last_message = message + self._last_message = self._snapshot_replay_message(message) # Set responses for replay if responses: diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 1abc9548869..da32b5c79e6 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -11,6 +11,7 @@ from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass +from types import FunctionType from typing import Any, overload import pytest @@ -35,9 +36,10 @@ step, workflow, ) -from agent_framework._workflows._functional import ( - RunContext as _RunContext, -) +from agent_framework._workflows._functional import RunContext as _RunContext +from agent_framework._workflows._functional import _get_step_wrapper_identity + +_factory_step_calls: list[str] = [] # --------------------------------------------------------------------------- # Helpers @@ -942,13 +944,13 @@ async def test_checkpoint_restore_preserves_concurrent_step_invocations(self): branch_a_shared_completed = asyncio.Event() calls: list[str] = [] - @step + @step(replay_key=lambda value: f"predecessor:{value}") async def predecessor(value: str) -> str: predecessor_started.set() await release_predecessor.wait() return value - @step + @step(replay_key=lambda value: value) async def shared_step(value: str) -> str: calls.append(value) if value == "B": @@ -980,13 +982,237 @@ async def parallel_workflow(_: str) -> list[str]: assert replayed.get_outputs() == initial.get_outputs() assert calls == ["B", "A"] + async def test_same_named_wrappers_keep_distinct_concurrent_identities(self): + storage = InMemoryCheckpointStorage() + predecessor_started = asyncio.Event() + release_predecessor = asyncio.Event() + read_completed = asyncio.Event() + calls: list[str] = [] + + @step(replay_key=lambda value: f"predecessor:{value}") + async def predecessor(value: str) -> str: + predecessor_started.set() + await release_predecessor.wait() + return value + + @step(name="authorization", replay_key=lambda value: f"read:{value}") + async def read_check(value: str) -> str: + calls.append("read") + read_completed.set() + return f"read:{value}" + + @step(name="authorization", replay_key=lambda value: f"delete:{value}") + async def delete_check(value: str) -> str: + calls.append("delete") + release_predecessor.set() + await read_completed.wait() + return f"delete:{value}" + + async def branch_a() -> str: + await predecessor("same") + return await read_check("same") + + async def branch_b() -> str: + await predecessor_started.wait() + return await delete_check("same") + + @built_workflow(checkpoint_storage=storage) + async def wf(_: str) -> list[str]: + return list(await asyncio.gather(branch_a(), branch_b())) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + replayed = await wf.run(checkpoint_id=checkpoints[-1].checkpoint_id) + + assert initial.get_outputs() == [["read:same", "delete:same"]] + assert replayed.get_outputs() == initial.get_outputs() + assert calls == ["delete", "read"] + + async def test_factory_generated_wrappers_use_canonical_default_identity(self): + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + def make_step(label: str) -> StepWrapper[str]: + async def generated(value: str, label: str = label) -> str: + _factory_step_calls.append(label) + return f"{label}:{value}" + + return step(name="generated")(generated) + + first = make_step("first") + second = make_step("second") + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> list[str]: + return list(await asyncio.gather(first(value), second(value))) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + replayed = await wf.run(checkpoint_id=checkpoints[-1].checkpoint_id) + + assert initial.get_outputs() == [["first:input", "second:input"]] + assert replayed.get_outputs() == initial.get_outputs() + assert _factory_step_calls == ["first", "second"] + + async def test_factory_generated_wrappers_with_opaque_state_require_replay_key(self): + @dataclass + class Token: + label: str + + calls: list[str] = [] + + def make_step(token: Token) -> StepWrapper[str]: + async def generated(value: str) -> str: + calls.append(token.label) + return f"{token.label}:{value}" + + return step(name="generated")(generated) + + first = make_step(Token("first")) + second = make_step(Token("second")) + + @built_workflow + async def wf(value: str) -> list[str]: + return list(await asyncio.gather(first(value), second(value))) + + with pytest.raises(ValueError, match="replay_key"): + await wf.run("input") + + assert calls == [] + + async def test_factory_generated_wrappers_with_opaque_state_replay_with_explicit_keys(self): + @dataclass + class Token: + label: str + + storage = InMemoryCheckpointStorage() + calls: list[str] = [] + + def make_step(token: Token) -> StepWrapper[str]: + async def generated(value: str) -> str: + calls.append(token.label) + return f"{token.label}:{value}" + + return step(name="generated", replay_key=lambda _value: token.label)(generated) + + first = make_step(Token("first")) + second = make_step(Token("second")) + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> list[str]: + return list(await asyncio.gather(first(value), second(value))) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + replayed = await wf.run(checkpoint_id=checkpoints[-1].checkpoint_id) + + assert initial.get_outputs() == [["first:input", "second:input"]] + assert replayed.get_outputs() == initial.get_outputs() + assert calls == ["first", "second"] + + async def test_existing_generated_wrapper_fails_closed_after_collision(self): + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + def make_step() -> StepWrapper[str]: + async def generated(value: str) -> str: + _factory_step_calls.append(value) + return value + + return step(name="generated")(generated) + + first = make_step() + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await first(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + second = make_step() + + with pytest.raises(ValueError, match="multiple wrappers"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + del second + gc.collect() + replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["input"] + assert replayed.get_outputs() == initial.get_outputs() + assert _factory_step_calls == ["input"] + + async def test_cyclic_arguments_use_sequential_fallback(self): + call_count = 0 + + @step + async def count_items(items: list[Any]) -> int: + nonlocal call_count + call_count += 1 + return len(items) + + @built_workflow + async def wf(items: list[Any]) -> int: + return await count_items(items) + + cyclic: list[Any] = [] + cyclic.append(cyclic) + result = await wf.run(cyclic) + + assert result.get_outputs() == [1] + assert call_count == 1 + + async def test_cyclic_concurrent_arguments_require_replay_key(self): + call_count = 0 + + @step + async def count_items(items: list[Any]) -> int: + nonlocal call_count + call_count += 1 + return len(items) + + @built_workflow + async def wf(items: list[Any]) -> int: + return (await asyncio.gather(count_items(items)))[0] + + cyclic: list[Any] = [] + cyclic.append(cyclic) + + with pytest.raises(ValueError, match="replay_key"): + await wf.run(cyclic) + + assert call_count == 0 + + async def test_mutated_workflow_input_keeps_original_replay_identity(self): + storage = InMemoryCheckpointStorage() + + @step + async def mutate(items: list[str]) -> list[str]: + items.append("mutated") + return items + + @built_workflow(checkpoint_storage=storage) + async def wf(items: list[str]) -> list[str]: + return await mutate(items) + + message = ["original"] + initial = await wf.run(message) + checkpoints = await storage.list_checkpoints(workflow_name="wf") + replayed = await wf.run(checkpoint_id=checkpoints[-1].checkpoint_id) + + assert message == ["original", "mutated"] + assert initial.get_outputs() == [["original", "mutated"]] + assert replayed.get_outputs() == initial.get_outputs() + assert len([event for event in replayed if event.type == "executor_bypassed"]) == 1 + async def test_intermediate_checkpoint_does_not_cross_branches(self): storage = InMemoryCheckpointStorage() release_branch_a = asyncio.Event() branch_b_returned = asyncio.Event() calls: list[str] = [] - @step + @step(replay_key=lambda value: value) async def shared_step(value: str) -> str: calls.append(value) return f"result:{value}" @@ -1019,12 +1245,9 @@ async def parallel_workflow(_: str) -> list[str]: async def test_repeated_automatic_identity_uses_occurrences(self): storage = InMemoryCheckpointStorage() - call_count = 0 @step async def echo(value: str) -> str: - nonlocal call_count - call_count += 1 return value @built_workflow(checkpoint_storage=storage) @@ -1038,7 +1261,7 @@ async def wf(value: str) -> list[str]: assert initial.get_outputs() == [["same", "same"]] assert replayed.get_outputs() == initial.get_outputs() - assert call_count == 2 + assert len([event for event in replayed if event.type == "executor_bypassed"]) == 2 async def test_automatic_identity_normalizes_bound_arguments(self): @step @@ -1052,12 +1275,9 @@ async def combine(first: int, second: int = 2) -> int: async def test_automatic_identity_ignores_explicit_context_keyword(self): storage = InMemoryCheckpointStorage() - call_count = 0 @step async def use_context(value: str, ctx: RunContext) -> str: - nonlocal call_count - call_count += 1 return f"{value}:{ctx.get_state('marker')}" @built_workflow(checkpoint_storage=storage) @@ -1072,7 +1292,7 @@ async def wf(value: str, ctx: RunContext) -> str: assert initial.get_outputs() == ["input:ok"] assert replayed.get_outputs() == initial.get_outputs() - assert call_count == 1 + assert len([event for event in replayed if event.type == "executor_bypassed"]) == 1 async def test_concurrent_opaque_arguments_require_replay_key(self): import threading @@ -1152,6 +1372,28 @@ async def wf(value: str) -> list[str]: with pytest.raises(ValueError, match="duplicate replay_key"): await wf.run("input") + async def test_changed_immutable_kwdefault_fails_closed(self): + storage = InMemoryCheckpointStorage() + + @step(replay_key=lambda value: value) + async def marker(value: str, *, suffix: bytes = b"A") -> str: + return f"{value}:{suffix.decode()}" + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await marker(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + assert marker._func.__kwdefaults__ is not None # pyright: ignore[reportPrivateUsage] + marker._func.__kwdefaults__["suffix"] = b"B" # pyright: ignore[reportPrivateUsage] + + with pytest.raises(ValueError, match="changed after decoration"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["input:A"] + async def test_legacy_checkpoint_replays_sequential_step(self): from agent_framework import WorkflowCheckpoint @@ -1332,7 +1574,7 @@ async def test_versioned_identity_precedes_unrelated_legacy_entry(self): storage = InMemoryCheckpointStorage() calls: list[str] = [] - @step + @step(replay_key=lambda value: value) async def mixed(value: str) -> str: calls.append(value) return f"live:{value}" @@ -1344,10 +1586,11 @@ async def wf(_: str) -> str: replay_identity = mixed._get_replay_identity(("child",), {}) # pyright: ignore[reportPrivateUsage] assert replay_identity is not None identity_kind, identity = replay_identity - assert identity_kind == "auto" + assert identity_kind == "explicit" key_ctx = _RunContext("mixed_wf") - versioned_key = key_ctx._get_automatic_step_cache_key( # pyright: ignore[reportPrivateUsage] + versioned_key = key_ctx._get_explicit_step_cache_key( # pyright: ignore[reportPrivateUsage] "mixed", + mixed._wrapper_identity, # pyright: ignore[reportPrivateUsage] identity, ) checkpoint = WorkflowCheckpoint( @@ -1714,6 +1957,56 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointValidation: + def test_wrapper_identity_includes_code_constants(self): + async def plus_one(value: int) -> int: + return value + 1 + + async def plus_two(value: int) -> int: + return value + 2 + + first = FunctionType(plus_one.__code__.replace(co_firstlineno=1), globals(), "generated") + second = FunctionType(plus_two.__code__.replace(co_firstlineno=1), globals(), "generated") + first.__module__ = second.__module__ = "test_module" + first.__qualname__ = second.__qualname__ = "generated" + + first_identity, first_is_durable = _get_step_wrapper_identity(first) + second_identity, second_is_durable = _get_step_wrapper_identity(second) + + assert first_is_durable is True + assert second_is_durable is True + assert first_identity != second_identity + + def test_wrapper_identity_includes_defaults_and_closure_state(self): + async def template(value: int = 0) -> int: + return value + + first_default = FunctionType(template.__code__, globals(), "generated", (1,)) + second_default = FunctionType(template.__code__, globals(), "generated", (2,)) + first_default.__module__ = second_default.__module__ = "test_module" + first_default.__qualname__ = second_default.__qualname__ = "generated" + + def make(label: str) -> Callable[[int], Awaitable[str]]: + async def generated(value: int) -> str: + return f"{label}:{value}" + + return generated + + first_closure = make("first") + second_closure = make("second") + + assert _get_step_wrapper_identity(first_default)[0] != _get_step_wrapper_identity(second_default)[0] + assert _get_step_wrapper_identity(first_closure)[0] != _get_step_wrapper_identity(second_closure)[0] + + def test_wrapper_identity_marks_mutable_closure_state_non_durable(self): + state = ["A"] + + async def generated(value: str) -> str: + return f"{state[0]}:{value}" + + _, is_durable = _get_step_wrapper_identity(generated) + + assert is_durable is False + async def test_checkpoint_signature_mismatch_raises(self): from agent_framework import WorkflowCheckpoint @@ -1750,9 +2043,11 @@ async def test_step_cache_round_trips_versioned_and_legacy_keys(self): automatic_key = ctx._get_automatic_step_cache_key( # pyright: ignore[reportPrivateUsage] "automatic", "a" * 64, + "b" * 64, ) explicit_key = ctx._get_explicit_step_cache_key( # pyright: ignore[reportPrivateUsage] "explicit", + "c" * 64, "b" * 64, ) ctx._step_cache = {automatic_key: "auto", explicit_key: "explicit", "legacy::0": "legacy"} From f849efa400194c4e7355245b7915c66ac9deef79 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 11:08:27 +0200 Subject: [PATCH 4/7] Python: fix replay test typing --- .../core/tests/workflow/test_functional_workflow.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index da32b5c79e6..7be955e9246 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -12,7 +12,7 @@ from contextlib import contextmanager from dataclasses import dataclass from types import FunctionType -from typing import Any, overload +from typing import Any, cast, overload import pytest @@ -1386,8 +1386,9 @@ async def wf(value: str) -> str: initial = await wf.run("input") checkpoints = await storage.list_checkpoints(workflow_name="wf") checkpoint = checkpoints[-1] - assert marker._func.__kwdefaults__ is not None # pyright: ignore[reportPrivateUsage] - marker._func.__kwdefaults__["suffix"] = b"B" # pyright: ignore[reportPrivateUsage] + marker_func = cast(FunctionType, marker._func) # pyright: ignore[reportPrivateUsage] + assert marker_func.__kwdefaults__ is not None + marker_func.__kwdefaults__["suffix"] = b"B" with pytest.raises(ValueError, match="changed after decoration"): await wf.run(checkpoint_id=checkpoint.checkpoint_id) From 40dbf529f541c6378efe5bae4dd17f64f69facf9 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 11:26:57 +0200 Subject: [PATCH 5/7] Python: finalize replay durability checks --- .../agent_framework/_workflows/_functional.py | 143 ++++++++++++------ .../workflow/test_functional_workflow.py | 102 ++++++++++++- 2 files changed, 200 insertions(+), 45 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 3cc9d7594c6..9ce7e4cdec9 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -288,14 +288,25 @@ def _is_immutable_wrapper_state(value: Any) -> bool: return False -def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str, bool]: - code = getattr(func, "__code__", None) +def _get_step_wrapper_defaults_state(func: Callable[..., Awaitable[Any]]) -> tuple[dict[str, Any], bool]: defaults, defaults_are_durable = _canonicalize_wrapper_state(func.__defaults__ or ()) raw_kwdefaults = func.__kwdefaults__ or {} kwdefaults, _ = _canonicalize_wrapper_state(raw_kwdefaults) kwdefaults_are_durable = all(_is_immutable_wrapper_state(value) for value in raw_kwdefaults.values()) + return ( + { + "defaults": defaults, + "kwdefaults": kwdefaults, + }, + defaults_are_durable and kwdefaults_are_durable, + ) + + +def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str, bool]: + code = getattr(func, "__code__", None) + defaults_state, defaults_are_durable = _get_step_wrapper_defaults_state(func) closure_items: list[Any] = [] - closure_is_durable = func.__closure__ is None + closure_is_durable = True if code is not None and func.__closure__ is not None: for name, cell in zip(code.co_freevars, func.__closure__): try: @@ -304,22 +315,22 @@ def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str closure_items.append([name, ["empty"]]) closure_is_durable = False continue - canonical, _ = _canonicalize_wrapper_state(value) + canonical, is_durable = _canonicalize_wrapper_state(value) closure_items.append([name, canonical]) + closure_is_durable = closure_is_durable and is_durable payload = { "module": func.__module__, "qualname": func.__qualname__, "firstlineno": code.co_firstlineno if code is not None else None, "code": _get_code_identity_payload(code) if code is not None else None, - "defaults": defaults, - "kwdefaults": kwdefaults, + **defaults_state, "closure": closure_items, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) return ( hashlib.sha256(encoded.encode("utf-8")).hexdigest(), - defaults_are_durable and kwdefaults_are_durable and closure_is_durable, + defaults_are_durable and closure_is_durable, ) @@ -338,29 +349,38 @@ def _encode_step_cache_key( return f"{_STEP_CACHE_KEY_V2_PREFIX}{json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}" +def _decode_versioned_step_cache_key( + key: str, +) -> tuple[Literal["auto", "explicit"], str, str, str, int | None] | None: + if not key.startswith(_STEP_CACHE_KEY_V2_PREFIX): + return None + try: + raw_payload: object = json.loads(key.removeprefix(_STEP_CACHE_KEY_V2_PREFIX)) + except json.JSONDecodeError: + return None + if not isinstance(raw_payload, list) or not raw_payload: + return None + payload = typing.cast(list[Any], raw_payload) + kind = payload[0] + if ( + kind == "auto" + and len(payload) == 5 + and all(isinstance(item, str) for item in payload[1:4]) + and isinstance(payload[4], int) + and payload[4] >= 0 + ): + return ("auto", payload[1], payload[2], payload[3], payload[4]) + if kind == "explicit" and len(payload) == 4 and all(isinstance(item, str) for item in payload[1:]): + return ("explicit", payload[1], payload[2], payload[3], None) + raise ValueError("Invalid versioned step cache key.") + + def _validate_step_cache_key(key: Any) -> str: if not isinstance(key, str): raise TypeError("Step cache keys must be strings.") - if key.startswith(_STEP_CACHE_KEY_V2_PREFIX): - try: - raw_payload: object = json.loads(key.removeprefix(_STEP_CACHE_KEY_V2_PREFIX)) - except json.JSONDecodeError: - raw_payload = None - if isinstance(raw_payload, list) and raw_payload: - payload = typing.cast(list[Any], raw_payload) - kind = payload[0] - if ( - kind == "auto" - and len(payload) == 5 - and all(isinstance(item, str) for item in payload[1:4]) - and isinstance(payload[4], int) - and payload[4] >= 0 - ): - return key - if kind == "explicit" and len(payload) == 4 and all(isinstance(item, str) for item in payload[1:]): - return key - raise ValueError("Invalid versioned step cache key.") + if _decode_versioned_step_cache_key(key) is not None: + return key name, idx_str = key.rsplit("::", 1) if not name or int(idx_str) < 0: @@ -368,6 +388,15 @@ def _validate_step_cache_key(key: Any) -> str: return key +def _snapshot_step_cache_value(value: Any) -> Any: + try: + return deepcopy(value) + except Exception as exc: + raise ValueError( + "@step results must support deepcopy so cached replay cannot be changed by downstream mutation." + ) from exc + + @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) def get_run_context() -> RunContext | None: """Return the active :class:`RunContext`, or ``None`` if not inside a ``@workflow``. @@ -662,6 +691,27 @@ def _get_cached_result(self, key: str) -> tuple[bool, Any]: return True, self._step_cache[key] return False, None + def _has_incompatible_versioned_cache_entry(self, key: str) -> bool: + current = _decode_versioned_step_cache_key(key) + if current is None: + return False + current_kind, current_step, current_wrapper, _current_identity, current_occurrence = current + for cached_key in self._step_cache: + cached = _decode_versioned_step_cache_key(cached_key) + if cached is None: + continue + cached_kind, cached_step, cached_wrapper, _cached_identity, cached_occurrence = cached + if ( + cached_kind == current_kind + and cached_step == current_step + and cached_occurrence == current_occurrence + and cached_wrapper != current_wrapper + ): + live_wrappers = _step_wrapper_identity_registry.get((cached_step, cached_wrapper)) + if not live_wrappers: + return True + return False + def _set_cached_result(self, key: str, value: Any) -> None: self._step_cache[key] = value @@ -806,14 +856,14 @@ def __init__( self.name: str = name or func.__name__ self._signature = inspect.signature(func) self._wrapper_identity, self._wrapper_identity_is_durable = _get_step_wrapper_identity(func) + defaults_state, self._wrapper_defaults_are_durable = _get_step_wrapper_defaults_state(func) + defaults_encoded = json.dumps(defaults_state, sort_keys=True, separators=(",", ":")) + self._wrapper_defaults_identity = hashlib.sha256(defaults_encoded.encode("utf-8")).hexdigest() if _active_run_ctx.get() is not None: self._wrapper_identity_is_durable = False - self._wrapper_registry_key: tuple[str, str] | None = None - if self._wrapper_identity_is_durable: - registry_key = (self.name, self._wrapper_identity) - matching_wrappers = _step_wrapper_identity_registry.setdefault(registry_key, WeakSet()) - matching_wrappers.add(self) - self._wrapper_registry_key = registry_key + self._wrapper_registry_key = (self.name, self._wrapper_identity) + matching_wrappers = _step_wrapper_identity_registry.setdefault(self._wrapper_registry_key, WeakSet()) + matching_wrappers.add(self) self._replay_key = replay_key functools.update_wrapper(self, func) @@ -845,11 +895,13 @@ def _get_replay_identity( args: tuple[Any, ...], kwargs: dict[str, Any], ) -> tuple[Literal["auto", "explicit"], str] | None: - if self._wrapper_identity_is_durable: - current_identity, current_is_durable = _get_step_wrapper_identity(self._func) - if not current_is_durable or current_identity != self._wrapper_identity: + if self._wrapper_defaults_are_durable: + current_defaults, current_defaults_are_durable = _get_step_wrapper_defaults_state(self._func) + current_encoded = json.dumps(current_defaults, sort_keys=True, separators=(",", ":")) + current_identity = hashlib.sha256(current_encoded.encode("utf-8")).hexdigest() + if not current_defaults_are_durable or current_identity != self._wrapper_defaults_identity: raise ValueError( - f"@step '{self.name}' defaults or captured state changed after decoration. " + f"@step '{self.name}' defaults changed after decoration. " "Create a new workflow definition or provide a versioned replay_key." ) @@ -859,10 +911,7 @@ def _get_replay_identity( raise ValueError(f"@step '{self.name}' replay_key must return a non-empty string.") return ("explicit", hashlib.sha256(explicit_key.encode("utf-8")).hexdigest()) - if self._wrapper_registry_key is not None: - matching_wrappers = _step_wrapper_identity_registry.get(self._wrapper_registry_key) - else: - matching_wrappers = None + matching_wrappers = _step_wrapper_identity_registry.get(self._wrapper_registry_key) if matching_wrappers is not None and len(matching_wrappers) > 1: raise ValueError( f"@step '{self.name}' was generated from a function definition used by multiple wrappers. " @@ -885,8 +934,9 @@ def _get_replay_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) - await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, cached)) - return cached + replayed = _snapshot_step_cache_value(cached) + await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, _snapshot_step_cache_value(replayed))) + return replayed def _build_call_args_with_ctx( self, @@ -942,6 +992,11 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: found, cached = ctx._get_cached_result(cache_key) if found: return await self._return_cached_result(ctx, cache_key, cached) + if ctx._has_incompatible_versioned_cache_entry(cache_key): + raise ValueError( + f"Checkpoint cache for @step '{self.name}' was created by a different step definition. " + "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: @@ -986,8 +1041,8 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: cache_key, ctx._auto_request_info_index - auto_request_info_index_before, ) - ctx._set_cached_result(cache_key, result) - await ctx.add_event(WorkflowEvent.executor_completed(self.name, result)) + ctx._set_cached_result(cache_key, _snapshot_step_cache_value(result)) + await ctx.add_event(WorkflowEvent.executor_completed(self.name, _snapshot_step_cache_value(result))) if ctx._on_step_completed is not None: await ctx._on_step_completed() return result diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 7be955e9246..b103ca8fd5f 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -1395,6 +1395,102 @@ async def wf(value: str) -> str: assert initial.get_outputs() == ["input:A"] + async def test_changed_step_definition_rejects_versioned_checkpoint(self): + storage = InMemoryCheckpointStorage() + _factory_step_calls.clear() + + async def old_marker(value: str) -> str: + _factory_step_calls.append("old") + return f"old:{value}" + + async def new_marker(value: str) -> str: + _factory_step_calls.append("new") + return f"new:{value}" + + marker = step(name="marker")(old_marker) + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await marker(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + marker = step(name="marker")(new_marker) + gc.collect() + + with pytest.raises(ValueError, match="different step definition"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["old:input"] + assert _factory_step_calls == ["old"] + + async def test_rebound_default_change_rejects_versioned_checkpoint(self): + storage = InMemoryCheckpointStorage() + + async def old_marker(value: str, *, suffix: bytes = b"A") -> str: + return f"{value}:{suffix.decode()}" + + async def new_marker(value: str, *, suffix: bytes = b"B") -> str: + return f"{value}:{suffix.decode()}" + + marker = step(name="marker")(old_marker) + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await marker(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + marker = step(name="marker")(new_marker) + gc.collect() + + with pytest.raises(ValueError, match="different step definition"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["input:A"] + + async def test_mutable_step_result_isolated_across_response_replay(self): + _factory_step_calls.clear() + + @step + async def create_items() -> list[str]: + _factory_step_calls.append("create") + return [] + + @built_workflow + async def wf(_: str, ctx: RunContext) -> list[str]: + items = await create_items() + items.append("after") + await ctx.request_info("continue", response_type=str, request_id="continue") + return items + + interrupted = await wf.run("input") + assert interrupted.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + resumed = await wf.run(responses={"continue": "yes"}) + + assert resumed.get_outputs() == [["after"]] + assert _factory_step_calls == ["create"] + + async def test_completed_event_isolated_from_mutable_result(self): + @step + async def create_items() -> list[str]: + return [] + + @built_workflow + async def wf(_: str) -> list[str]: + items = await create_items() + items.append("after") + return items + + result = await wf.run("input") + completed = [event for event in result if event.type == "executor_completed"] + + assert result.get_outputs() == [["after"]] + assert len(completed) == 1 + assert completed[0].data == [] + async def test_legacy_checkpoint_replays_sequential_step(self): from agent_framework import WorkflowCheckpoint @@ -1996,7 +2092,11 @@ async def generated(value: int) -> str: second_closure = make("second") assert _get_step_wrapper_identity(first_default)[0] != _get_step_wrapper_identity(second_default)[0] - assert _get_step_wrapper_identity(first_closure)[0] != _get_step_wrapper_identity(second_closure)[0] + first_closure_identity, first_closure_is_durable = _get_step_wrapper_identity(first_closure) + second_closure_identity, second_closure_is_durable = _get_step_wrapper_identity(second_closure) + assert first_closure_is_durable is True + assert second_closure_is_durable is True + assert first_closure_identity != second_closure_identity def test_wrapper_identity_marks_mutable_closure_state_non_durable(self): state = ["A"] From 0e02f754900bd067e79e98a716cd27778e2f505e Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 30 Sep 2026 11:33:49 +0200 Subject: [PATCH 6/7] Python: close replay durability gaps --- .../packages/core/agent_framework/_workflows/_functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 9ce7e4cdec9..88cfe76d29f 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -47,7 +47,7 @@ import logging import math import typing -from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence from contextvars import Context, ContextVar from copy import deepcopy from types import CodeType @@ -99,7 +99,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None: def _create_task( self, loop: asyncio.AbstractEventLoop, - coro: Coroutine[Any, Any, Any], + coro: Any, **kwargs: Any, ) -> asyncio.Future[Any]: if self.original_factory is None: From c4d4685d47647883a298f81bdc7333fa7288bc95 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 1 Oct 2026 17:55:42 +0200 Subject: [PATCH 7/7] Python: scope replay compatibility to workflows --- .../agent_framework/_workflows/_functional.py | 289 +++++++++++++----- .../workflow/test_functional_workflow.py | 289 ++++++++++++++++-- 2 files changed, 469 insertions(+), 109 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 88cfe76d29f..1920b499094 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -40,6 +40,7 @@ # Classes in this module (RunContext, StepWrapper, FunctionalWorkflow) form a # cohesive unit and intentionally access each other's underscore-prefixed members. import asyncio +import dis import functools import hashlib import inspect @@ -47,12 +48,11 @@ import logging import math import typing -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence from contextvars import Context, ContextVar from copy import deepcopy from types import CodeType from typing import Any, Generic, Literal, TypeVar, overload -from weakref import WeakSet from .._agents import BaseAgent from .._feature_stage import ExperimentalFeature, experimental @@ -79,7 +79,6 @@ _STEP_CACHE_KEY_V2_PREFIX = "v2::" _UNSUPPORTED_STEP_IDENTITY = object() -_step_wrapper_identity_registry: dict[tuple[str, str], WeakSet[Any]] = {} # ContextVar holding the active RunContext during workflow execution. # ContextVar is per-asyncio-Task, so concurrent workflows each get their own context. @@ -288,7 +287,7 @@ def _is_immutable_wrapper_state(value: Any) -> bool: return False -def _get_step_wrapper_defaults_state(func: Callable[..., Awaitable[Any]]) -> tuple[dict[str, Any], bool]: +def _get_step_wrapper_defaults_state(func: Callable[..., Any]) -> tuple[dict[str, Any], bool]: defaults, defaults_are_durable = _canonicalize_wrapper_state(func.__defaults__ or ()) raw_kwdefaults = func.__kwdefaults__ or {} kwdefaults, _ = _canonicalize_wrapper_state(raw_kwdefaults) @@ -302,11 +301,23 @@ def _get_step_wrapper_defaults_state(func: Callable[..., Awaitable[Any]]) -> tup ) -def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str, bool]: +def _get_step_wrapper_source_identity(func: Callable[..., Any]) -> str: + code = getattr(func, "__code__", None) + payload = { + "module": func.__module__, + "qualname": func.__qualname__, + "firstlineno": code.co_firstlineno if code is not None else None, + "code": _get_code_identity_payload(code) if code is not None else None, + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def _get_step_wrapper_identity(func: Callable[..., Any]) -> tuple[str, bool]: code = getattr(func, "__code__", None) defaults_state, defaults_are_durable = _get_step_wrapper_defaults_state(func) closure_items: list[Any] = [] - closure_is_durable = True + closure_is_durable = func.__closure__ is None if code is not None and func.__closure__ is not None: for name, cell in zip(code.co_freevars, func.__closure__): try: @@ -320,10 +331,7 @@ def _get_step_wrapper_identity(func: Callable[..., Awaitable[Any]]) -> tuple[str closure_is_durable = closure_is_durable and is_durable payload = { - "module": func.__module__, - "qualname": func.__qualname__, - "firstlineno": code.co_firstlineno if code is not None else None, - "code": _get_code_identity_payload(code) if code is not None else None, + "source": _get_step_wrapper_source_identity(func), **defaults_state, "closure": closure_items, } @@ -388,13 +396,18 @@ def _validate_step_cache_key(key: Any) -> str: return key -def _snapshot_step_cache_value(value: Any) -> Any: +def _try_snapshot_step_cache_value(value: Any) -> tuple[bool, Any]: try: - return deepcopy(value) - except Exception as exc: - raise ValueError( - "@step results must support deepcopy so cached replay cannot be changed by downstream mutation." - ) from exc + return (True, deepcopy(value)) + except Exception: + return (False, value) + + +def _snapshot_replayed_step_cache_value(value: Any) -> Any: + copied, snapshot = _try_snapshot_step_cache_value(value) + if copied: + return snapshot + raise ValueError("@step result cannot be safely copied for replay.") @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) @@ -498,8 +511,12 @@ def __init__( # Step result cache. Keys are already checkpoint-safe strings so result # and metadata maps cannot drift during serialization. self._step_cache: dict[str, Any] = {} + # Keys whose live values could not be copied and therefore cannot be replayed. + self._non_replayable_step_cache_keys: set[str] = set() # Cached step metadata used to keep auto-generated request_info IDs in sync on bypass. 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. self._step_call_counters: dict[str, int] = {} # Per-identity counters keep repeated automatic invocations distinct. @@ -688,6 +705,8 @@ def _is_concurrent_execution(self) -> bool: def _get_cached_result(self, key: str) -> tuple[bool, Any]: if key in self._step_cache: + if key in self._non_replayable_step_cache_keys: + raise ValueError(f"Cached result for @step key {key!r} cannot be safely replayed.") return True, self._step_cache[key] return False, None @@ -707,13 +726,17 @@ def _has_incompatible_versioned_cache_entry(self, key: str) -> bool: and cached_occurrence == current_occurrence and cached_wrapper != current_wrapper ): - live_wrappers = _step_wrapper_identity_registry.get((cached_step, cached_wrapper)) - if not live_wrappers: + expected_wrappers = self._expected_step_wrapper_identities.get(cached_step, set()) + if cached_wrapper not in expected_wrappers: return True return False - def _set_cached_result(self, key: str, value: Any) -> None: + def _set_cached_result(self, key: str, value: Any, *, replayable: bool = True) -> None: self._step_cache[key] = value + if replayable: + self._non_replayable_step_cache_keys.discard(key) + else: + self._non_replayable_step_cache_keys.add(key) def _set_cached_step_auto_request_info_count(self, key: str, count: int) -> None: self._step_cache_auto_request_info_counts[key] = count @@ -751,6 +774,11 @@ def _get_response(self, request_id: str) -> tuple[bool, Any]: def _export_step_cache(self) -> dict[str, Any]: """Serialize the step cache for checkpointing.""" + 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." + ) return dict(self._step_cache) def _export_step_cache_auto_request_info_counts(self) -> dict[str, int]: @@ -760,6 +788,7 @@ def _export_step_cache_auto_request_info_counts(self) -> dict[str, int]: def _import_step_cache(self, data: dict[str, Any]) -> None: """Restore step cache from checkpoint data.""" self._step_cache = {} + self._non_replayable_step_cache_keys = set() for k, v in data.items(): try: self._step_cache[_validate_step_cache_key(k)] = v @@ -821,9 +850,9 @@ 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 captured defaults/closure state, and multiple wrappers whose - canonical source and captured state are identical, require ``replay_key`` - for concurrent use. + 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. @@ -855,15 +884,13 @@ def __init__( self._func = func self.name: str = name or func.__name__ self._signature = inspect.signature(func) + self._wrapper_source_identity = _get_step_wrapper_source_identity(func) self._wrapper_identity, self._wrapper_identity_is_durable = _get_step_wrapper_identity(func) defaults_state, self._wrapper_defaults_are_durable = _get_step_wrapper_defaults_state(func) defaults_encoded = json.dumps(defaults_state, sort_keys=True, separators=(",", ":")) self._wrapper_defaults_identity = hashlib.sha256(defaults_encoded.encode("utf-8")).hexdigest() if _active_run_ctx.get() is not None: self._wrapper_identity_is_durable = False - self._wrapper_registry_key = (self.name, self._wrapper_identity) - matching_wrappers = _step_wrapper_identity_registry.setdefault(self._wrapper_registry_key, WeakSet()) - matching_wrappers.add(self) self._replay_key = replay_key functools.update_wrapper(self, func) @@ -911,13 +938,6 @@ def _get_replay_identity( raise ValueError(f"@step '{self.name}' replay_key must return a non-empty string.") return ("explicit", hashlib.sha256(explicit_key.encode("utf-8")).hexdigest()) - matching_wrappers = _step_wrapper_identity_registry.get(self._wrapper_registry_key) - if matching_wrappers is not None and len(matching_wrappers) > 1: - raise ValueError( - f"@step '{self.name}' was generated from a function definition used by multiple wrappers. " - "Provide replay_key so each logical invocation has an explicit durable identity." - ) - if not self._wrapper_identity_is_durable: return None @@ -934,8 +954,9 @@ def _get_replay_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) - replayed = _snapshot_step_cache_value(cached) - await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, _snapshot_step_cache_value(replayed))) + replayed = _snapshot_replayed_step_cache_value(cached) + _, event_data = _try_snapshot_step_cache_value(replayed) + await ctx.add_event(WorkflowEvent.executor_bypassed(self.name, event_data)) return replayed def _build_call_args_with_ctx( @@ -985,7 +1006,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: if replay_identity is not None: identity_kind, identity = replay_identity if identity_kind == "explicit": - cache_key = ctx._get_explicit_step_cache_key(self.name, self._wrapper_identity, identity) + cache_key = ctx._get_explicit_step_cache_key(self.name, self._wrapper_source_identity, identity) else: cache_key = ctx._get_automatic_step_cache_key(self.name, self._wrapper_identity, identity) @@ -1041,8 +1062,11 @@ async def __call__(self, *args: Any, **kwargs: Any) -> R: cache_key, ctx._auto_request_info_index - auto_request_info_index_before, ) - ctx._set_cached_result(cache_key, _snapshot_step_cache_value(result)) - await ctx.add_event(WorkflowEvent.executor_completed(self.name, _snapshot_step_cache_value(result))) + copied, cached_result = _try_snapshot_step_cache_value(result) + identity_is_replayable = self._replay_key is not None or self._wrapper_identity_is_durable + 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)) if ctx._on_step_completed is not None: await ctx._on_step_completed() return result @@ -1092,7 +1116,8 @@ 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. Legacy order-based checkpoints replay sequential + 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. @@ -1254,13 +1279,11 @@ def __init__( self._last_step_cache_auto_request_info_counts: dict[str, int] = {} self._last_state: dict[str, Any] = {} self._last_pending_request_ids: set[str] = set() + self._last_graph_signature_hash: str | None = None # Signature arity is validated once at decoration time. self._non_ctx_param_names = self._classify_signature(func) - # Discover step names referenced in the function for signature hash - self._step_names = self._discover_step_names(func) - # Compute a stable signature hash self.graph_signature_hash = self._compute_signature_hash() @@ -1277,14 +1300,20 @@ def _capture_replay_state(self, ctx: RunContext, message: Any | None = None) -> """Capture the state needed to continue a response-only HITL replay.""" if message is not None and self._last_message is None: self._last_message = self._snapshot_replay_message(message) - self._last_step_cache = dict(ctx._step_cache) + self._last_step_cache = ctx._export_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._last_graph_signature_hash = self.graph_signature_hash def _restore_replay_state(self, ctx: RunContext) -> Any: """Restore cached execution state and return the message used by the replay.""" + if self._last_graph_signature_hash is not None and self._last_graph_signature_hash != self.graph_signature_hash: + raise ValueError( + f"Response-only replay state for workflow '{self.name}' was created by a different workflow version." + ) ctx._step_cache = dict(self._last_step_cache) + ctx._non_replayable_step_cache_keys = set() 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 @@ -1296,6 +1325,7 @@ def _clear_replay_state(self) -> None: self._last_step_cache_auto_request_info_counts = {} self._last_state = {} self._last_pending_request_ids = set() + self._last_graph_signature_hash = None @staticmethod def _classify_signature(func: Callable[..., Any]) -> list[str]: @@ -1521,9 +1551,18 @@ async def _run_core( **kwargs: Any, ) -> AsyncIterable[WorkflowEvent[Any]]: storage = checkpoint_storage or self._checkpoint_storage + dependencies = self._discover_workflow_dependencies(self._func) + step_wrappers, _ = dependencies + self.graph_signature_hash = self._compute_signature_hash(dependencies) # Build context ctx = RunContext(self.name, streaming=streaming, run_kwargs=kwargs if kwargs else None) + for wrapper in step_wrappers: + if wrapper._replay_key is not None: + wrapper_identity = wrapper._wrapper_source_identity + else: + wrapper_identity, _ = _get_step_wrapper_identity(wrapper._func) + ctx._expected_step_wrapper_identities.setdefault(wrapper.name, set()).add(wrapper_identity) # Restore from checkpoint if requested prev_checkpoint_id: str | None = None @@ -1605,9 +1644,6 @@ async def _on_step_completed() -> None: _framework_event(WorkflowEvent, "output", executor_id=self.name, data=return_value) ) - # Persist step cache for response-only replay - self._capture_replay_state(ctx, message) - # Yield collected events. # NOTE: Events are buffered during _execute() and yielded after # the user function completes. This is *not* true streaming — @@ -1755,44 +1791,143 @@ async def _save_checkpoint( ) return await storage.save(checkpoint) - def _compute_signature_hash(self) -> str: - """Stable hash of the workflow's code shape. - - Mixes workflow name, statically-discovered step names, and a digest - of ``__code__.co_code`` + ``co_names``. The code digest catches - body changes that step-name discovery misses (e.g. attribute-access - step references). - """ - code = getattr(self._func, "__code__", None) - co_code_hex = hashlib.sha256(code.co_code).hexdigest() if code is not None else "" - co_names = tuple(sorted(code.co_names)) if code is not None else () + def _compute_signature_hash( + self, + dependencies: tuple[set[StepWrapper[Any]], set[tuple[str, str, str]]] | None = None, + ) -> str: + """Build a stable signature from the workflow code and reachable step definitions.""" + discovered, function_signatures = ( + dependencies if dependencies is not None else self._discover_workflow_dependencies(self._func) + ) + step_signatures: list[tuple[str, str, str]] = [] + for wrapper in discovered: + if wrapper._replay_key is not None: + defaults_state, _ = _get_step_wrapper_defaults_state(wrapper._func) + defaults_encoded = json.dumps(defaults_state, sort_keys=True, separators=(",", ":")) + identity = hashlib.sha256(defaults_encoded.encode("utf-8")).hexdigest() + step_signatures.append((wrapper.name, wrapper._wrapper_source_identity, identity)) + else: + identity, _ = _get_step_wrapper_identity(wrapper._func) + step_signatures.append((wrapper.name, "automatic", identity)) sig_data = { "workflow": self.name, - "steps": sorted(self._step_names), - "co_code": co_code_hex, - "co_names": list(co_names), + "steps": sorted(step_signatures), + "functions": sorted(function_signatures), } canonical = json.dumps(sig_data, sort_keys=True, separators=(",", ":")) return hashlib.sha256(canonical.encode("utf-8")).hexdigest() - @staticmethod - def _discover_step_names(func: Callable[..., Any]) -> list[str]: - """Extract step names referenced by the workflow function. - - Inspects the function's ``__code__.co_names`` and global scope for - ``StepWrapper`` instances. Steps accessed via module or class - attributes (``my_steps.fetch``) are missed here, but - :meth:`_compute_signature_hash` still captures them through the - ``co_code`` digest. - """ - names: list[str] = [] - globs = getattr(func, "__globals__", {}) - code_names = getattr(getattr(func, "__code__", None), "co_names", ()) - for n in code_names: - obj = globs.get(n) - if isinstance(obj, StepWrapper): - names.append(obj.name) - return names + @classmethod + def _discover_workflow_dependencies( + cls, + func: Callable[..., Any], + ) -> tuple[set[StepWrapper[Any]], set[tuple[str, str, str]]]: + """Find step wrappers and helper identities reachable from the workflow.""" + wrappers: set[StepWrapper[Any]] = set() + functions: set[tuple[str, str, str]] = set() + visited: set[int] = set() + recorded_functions: set[int] = set() + visited_code_contexts: set[tuple[int, int, str | None]] = set() + workflow_module = getattr(func, "__module__", None) + + def visit_code( + code: CodeType, + globals_dict: Mapping[str, Any], + *, + allowed_module: str | None, + ) -> None: + context_key = (id(code), id(globals_dict), allowed_module) + if context_key in visited_code_contexts: + return + visited_code_contexts.add(context_key) + + global_names = { + instruction.argval + for instruction in dis.get_instructions(code) + if instruction.opname in {"LOAD_GLOBAL", "LOAD_NAME"} and isinstance(instruction.argval, str) + } + for name in global_names: + if name in globals_dict: + reference = globals_dict[name] + reference_module = ( + getattr(reference, "__module__", allowed_module) + if inspect.isfunction(reference) + else allowed_module + ) + visit(reference, allowed_module=reference_module) + for constant in code.co_consts: + if isinstance(constant, CodeType): + visit_code(constant, globals_dict, allowed_module=allowed_module) + + def visit( + value: Any, + *, + record_function: bool = True, + allowed_module: str | None = workflow_module, + ) -> None: + value_id = id(value) + if inspect.isfunction(value) and value is not func and getattr(value, "__module__", None) != allowed_module: + return + if inspect.isfunction(value) and record_function and value_id not in recorded_functions: + function_identity, _ = _get_step_wrapper_identity(value) + functions.add(( + getattr(value, "__module__", ""), + getattr(value, "__qualname__", getattr(value, "__name__", "")), + function_identity, + )) + recorded_functions.add(value_id) + if value_id in visited: + return + visited.add(value_id) + + if isinstance(value, StepWrapper): + wrapper = typing.cast(StepWrapper[Any], value) + wrappers.add(wrapper) + visit( + wrapper._func, + record_function=False, + allowed_module=getattr(wrapper._func, "__module__", None), + ) + return + if isinstance(value, (staticmethod, classmethod)): + descriptor = typing.cast(Any, value) + visit(descriptor.__func__, allowed_module=allowed_module) + return + if inspect.ismethod(value): + visit(value.__func__, allowed_module=allowed_module) + return + if inspect.isfunction(value): + code = getattr(value, "__code__", None) + if code is None: + return + references: list[Any] = [] + closure = getattr(value, "__closure__", None) + if closure is not None: + for cell in closure: + try: + references.append(cell.cell_contents) + except ValueError: + continue + globals_dict = getattr(value, "__globals__", {}) + for reference in references: + reference_module = ( + getattr(reference, "__module__", allowed_module) + if inspect.isfunction(reference) + else allowed_module + ) + visit(reference, allowed_module=reference_module) + visit_code(code, globals_dict, allowed_module=allowed_module) + return + if isinstance(value, Mapping): + for item in typing.cast(Mapping[Any, Any], value).values(): + visit(item, allowed_module=allowed_module) + return + if isinstance(value, (list, tuple, set, frozenset)): + for item in typing.cast(Iterable[Any], value): + visit(item, allowed_module=allowed_module) + + visit(func) + return wrappers, functions # ------------------------------------------------------------------ # Finalize / cleanup / validation (mirrors Workflow) diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index b103ca8fd5f..de9ec51f2d1 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -759,7 +759,7 @@ async def test_checkpoint_restore_replays_cached_tasks(self): storage = InMemoryCheckpointStorage() call_count = 0 - @step + @step(replay_key=lambda x: str(x)) async def counting_task(x: int) -> int: nonlocal call_count call_count += 1 @@ -840,13 +840,13 @@ async def test_per_step_checkpoint_enables_crash_recovery(self): step1_calls = 0 step2_calls = 0 - @step + @step(replay_key=lambda x: str(x)) async def slow_step1(x: int) -> int: nonlocal step1_calls step1_calls += 1 return x + 10 - @step + @step(replay_key=lambda x: str(x)) async def crashing_step2(x: int) -> int: nonlocal step2_calls step2_calls += 1 @@ -1110,8 +1110,9 @@ async def wf(value: str) -> list[str]: assert replayed.get_outputs() == initial.get_outputs() assert calls == ["first", "second"] - async def test_existing_generated_wrapper_fails_closed_after_collision(self): - storage = InMemoryCheckpointStorage() + async def test_identical_generated_wrappers_do_not_cross_workflows(self): + first_storage = InMemoryCheckpointStorage() + second_storage = InMemoryCheckpointStorage() _factory_step_calls.clear() def make_step() -> StepWrapper[str]: @@ -1122,26 +1123,26 @@ async def generated(value: str) -> str: return step(name="generated")(generated) first = make_step() + second = make_step() - @built_workflow(checkpoint_storage=storage) - async def wf(value: str) -> str: + @built_workflow(name="first_wf", checkpoint_storage=first_storage) + async def first_wf(value: str) -> str: return await first(value) - initial = await wf.run("input") - checkpoints = await storage.list_checkpoints(workflow_name="wf") - checkpoint = checkpoints[-1] - second = make_step() - - with pytest.raises(ValueError, match="multiple wrappers"): - await wf.run(checkpoint_id=checkpoint.checkpoint_id) + @built_workflow(name="second_wf", checkpoint_storage=second_storage) + async def second_wf(value: str) -> str: + return await second(value) - del second - gc.collect() - replayed = await wf.run(checkpoint_id=checkpoint.checkpoint_id) + first_initial = await first_wf.run("first") + second_initial = await second_wf.run("second") + first_checkpoints = await first_storage.list_checkpoints(workflow_name="first_wf") + second_checkpoints = await second_storage.list_checkpoints(workflow_name="second_wf") + first_replayed = await first_wf.run(checkpoint_id=first_checkpoints[-1].checkpoint_id) + second_replayed = await second_wf.run(checkpoint_id=second_checkpoints[-1].checkpoint_id) - assert initial.get_outputs() == ["input"] - assert replayed.get_outputs() == initial.get_outputs() - assert _factory_step_calls == ["input"] + assert first_replayed.get_outputs() == first_initial.get_outputs() == ["first"] + assert second_replayed.get_outputs() == second_initial.get_outputs() == ["second"] + assert _factory_step_calls == ["first", "second"] async def test_cyclic_arguments_use_sequential_fallback(self): call_count = 0 @@ -1390,7 +1391,7 @@ async def wf(value: str) -> str: assert marker_func.__kwdefaults__ is not None marker_func.__kwdefaults__["suffix"] = b"B" - with pytest.raises(ValueError, match="changed after decoration"): + with pytest.raises(ValueError, match="not compatible"): await wf.run(checkpoint_id=checkpoint.checkpoint_id) assert initial.get_outputs() == ["input:A"] @@ -1416,14 +1417,15 @@ async def wf(value: str) -> str: initial = await wf.run("input") checkpoints = await storage.list_checkpoints(workflow_name="wf") checkpoint = checkpoints[-1] + old_alias = marker marker = step(name="marker")(new_marker) - gc.collect() - with pytest.raises(ValueError, match="different step definition"): + with pytest.raises(ValueError, match="not compatible"): await wf.run(checkpoint_id=checkpoint.checkpoint_id) assert initial.get_outputs() == ["old:input"] assert _factory_step_calls == ["old"] + assert old_alias.name == "marker" async def test_rebound_default_change_rejects_versioned_checkpoint(self): storage = InMemoryCheckpointStorage() @@ -1443,13 +1445,177 @@ async def wf(value: str) -> str: initial = await wf.run("input") checkpoints = await storage.list_checkpoints(workflow_name="wf") checkpoint = checkpoints[-1] + old_alias = marker marker = step(name="marker")(new_marker) - gc.collect() - with pytest.raises(ValueError, match="different step definition"): + with pytest.raises(ValueError, match="not compatible"): await wf.run(checkpoint_id=checkpoint.checkpoint_id) assert initial.get_outputs() == ["input:A"] + assert old_alias.name == "marker" + + async def test_rebound_helper_body_rejects_checkpoint(self): + storage = InMemoryCheckpointStorage() + + @step + async def marker(value: str) -> str: + return value + + async def old_helper(value: str) -> str: + return await marker(value) + + async def new_helper(value: str) -> str: + return f"{await marker(value)}:changed" + + helper = old_helper + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await helper(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + old_alias = helper + helper = new_helper + + with pytest.raises(ValueError, match="not compatible"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["input"] + assert old_alias.__name__ == "old_helper" + + async def test_rebound_imported_workflow_helper_rejects_checkpoint(self): + storage = InMemoryCheckpointStorage() + + @step + async def marker(value: str) -> str: + return f"marker:{value}" + + async def old_helper(value: str) -> str: + return await marker(value) + + async def new_helper(value: str) -> str: + return f"changed:{await marker(value)}" + + old_helper.__module__ = "external_helpers" + new_helper.__module__ = "external_helpers" + helper = old_helper + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await helper(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + old_alias = helper + helper = new_helper + + with pytest.raises(ValueError, match="not compatible"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["marker:input"] + assert old_alias.__name__ == "old_helper" + + async def test_rebound_step_helper_body_rejects_checkpoint(self): + storage = InMemoryCheckpointStorage() + + async def old_helper(value: str) -> str: + return f"old:{value}" + + async def new_helper(value: str) -> str: + return f"new:{value}" + + helper = old_helper + + @step(replay_key=lambda value: value) + async def marker(value: str) -> str: + return await helper(value) + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await marker(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + old_alias = helper + helper = new_helper + + with pytest.raises(ValueError, match="not compatible"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["old:input"] + assert old_alias.__name__ == "old_helper" + + async def test_rebound_imported_step_helper_rejects_checkpoint(self): + storage = InMemoryCheckpointStorage() + + async def old_helper(value: str) -> str: + return f"old:{value}" + + async def new_helper(value: str) -> str: + return f"new:{value}" + + old_helper.__module__ = "external_steps" + new_helper.__module__ = "external_steps" + helper = old_helper + + async def marker_func(value: str) -> str: + return await helper(value) + + marker_func.__module__ = "external_steps" + imported_marker = step(replay_key=lambda value: value)(marker_func) + + @built_workflow(checkpoint_storage=storage) + async def wf(value: str) -> str: + return await imported_marker(value) + + initial = await wf.run("input") + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + old_alias = helper + helper = new_helper + + with pytest.raises(ValueError, match="not compatible"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + + assert initial.get_outputs() == ["old:input"] + assert old_alias.__name__ == "old_helper" + + async def test_nested_helper_rebound_global_step_rejects_checkpoint(self): + storage = InMemoryCheckpointStorage() + + @step(name="old") + async def old_step(value: int) -> int: + return value + 1 + + @step(name="new") + async def new_step(value: int) -> int: + return value + 100 + + globals()["nested_current_step"] = old_step + try: + + @built_workflow(checkpoint_storage=storage) + async def wf(value: int) -> int: + async def nested_helper() -> int: + return await nested_current_step(value) # type: ignore[name-defined] # pyright: ignore[reportUndefinedVariable] # ty: ignore[unresolved-reference] # noqa: F821 + + return await nested_helper() + + initial = await wf.run(1) + checkpoints = await storage.list_checkpoints(workflow_name="wf") + checkpoint = checkpoints[-1] + globals()["nested_current_step"] = new_step + + with pytest.raises(ValueError, match="not compatible"): + await wf.run(checkpoint_id=checkpoint.checkpoint_id) + finally: + del globals()["nested_current_step"] + + assert initial.get_outputs() == [2] async def test_mutable_step_result_isolated_across_response_replay(self): _factory_step_calls.clear() @@ -1491,6 +1657,39 @@ async def wf(_: str) -> list[str]: assert len(completed) == 1 assert completed[0].data == [] + async def test_non_deepcopyable_result_succeeds_without_replay(self): + import threading + + @step + async def create_lock() -> threading.Lock: + return threading.Lock() + + @built_workflow + async def wf(_: str) -> threading.Lock: + return await create_lock() + + result = await wf.run("input") + lock = result.get_outputs()[0] + + assert lock.acquire(blocking=False) is True + lock.release() + + async def test_non_deepcopyable_result_rejects_checkpointing(self): + import threading + + storage = InMemoryCheckpointStorage() + + @step + async def create_lock() -> threading.Lock: + return threading.Lock() + + @built_workflow(checkpoint_storage=storage) + async def wf(_: str) -> threading.Lock: + return await create_lock() + + with pytest.raises(ValueError, match="Cannot checkpoint"): + await wf.run("input") + async def test_legacy_checkpoint_replays_sequential_step(self): from agent_framework import WorkflowCheckpoint @@ -1687,7 +1886,7 @@ async def wf(_: str) -> str: key_ctx = _RunContext("mixed_wf") versioned_key = key_ctx._get_explicit_step_cache_key( # pyright: ignore[reportPrivateUsage] "mixed", - mixed._wrapper_identity, # pyright: ignore[reportPrivateUsage] + mixed._wrapper_source_identity, # pyright: ignore[reportPrivateUsage] identity, ) checkpoint = WorkflowCheckpoint( @@ -2094,8 +2293,8 @@ async def generated(value: int) -> str: assert _get_step_wrapper_identity(first_default)[0] != _get_step_wrapper_identity(second_default)[0] first_closure_identity, first_closure_is_durable = _get_step_wrapper_identity(first_closure) second_closure_identity, second_closure_is_durable = _get_step_wrapper_identity(second_closure) - assert first_closure_is_durable is True - assert second_closure_is_durable is True + assert first_closure_is_durable is False + assert second_closure_is_durable is False assert first_closure_identity != second_closure_identity def test_wrapper_identity_marks_mutable_closure_state_non_durable(self): @@ -2108,6 +2307,32 @@ async def generated(value: str) -> str: assert is_durable is False + def test_attribute_name_does_not_discover_same_named_global_helper(self): + @step + async def unrelated_step(value: str) -> str: + return value + + async def unrelated_helper(value: str) -> str: + return await unrelated_step(value) + + class Client: + async def signature_probe(self, value: str) -> str: + return value + + client = Client() + globals()["signature_probe"] = unrelated_helper + try: + + @built_workflow + async def wf(value: str) -> str: + return await client.signature_probe(value) + + wrappers, _ = wf._discover_workflow_dependencies(wf._func) # pyright: ignore[reportPrivateUsage] + finally: + del globals()["signature_probe"] + + assert unrelated_step not in wrappers + async def test_checkpoint_signature_mismatch_raises(self): from agent_framework import WorkflowCheckpoint @@ -2182,7 +2407,7 @@ async def test_cached_step_emits_bypassed_event(self): storage = InMemoryCheckpointStorage() call_count = 0 - @step + @step(replay_key=lambda x: str(x)) async def tracked(x: int) -> int: nonlocal call_count call_count += 1 @@ -2427,7 +2652,7 @@ 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 + @step(replay_key=lambda: "seed") async def seed_state(ctx: RunContext) -> str: nonlocal seed_calls seed_calls += 1 @@ -2453,7 +2678,7 @@ async def test_preceding_step_bypassed_on_hitl_resume(self): resuming should bypass the first step (cached) and re-execute the HITL step.""" call_count_a = 0 - @step + @step(replay_key=lambda x: str(x)) async def step_a(x: int) -> int: nonlocal call_count_a call_count_a += 1 @@ -2600,7 +2825,7 @@ async def wf(x: int, ctx: RunContext) -> str: async def test_cached_step_advances_auto_request_id_counter(self): call_count = 0 - @step + @step(replay_key=lambda value: str(value)) async def first_review(value: int, ctx: RunContext) -> str: nonlocal call_count call_count += 1