From 749701c8714f3ead1ac43e7faa30c46349560047 Mon Sep 17 00:00:00 2001 From: Joao Moura Date: Mon, 13 Jul 2026 21:52:38 -0700 Subject: [PATCH] [EPD-196] Add strict mode + machine-readable outcome for restore_from_state_id misses - Add opt-in raise_on_missing_state flag to kickoff, kickoff_async, akickoff, stream_events, and astream: when True and the restore_from_state_id lookup misses (or no persistence backend is configured), kickoff raises a clear ValueError instead of silently running a fresh session. - Expose flow.last_restore_succeeded (None = no restore requested, True = hydrated, False = miss) so API wrappers can detect silent restore misses without parsing logs. - Default behavior is unchanged: misses still fall back silently. Co-Authored-By: Claude Fable 5 --- .../en/guides/flows/mastering-flow-state.mdx | 2 +- .../src/crewai/flow/runtime/__init__.py | 65 ++++++- lib/crewai/tests/test_flow_persistence.py | 171 ++++++++++++++++++ 3 files changed, 233 insertions(+), 5 deletions(-) diff --git a/docs/edge/en/guides/flows/mastering-flow-state.mdx b/docs/edge/en/guides/flows/mastering-flow-state.mdx index 648a82dbd..fed8f6acd 100644 --- a/docs/edge/en/guides/flows/mastering-flow-state.mdx +++ b/docs/edge/en/guides/flows/mastering-flow-state.mdx @@ -385,7 +385,7 @@ flow_2.kickoff(restore_from_state_id=flow_1.state.id) Behavior notes: -- `restore_from_state_id` not found in persistence → the kickoff falls back silently to default behavior (mirrors the existing `inputs["id"]` resume not-found behavior). No exception is raised. +- `restore_from_state_id` not found in persistence → the kickoff falls back silently to default behavior (mirrors the existing `inputs["id"]` resume not-found behavior). No exception is raised by default. Pass `raise_on_missing_state=True` to get a `ValueError` instead of a silent fresh run, or check `flow.last_restore_succeeded` after kickoff (`None` = no restore requested, `True` = hydrated, `False` = miss). - Combining `restore_from_state_id` with `from_checkpoint` raises a `ValueError` — they target different state systems (`@persist` vs. Checkpointing) and cannot be combined. - `restore_from_state_id=None` (default) is byte-identical to a kickoff without the parameter. - Pinning `inputs["id"]` while forking means the new run shares a persistence key with another flow — usually you want only `restore_from_state_id`. diff --git a/lib/crewai/src/crewai/flow/runtime/__init__.py b/lib/crewai/src/crewai/flow/runtime/__init__.py index a6c73e697..7f74412b5 100644 --- a/lib/crewai/src/crewai/flow/runtime/__init__.py +++ b/lib/crewai/src/crewai/flow/runtime/__init__.py @@ -730,6 +730,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): _usage_aggregation_handler: Callable[..., Any] | None = PrivateAttr(default=None) _persist_backends: dict[int, FlowPersistence] = PrivateAttr(default_factory=dict) _instance_persistence: bool = PrivateAttr(default=False) + _last_restore_succeeded: bool | None = PrivateAttr(default=None) def __class_getitem__(cls: type[Flow[T]], item: type[T]) -> type[Flow[T]]: # type: ignore[override] class _FlowGeneric(cls): # type: ignore[valid-type,misc] @@ -1660,6 +1661,18 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): def state(self) -> T: return cast(T, self._state) + @property + def last_restore_succeeded(self) -> bool | None: + """Outcome of the most recent ``restore_from_state_id`` request. + + ``None`` when the latest kickoff did not request a restore, ``True`` + when the referenced state was found and hydrated, and ``False`` when + the lookup missed (or no persistence backend was configured) and the + run proceeded without hydration. Lets callers detect silent restore + misses without parsing logs. + """ + return self._last_restore_succeeded + @property def method_outputs(self) -> list[Any]: """Returns the list of all outputs from executed methods.""" @@ -1857,6 +1870,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files: dict[str, FileInput] | None = None, from_checkpoint: CheckpointConfig | None = None, restore_from_state_id: str | None = None, + raise_on_missing_state: bool = False, ) -> StreamSession[Any]: """Run the flow and stream all scoped public ``StreamFrame`` events.""" result_holder: list[Any] = [] @@ -1872,6 +1886,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files=input_files, from_checkpoint=from_checkpoint, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) except HumanFeedbackPending as e: return e @@ -1890,6 +1905,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files: dict[str, FileInput] | None = None, from_checkpoint: CheckpointConfig | None = None, restore_from_state_id: str | None = None, + raise_on_missing_state: bool = False, ) -> AsyncStreamSession[Any]: """Run the flow asynchronously and stream scoped public frames.""" result_holder: list[Any] = [] @@ -1905,6 +1921,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files=input_files, from_checkpoint=from_checkpoint, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) except HumanFeedbackPending as e: return e @@ -1923,6 +1940,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files: dict[str, FileInput] | None = None, from_checkpoint: CheckpointConfig | None = None, restore_from_state_id: str | None = None, + raise_on_missing_state: bool = False, ) -> Any | StreamSession[Any]: """Start the flow execution in a synchronous context. @@ -1940,8 +1958,14 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): pinned), so its ``@persist`` writes land under a separate persistence key and the source flow's history is preserved. If the referenced state is not found, the kickoff falls back - silently to baseline behavior. Cannot be combined with - ``from_checkpoint``; passing both raises ``ValueError``. + silently to baseline behavior (see ``raise_on_missing_state`` + to opt out). Cannot be combined with ``from_checkpoint``; + passing both raises ``ValueError``. + raise_on_missing_state: When True and ``restore_from_state_id`` is + provided but the referenced state cannot be loaded, raise + ``ValueError`` instead of silently running without hydration. + Defaults to False. ``last_restore_succeeded`` reports the + restore outcome either way. Returns: The final output from the flow or StreamSession if streaming. @@ -1961,6 +1985,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files=input_files, from_checkpoint=from_checkpoint, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) async def _run_flow() -> Any: @@ -1968,6 +1993,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): inputs, input_files, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) runtime_scope = crewai_event_bus._enter_runtime_scope() @@ -1988,6 +2014,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files: dict[str, FileInput] | None = None, from_checkpoint: CheckpointConfig | None = None, restore_from_state_id: str | None = None, + raise_on_missing_state: bool = False, ) -> Any | AsyncStreamSession[Any]: """Start the flow execution asynchronously. @@ -2006,8 +2033,14 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): run is assigned a fresh ``state.id`` (or ``inputs["id"]`` if pinned), so subsequent ``@persist`` writes land under a separate persistence key. If the referenced state is not - found, falls back silently to baseline. Cannot be combined + found, falls back silently to baseline (see + ``raise_on_missing_state`` to opt out). Cannot be combined with ``from_checkpoint``; passing both raises ``ValueError``. + raise_on_missing_state: When True and ``restore_from_state_id`` is + provided but the referenced state cannot be loaded, raise + ``ValueError`` instead of silently running without hydration. + Defaults to False. ``last_restore_succeeded`` reports the + restore outcome either way. Returns: The final output from the flow, which is the result of the last executed method. @@ -2027,6 +2060,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files=input_files, from_checkpoint=from_checkpoint, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) ctx = baggage.set_baggage("flow_inputs", inputs or {}) @@ -2090,7 +2124,10 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): # available, hydrate self._state from the source UUID's latest snapshot # and reassign state.id to a fresh value so subsequent @persist writes # don't extend the source flow's history. If the source state is not - # found, fall through silently to the existing inputs handling. + # found, fall through silently to the existing inputs handling + # (unless raise_on_missing_state is set). last_restore_succeeded + # records the outcome for callers either way. + self._last_restore_succeeded = None fork_succeeded = False if restore_from_state_id is not None and self.persistence is not None: stored_state = self.persistence.load_state(restore_from_state_id) @@ -2111,12 +2148,27 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): elif isinstance(self._state, BaseModel): setattr(self._state, "id", new_state_id) # noqa: B010 fork_succeeded = True + self._last_restore_succeeded = True else: + self._last_restore_succeeded = False + if raise_on_missing_state: + raise ValueError( + "No flow state found for restore_from_state_id: " + f"'{restore_from_state_id}'" + ) self._log_flow_event( "No flow state found for restore_from_state_id: " f"{restore_from_state_id}; proceeding without hydration", color="yellow", ) + elif restore_from_state_id is not None: + self._last_restore_succeeded = False + if raise_on_missing_state: + raise ValueError( + "Cannot restore from restore_from_state_id: " + f"'{restore_from_state_id}'; no persistence backend " + "is configured for this flow." + ) if inputs: # Override the id in the state if it exists in inputs. @@ -2379,6 +2431,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files: dict[str, FileInput] | None = None, from_checkpoint: CheckpointConfig | None = None, restore_from_state_id: str | None = None, + raise_on_missing_state: bool = False, ) -> Any | AsyncStreamSession[Any]: """Native async method to start the flow execution. Alias for kickoff_async. @@ -2390,6 +2443,9 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): restore_from_state_id: Optional UUID of a previously-persisted flow whose latest snapshot should hydrate this run's state. See ``kickoff_async`` for full semantics. + raise_on_missing_state: When True, a ``restore_from_state_id`` + lookup miss raises ``ValueError`` instead of silently running + without hydration. See ``kickoff_async``. Returns: The final output from the flow, which is the result of the last executed method. @@ -2399,6 +2455,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): input_files, from_checkpoint, restore_from_state_id=restore_from_state_id, + raise_on_missing_state=raise_on_missing_state, ) async def _replay_recorded_events(self) -> None: diff --git a/lib/crewai/tests/test_flow_persistence.py b/lib/crewai/tests/test_flow_persistence.py index b405cc64d..3bbb7b670 100644 --- a/lib/crewai/tests/test_flow_persistence.py +++ b/lib/crewai/tests/test_flow_persistence.py @@ -331,6 +331,118 @@ def test_restore_from_state_id_not_found_silent_fallback(tmp_path): assert flow.state.id != "no-such-uuid" +def test_restore_from_state_id_not_found_raises_when_strict(tmp_path): + """raise_on_missing_state=True turns a restore_from_state_id lookup miss into a + ValueError instead of a silent fresh run.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class StrictFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow = StrictFlow(persistence=persistence) + with pytest.raises(ValueError) as excinfo: + flow.kickoff( + restore_from_state_id="no-such-uuid", + raise_on_missing_state=True, + ) + + msg = str(excinfo.value) + assert "restore_from_state_id" in msg + assert "no-such-uuid" in msg + # The flow never ran: state is untouched and the outcome signal reports the miss. + assert flow.state.counter == 0 + assert flow.last_restore_succeeded is False + + +def test_restore_from_state_id_strict_with_valid_id_restores(tmp_path): + """raise_on_missing_state=True does not change behavior when the source state + exists: hydrate from the snapshot, mint a fresh state.id, report success.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class StrictForkFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow1 = StrictForkFlow(persistence=persistence) + flow1.kickoff() + source_uuid = flow1.state.id + assert flow1.state.counter == 1 + + flow2 = StrictForkFlow(persistence=persistence) + flow2.kickoff( + restore_from_state_id=source_uuid, + raise_on_missing_state=True, + ) + + assert flow2.state.id != source_uuid + assert flow2.state.counter == 2 + assert flow2.last_restore_succeeded is True + + +def test_strict_restore_without_persistence_raises(): + """raise_on_missing_state=True with restore_from_state_id but no persistence + backend raises instead of silently skipping hydration.""" + + class NoPersistenceFlow(Flow[TestState]): + @start() + def step(self): + self.state.counter += 1 + + flow = NoPersistenceFlow() + with pytest.raises(ValueError) as excinfo: + flow.kickoff( + restore_from_state_id="some-uuid", + raise_on_missing_state=True, + ) + msg = str(excinfo.value) + assert "restore_from_state_id" in msg + assert "persistence" in msg + assert flow.last_restore_succeeded is False + + +def test_last_restore_succeeded_signal(tmp_path): + """last_restore_succeeded is None when no restore was requested, False after a + silent miss, and True after a successful restore.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class SignalFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow1 = SignalFlow(persistence=persistence) + assert flow1.last_restore_succeeded is None + flow1.kickoff() + # No restore requested: signal stays None. + assert flow1.last_restore_succeeded is None + source_uuid = flow1.state.id + + # Silent miss (default non-strict): the run proceeds but the signal is False. + flow2 = SignalFlow(persistence=persistence) + flow2.kickoff(restore_from_state_id="no-such-uuid") + assert flow2.state.counter == 1 + assert flow2.last_restore_succeeded is False + + # Successful restore: the signal is True. + flow3 = SignalFlow(persistence=persistence) + flow3.kickoff(restore_from_state_id=source_uuid) + assert flow3.state.counter == 2 + assert flow3.last_restore_succeeded is True + + # A later kickoff without a restore request resets the signal to None. + flow3.kickoff() + assert flow3.last_restore_succeeded is None + + def test_restore_from_state_id_none_is_no_op(tmp_path): """restore_from_state_id=None (default) preserves baseline kickoff behavior.""" db_path = os.path.join(tmp_path, "test_flows.db") @@ -452,6 +564,65 @@ async def test_akickoff_pinned_fork(tmp_path): assert persistence.load_state(pinned_uuid)["counter"] == 2 +@pytest.mark.asyncio +async def test_restore_from_state_id_not_found_raises_when_strict_async(tmp_path): + """kickoff_async and akickoff honor raise_on_missing_state on a lookup miss.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class AsyncStrictFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow = AsyncStrictFlow(persistence=persistence) + with pytest.raises(ValueError) as excinfo: + await flow.kickoff_async( + restore_from_state_id="no-such-uuid", + raise_on_missing_state=True, + ) + assert "no-such-uuid" in str(excinfo.value) + assert flow.state.counter == 0 + assert flow.last_restore_succeeded is False + + flow_alias = AsyncStrictFlow(persistence=persistence) + with pytest.raises(ValueError): + await flow_alias.akickoff( + restore_from_state_id="no-such-uuid", + raise_on_missing_state=True, + ) + assert flow_alias.last_restore_succeeded is False + + +@pytest.mark.asyncio +async def test_restore_from_state_id_strict_with_valid_id_restores_async(tmp_path): + """kickoff_async with raise_on_missing_state=True restores normally when the + source state exists and reports success.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class AsyncStrictForkFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow1 = AsyncStrictForkFlow(persistence=persistence) + await flow1.kickoff_async() + source_uuid = flow1.state.id + + flow2 = AsyncStrictForkFlow(persistence=persistence) + await flow2.kickoff_async( + restore_from_state_id=source_uuid, + raise_on_missing_state=True, + ) + + assert flow2.state.id != source_uuid + assert flow2.state.counter == 2 + assert flow2.last_restore_succeeded is True + + @pytest.mark.asyncio async def test_akickoff_fork_conflict_with_from_checkpoint_raises(): """akickoff must raise the same conflict ValueError as kickoff/kickoff_async when