diff --git a/lib/crewai/src/crewai/flow/runtime/__init__.py b/lib/crewai/src/crewai/flow/runtime/__init__.py index 7f74412b5..a76a5bace 100644 --- a/lib/crewai/src/crewai/flow/runtime/__init__.py +++ b/lib/crewai/src/crewai/flow/runtime/__init__.py @@ -1864,6 +1864,40 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): if hasattr(self._state, key): object.__setattr__(self._state, key, value) + @staticmethod + def _missing_state_error(restore_from_state_id: str) -> ValueError: + return ValueError( + f"No flow state found for restore_from_state_id: '{restore_from_state_id}'" + ) + + @staticmethod + def _no_persistence_error(restore_from_state_id: str) -> ValueError: + return ValueError( + "Cannot restore from restore_from_state_id: " + f"'{restore_from_state_id}'; no persistence backend " + "is configured for this flow." + ) + + def _ensure_restorable_state( + self, + restore_from_state_id: str | None, + raise_on_missing_state: bool, + ) -> None: + """Eagerly validate a strict restore request before deferred execution. + + Streaming kickoffs return a session before the runtime evaluates + ``restore_from_state_id``; this pre-check keeps + ``raise_on_missing_state=True`` failing at call time, matching the + non-streaming path. The runtime still performs the authoritative load + (and raises again) during execution. + """ + if not raise_on_missing_state or restore_from_state_id is None: + return + if self.persistence is None: + raise self._no_persistence_error(restore_from_state_id) + if not self.persistence.load_state(restore_from_state_id): + raise self._missing_state_error(restore_from_state_id) + def stream_events( self, inputs: dict[str, Any] | None = None, @@ -1872,7 +1906,13 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): 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.""" + """Run the flow and stream all scoped public ``StreamFrame`` events. + + With ``raise_on_missing_state=True``, a ``restore_from_state_id`` + lookup miss raises ``ValueError`` here, before the session is + returned, so callers do not have to consume frames to see the error. + """ + self._ensure_restorable_state(restore_from_state_id, raise_on_missing_state) result_holder: list[Any] = [] state = create_frame_streaming_state(result_holder, use_async=False) output_holder: list[StreamSession[Any]] = [] @@ -1907,7 +1947,13 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): restore_from_state_id: str | None = None, raise_on_missing_state: bool = False, ) -> AsyncStreamSession[Any]: - """Run the flow asynchronously and stream scoped public frames.""" + """Run the flow asynchronously and stream scoped public frames. + + With ``raise_on_missing_state=True``, a ``restore_from_state_id`` + lookup miss raises ``ValueError`` here, before the session is + returned, so callers do not have to consume frames to see the error. + """ + self._ensure_restorable_state(restore_from_state_id, raise_on_missing_state) result_holder: list[Any] = [] state = create_frame_streaming_state(result_holder, use_async=True) output_holder: list[AsyncStreamSession[Any]] = [] @@ -2152,10 +2198,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): 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}'" - ) + raise self._missing_state_error(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", @@ -2164,11 +2207,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): 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." - ) + raise self._no_persistence_error(restore_from_state_id) if inputs: # Override the id in the state if it exists in inputs. diff --git a/lib/crewai/tests/test_flow_persistence.py b/lib/crewai/tests/test_flow_persistence.py index 3bbb7b670..9fcce5c3d 100644 --- a/lib/crewai/tests/test_flow_persistence.py +++ b/lib/crewai/tests/test_flow_persistence.py @@ -407,6 +407,112 @@ def test_strict_restore_without_persistence_raises(): assert flow.last_restore_succeeded is False +def test_stream_events_strict_restore_miss_raises_eagerly(tmp_path): + """stream_events with raise_on_missing_state=True raises at session creation + time on a restore_from_state_id miss, not while consuming frames.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class StreamStrictFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow = StreamStrictFlow(persistence=persistence) + with pytest.raises(ValueError) as excinfo: + flow.stream_events( + restore_from_state_id="no-such-uuid", + raise_on_missing_state=True, + ) + assert "no-such-uuid" in str(excinfo.value) + # Nothing ran. + assert flow.state.counter == 0 + + +def test_kickoff_stream_strict_restore_miss_raises_eagerly(tmp_path): + """kickoff with stream=True and raise_on_missing_state=True raises immediately + on a restore miss instead of returning a StreamSession.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class StreamingKickoffFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow = StreamingKickoffFlow(persistence=persistence) + flow.stream = True + with pytest.raises(ValueError) as excinfo: + flow.kickoff( + 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 + + +def test_stream_events_strict_restore_valid_id_streams(tmp_path): + """stream_events with raise_on_missing_state=True streams normally when the + source state exists: hydrates, mints a fresh state.id, reports success.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class StreamForkFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow1 = StreamForkFlow(persistence=persistence) + flow1.kickoff() + source_uuid = flow1.state.id + + flow2 = StreamForkFlow(persistence=persistence) + with flow2.stream_events( + restore_from_state_id=source_uuid, + raise_on_missing_state=True, + ) as stream: + list(stream.events) + + assert flow2.state.id != source_uuid + assert flow2.state.counter == 2 + assert flow2.last_restore_succeeded is True + + +@pytest.mark.asyncio +async def test_astream_strict_restore_miss_raises_eagerly(tmp_path): + """astream and kickoff_async(stream=True) raise at call time on a restore miss + with raise_on_missing_state=True.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class AsyncStreamStrictFlow(Flow[TestState]): + @start() + @persist(persistence) + def step(self): + self.state.counter += 1 + + flow = AsyncStreamStrictFlow(persistence=persistence) + with pytest.raises(ValueError) as excinfo: + flow.astream( + 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 + + flow2 = AsyncStreamStrictFlow(persistence=persistence) + flow2.stream = True + with pytest.raises(ValueError): + await flow2.kickoff_async( + restore_from_state_id="no-such-uuid", + raise_on_missing_state=True, + ) + assert flow2.state.counter == 0 + + 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."""