diff --git a/api/core/app/layers/pause_state_persist_layer.py b/api/core/app/layers/pause_state_persist_layer.py
index 3dee75c082..eb5dfe9c2d 100644
--- a/api/core/app/layers/pause_state_persist_layer.py
+++ b/api/core/app/layers/pause_state_persist_layer.py
@@ -1,5 +1,5 @@
 from sqlalchemy import Engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_events.base import GraphEngineEvent
@@ -9,7 +9,7 @@ from repositories.factory import DifyAPIRepositoryFactory
 
 
 class PauseStatePersistenceLayer(GraphEngineLayer):
-    def __init__(self, session_factory: Engine | sessionmaker, state_owner_user_id: str):
+    def __init__(self, session_factory: Engine | sessionmaker[Session], state_owner_user_id: str):
         """Create a PauseStatePersistenceLayer.
 
         The `state_owner_user_id` is used when creating state file for pause.
diff --git a/api/repositories/factory.py b/api/repositories/factory.py
index 96f9f886a4..8e098a7059 100644
--- a/api/repositories/factory.py
+++ b/api/repositories/factory.py
@@ -5,7 +5,7 @@ This factory is specifically designed for DifyAPI repositories that handle
 service-layer operations with dependency injection patterns.
 """
 
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from configs import dify_config
 from core.repositories import DifyCoreRepositoryFactory, RepositoryImportError
@@ -25,7 +25,7 @@ class DifyAPIRepositoryFactory(DifyCoreRepositoryFactory):
 
     @classmethod
     def create_api_workflow_node_execution_repository(
-        cls, session_maker: sessionmaker
+        cls, session_maker: sessionmaker[Session]
     ) -> DifyAPIWorkflowNodeExecutionRepository:
         """
         Create a DifyAPIWorkflowNodeExecutionRepository instance based on configuration.
@@ -55,7 +55,7 @@ class DifyAPIRepositoryFactory(DifyCoreRepositoryFactory):
             ) from e
 
     @classmethod
-    def create_api_workflow_run_repository(cls, session_maker: sessionmaker) -> APIWorkflowRunRepository:
+    def create_api_workflow_run_repository(cls, session_maker: sessionmaker[Session]) -> APIWorkflowRunRepository:
         """
         Create an APIWorkflowRunRepository instance based on configuration.