diff --git a/src/aws_durable_execution_sdk_python/context.py b/src/aws_durable_execution_sdk_python/context.py
index 46d48a1..a1a49a5 100644
--- a/src/aws_durable_execution_sdk_python/context.py
+++ b/src/aws_durable_execution_sdk_python/context.py
@@ -44,21 +44,21 @@
 )
 from aws_durable_execution_sdk_python.state import ExecutionState  # noqa: TCH001
 from aws_durable_execution_sdk_python.threading import OrderedCounter
+from aws_durable_execution_sdk_python.types import Callback as CallbackProtocol
+from aws_durable_execution_sdk_python.types import (
+    DurableContext as DurableContextProtocol,
+)
 from aws_durable_execution_sdk_python.types import (
-    BatchResult,
     LoggerInterface,
     StepContext,
     WaitForCallbackContext,
     WaitForConditionCheckContext,
 )
-from aws_durable_execution_sdk_python.types import Callback as CallbackProtocol
-from aws_durable_execution_sdk_python.types import (
-    DurableContext as DurableContextProtocol,
-)
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
+    from aws_durable_execution_sdk_python.concurrency.models import BatchResult
     from aws_durable_execution_sdk_python.state import CheckpointedResult
     from aws_durable_execution_sdk_python.types import LambdaContext
     from aws_durable_execution_sdk_python.waits import WaitForConditionConfig