diff --git a/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSMBuilder.kt b/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSMBuilder.kt
index 96ca461801..17015fe5ee 100644
--- a/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSMBuilder.kt
+++ b/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/BehaviourContextWithFSMBuilder.kt
@@ -89,7 +89,7 @@ suspend fun TelegramBot.buildBehaviourWithFSM(
     scope: CoroutineScope = defaultCoroutineScopeProvider(),
     statesManager: StatesManager = DefaultStatesManager(InMemoryDefaultStatesManagerRepo()),
     handlersPreset: MutableList<BehaviourWithFSMStateHandlerHolder<*>> = mutableListOf(),
-    block: suspend BehaviourContextWithFSMBuilder.() -> Unit
+    block: CustomBehaviourContextReceiver<BehaviourContextWithFSMBuilder, Unit>
 ) = BehaviourContextWithFSMBuilder(
     DefaultBehaviourContext(this, scope, upstreamUpdatesFlow = upstreamUpdatesFlow),
     statesManager,
@@ -108,7 +108,7 @@ suspend fun TelegramBot.buildBehaviourWithFSMAndLongPolling(
     scope: CoroutineScope = defaultCoroutineScopeProvider(),
     statesManager: StatesManager = DefaultStatesManager(InMemoryDefaultStatesManagerRepo()),
     presetHandlers: MutableList<BehaviourWithFSMStateHandlerHolder<*>> = mutableListOf(),
-    block: suspend BehaviourContextWithFSMBuilder.() -> Unit
+    block: CustomBehaviourContextReceiver<BehaviourContextWithFSMBuilder, Unit>
 ) = buildBehaviourWithFSM(upstreamUpdatesFlow, scope, statesManager, presetHandlers, block).run {
     this to scope.launch {
         start()
@@ -169,7 +169,7 @@ suspend fun TelegramBot.buildBehaviourWithFSM(
     defaultExceptionsHandler: ExceptionHandler<Unit>? = null,
     statesManager: StatesManager = DefaultStatesManager(InMemoryDefaultStatesManagerRepo()),
     presetHandlers: MutableList<BehaviourWithFSMStateHandlerHolder<*>> = mutableListOf(),
-    block: BehaviourContextReceiver<Unit>
+    block: CustomBehaviourContextReceiver<BehaviourContextWithFSMBuilder, Unit>
 ) = FlowsUpdatesFilter().let {
     buildBehaviourWithFSM(
         it,
diff --git a/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/TelegramBotWithFSM.kt b/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/TelegramBotWithFSM.kt
index 23eb0fe7bc..39f4ba3c69 100644
--- a/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/TelegramBotWithFSM.kt
+++ b/tgbotapi.extensions.behaviour_builder.fsm/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/behaviour_builder/TelegramBotWithFSM.kt
@@ -37,7 +37,7 @@ suspend fun telegramBotWithBehaviourAndFSM(
     defaultExceptionsHandler: ExceptionHandler<Unit>? = null,
     statesManager: StatesManager = DefaultStatesManager(InMemoryDefaultStatesManagerRepo()),
     presetHandlers: MutableList<BehaviourWithFSMStateHandlerHolder<*>> = mutableListOf(),
-    block: BehaviourContextReceiver<Unit>
+    block: CustomBehaviourContextReceiver<BehaviourContextWithFSMBuilder, Unit>
 ): TelegramBot = telegramBot(
     token,
     apiUrl,
@@ -75,7 +75,7 @@ suspend fun telegramBotWithBehaviourAndFSM(
     defaultExceptionsHandler: ExceptionHandler<Unit>? = null,
     statesManager: StatesManager = DefaultStatesManager(InMemoryDefaultStatesManagerRepo()),
     presetHandlers: MutableList<BehaviourWithFSMStateHandlerHolder<*>> = mutableListOf(),
-    block: BehaviourContextReceiver<Unit>
+    block: CustomBehaviourContextReceiver<BehaviourContextWithFSMBuilder, Unit>
 ): Pair<TelegramBot, Job> {
     return telegramBot(
         token,