diff --git a/lib/crewai/src/crewai/flow/persistence/base.py b/lib/crewai/src/crewai/flow/persistence/base.py index fd31c9af1..26bc400e6 100644 --- a/lib/crewai/src/crewai/flow/persistence/base.py +++ b/lib/crewai/src/crewai/flow/persistence/base.py @@ -46,7 +46,11 @@ class FlowPersistence(BaseModel, ABC): @abstractmethod def save_state( - self, flow_uuid: str, method_name: str, state_data: dict[str, Any] | BaseModel + self, + flow_uuid: str, + method_name: str, + state_data: dict[str, Any] | BaseModel, + flow_class: str | None = None, ) -> None: """Persist the flow state after method completion. @@ -54,6 +58,7 @@ class FlowPersistence(BaseModel, ABC): flow_uuid: Unique identifier for the flow instance method_name: Name of the method that just completed state_data: Current state data (either dict or Pydantic model) + flow_class: Optional name of the flow class for auto-restore support """ @abstractmethod diff --git a/lib/crewai/src/crewai/flow/persistence/decorators.py b/lib/crewai/src/crewai/flow/persistence/decorators.py index 126bbd39a..f366a599a 100644 --- a/lib/crewai/src/crewai/flow/persistence/decorators.py +++ b/lib/crewai/src/crewai/flow/persistence/decorators.py @@ -110,21 +110,12 @@ class PersistenceDecorator: try: state_data = state._unwrap() if hasattr(state, "_unwrap") else state flow_class_name = type(flow_instance).__name__ - try: - persistence_instance.save_state( - flow_uuid=flow_uuid, - method_name=method_name, - state_data=state_data, - flow_class=flow_class_name, - ) - except TypeError: - # Fallback for custom persistence backends that - # don't accept the flow_class parameter - persistence_instance.save_state( - flow_uuid=flow_uuid, - method_name=method_name, - state_data=state_data, - ) + persistence_instance.save_state( + flow_uuid=flow_uuid, + method_name=method_name, + state_data=state_data, + flow_class=flow_class_name, + ) except Exception as e: error_msg = LOG_MESSAGES["save_error"].format(method_name, str(e)) if verbose: