From 378b59a471d6b4be6b41a04a41f11d214c30516d Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 9 Apr 2026 06:42:28 +0000
Subject: [PATCH] fix: add flow_class param to base save_state signature for
 mypy
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: João <joao@crewai.com>
---
 .../src/crewai/flow/persistence/base.py       |  7 ++++++-
 .../src/crewai/flow/persistence/decorators.py | 21 ++++++-------------
 2 files changed, 12 insertions(+), 16 deletions(-)

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: