[EPD-196] Fix eager restore pre-check: record miss signal, conflict check first

- _ensure_restorable_state now sets last_restore_succeeded=False before
  raising, matching the runtime path, so streaming callers failing in the
  pre-check see the miss signal instead of None or a stale True.
- stream_events/astream validate the from_checkpoint/restore_from_state_id
  mutual exclusion before the strict restore pre-check, so the conflict
  ValueError wins over the restore-miss one (matching non-streaming
  precedence). Conflict check extracted to _check_restore_conflict and
  reused by kickoff/kickoff_async.
- Failing-first tests for the stale signal (incl. prior-success staleness
  and no-persistence path) and for conflict precedence on both stream
  entry points.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Joao Moura
2026-07-13 22:52:45 -07:00
parent 2aae1fcf79
commit 79488fda15
2 changed files with 112 additions and 14 deletions

View File

@@ -1878,6 +1878,19 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
"is configured for this flow."
)
@staticmethod
def _check_restore_conflict(
from_checkpoint: CheckpointConfig | None,
restore_from_state_id: str | None,
) -> None:
"""Reject combining ``from_checkpoint`` with ``restore_from_state_id``."""
if from_checkpoint is not None and restore_from_state_id is not None:
raise ValueError(
"Cannot combine `from_checkpoint` and `restore_from_state_id`. "
"These parameters target different state systems "
"(Checkpointing and @persist) and cannot be used together."
)
def _ensure_restorable_state(
self,
restore_from_state_id: str | None,
@@ -1888,14 +1901,17 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
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.
non-streaming path. It records the miss in ``last_restore_succeeded``
before raising, exactly like the runtime 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:
self._last_restore_succeeded = False
raise self._no_persistence_error(restore_from_state_id)
if not self.persistence.load_state(restore_from_state_id):
self._last_restore_succeeded = False
raise self._missing_state_error(restore_from_state_id)
def stream_events(
@@ -1912,6 +1928,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
lookup miss raises ``ValueError`` here, before the session is
returned, so callers do not have to consume frames to see the error.
"""
self._check_restore_conflict(from_checkpoint, restore_from_state_id)
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)
@@ -1953,6 +1970,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
lookup miss raises ``ValueError`` here, before the session is
returned, so callers do not have to consume frames to see the error.
"""
self._check_restore_conflict(from_checkpoint, restore_from_state_id)
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)
@@ -2016,12 +2034,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
Returns:
The final output from the flow or StreamSession if streaming.
"""
if from_checkpoint is not None and restore_from_state_id is not None:
raise ValueError(
"Cannot combine `from_checkpoint` and `restore_from_state_id`. "
"These parameters target different state systems "
"(Checkpointing and @persist) and cannot be used together."
)
self._check_restore_conflict(from_checkpoint, restore_from_state_id)
restored = apply_checkpoint(self, from_checkpoint)
if restored is not None:
return restored.kickoff(inputs=inputs, input_files=input_files)
@@ -2091,12 +2104,7 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta):
Returns:
The final output from the flow, which is the result of the last executed method.
"""
if from_checkpoint is not None and restore_from_state_id is not None:
raise ValueError(
"Cannot combine `from_checkpoint` and `restore_from_state_id`. "
"These parameters target different state systems "
"(Checkpointing and @persist) and cannot be used together."
)
self._check_restore_conflict(from_checkpoint, restore_from_state_id)
restored = apply_checkpoint(self, from_checkpoint)
if restored is not None:
return await restored.kickoff_async(inputs=inputs, input_files=input_files)

View File

@@ -513,6 +513,96 @@ async def test_astream_strict_restore_miss_raises_eagerly(tmp_path):
assert flow2.state.counter == 0
def test_stream_eager_strict_miss_sets_last_restore_succeeded_false(tmp_path):
"""The eager streaming pre-check must record the miss in
last_restore_succeeded before raising — including overwriting a stale True
from a previous successful restore."""
db_path = os.path.join(tmp_path, "test_flows.db")
persistence = SQLiteFlowPersistence(db_path)
class EagerSignalFlow(Flow[TestState]):
@start()
@persist(persistence)
def step(self):
self.state.counter += 1
flow1 = EagerSignalFlow(persistence=persistence)
flow1.kickoff()
source_uuid = flow1.state.id
# Prime the signal with a successful restore first.
flow2 = EagerSignalFlow(persistence=persistence)
flow2.kickoff(restore_from_state_id=source_uuid)
assert flow2.last_restore_succeeded is True
# Eager streaming miss on the same instance: signal flips to False, not stale True.
with pytest.raises(ValueError):
flow2.stream_events(
restore_from_state_id="no-such-uuid",
raise_on_missing_state=True,
)
assert flow2.last_restore_succeeded is False
# Fresh instance failing in the astream pre-check: False, not None.
flow3 = EagerSignalFlow(persistence=persistence)
with pytest.raises(ValueError):
flow3.astream(
restore_from_state_id="no-such-uuid",
raise_on_missing_state=True,
)
assert flow3.last_restore_succeeded is False
# No-persistence eager failure records the miss too.
class NoPersistenceStreamFlow(Flow[TestState]):
@start()
def step(self):
self.state.counter += 1
flow4 = NoPersistenceStreamFlow()
with pytest.raises(ValueError):
flow4.stream_events(
restore_from_state_id="some-uuid",
raise_on_missing_state=True,
)
assert flow4.last_restore_succeeded is False
def test_stream_conflict_check_wins_over_strict_restore_miss(tmp_path):
"""stream_events/astream with both from_checkpoint and restore_from_state_id
raise the mutual-exclusion ValueError, not the restore-miss one, even with
raise_on_missing_state=True and a missing id (matches non-streaming precedence)."""
from crewai.state import CheckpointConfig
db_path = os.path.join(tmp_path, "test_flows.db")
persistence = SQLiteFlowPersistence(db_path)
class StreamConflictFlow(Flow[TestState]):
@start()
@persist(persistence)
def step(self):
self.state.counter += 1
flow = StreamConflictFlow(persistence=persistence)
with pytest.raises(ValueError) as excinfo:
flow.stream_events(
from_checkpoint=CheckpointConfig(),
restore_from_state_id="no-such-uuid",
raise_on_missing_state=True,
)
msg = str(excinfo.value)
assert "Cannot combine" in msg
assert "from_checkpoint" in msg
flow2 = StreamConflictFlow(persistence=persistence)
with pytest.raises(ValueError) as excinfo2:
flow2.astream(
from_checkpoint=CheckpointConfig(),
restore_from_state_id="no-such-uuid",
raise_on_missing_state=True,
)
assert "Cannot combine" in str(excinfo2.value)
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."""