diff --git a/kb_pipeline/__init__.py b/kb_pipeline/__init__.py
index 889531f..b470a31 100644
--- a/kb_pipeline/__init__.py
+++ b/kb_pipeline/__init__.py
@@ -12,6 +12,7 @@ from .schemas import (
     WorkflowState,
     StepDefinition,
 )
+from .state import WorkflowStateManager
 
 __all__ = [
     "StepInput",
@@ -21,6 +22,7 @@ __all__ = [
     "StepRecord",
     "WorkflowState",
     "StepDefinition",
+    "WorkflowStateManager",
 ]
 
 __version__ = "0.1.0"