fix: add flow_class param to base save_state signature for mypy

Co-Authored-By: João <joao@crewai.com>

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

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: