diff --git a/bot.py b/bot.py
index 7e56d9e..f421577 100644
--- a/bot.py
+++ b/bot.py
@@ -773,14 +773,16 @@ class MenuBuilder:
     def get_menu(self, menu_name: str) -> List[MenuItem]:
         return self._menus.get(menu_name, [])
     
-    def get_keyboard(self, menu_name: str, user_id: int = None) -> InlineKeyboardMarkup:
+    def get_keyboard(self, menu_name: str, user_id: int = None, state: UserState = None) -> InlineKeyboardMarkup:
         """Создает InlineKeyboard для меню."""
         items = self._menus.get(menu_name, [])
         keyboard = []
         
         # Для главного меню — динамически меняем кнопку ИИ
         if menu_name == "main" and user_id:
-            state = state_manager.get(user_id)
+            # Используем переданное состояние или получаем из менеджера
+            if state is None:
+                state = state_manager.get(user_id)
             logger.info(f"get_keyboard: user_id={user_id}, ai_chat_mode={state.ai_chat_mode}")
             
             for item in items:
@@ -979,7 +981,7 @@ async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
         f"Или выберите сервер в меню.\n"
         f"Команда /help покажет справку.",
         parse_mode="Markdown",
-        reply_markup=menu_builder.get_keyboard("main", user_id=update.effective_user.id)
+        reply_markup=menu_builder.get_keyboard("main", user_id=update.effective_user.id, state=state)
     )
 
 
@@ -1004,7 +1006,7 @@ async def menu_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
         f"📁 *Директория:* `{working_dir}`\n\n"
         f"Выберите действие:",
         parse_mode="Markdown",
-        reply_markup=menu_builder.get_keyboard("main", user_id=update.effective_user.id)
+        reply_markup=menu_builder.get_keyboard("main", user_id=update.effective_user.id, state=state)
     )
 
 
@@ -1081,7 +1083,7 @@ async def menu_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
             f"🏠 *Главное меню*\n\n"
             f"💬 *Чат с ИИ:* {ai_status}",
             parse_mode="Markdown",
-            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id)
+            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id, state=state)
         )
     
     elif callback == "preset_menu":
@@ -1327,7 +1329,7 @@ async def menu_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
                 f"📍 `{server.description}`\n\n"
                 f"Теперь команды выполняются на этом сервере.",
                 parse_mode="Markdown",
-                reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id)
+                reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id, state=state)
             )
             state.current_menu = "main"
         else:
@@ -1443,7 +1445,7 @@ async def menu_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
             f"Бот позволяет безопасно выполнять команды\n"
             f"на вашем сервере через интерфейс Telegram.",
             parse_mode="Markdown",
-            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id)
+            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id, state=state)
         )
         state.current_menu = "main"
     
@@ -1461,7 +1463,7 @@ async def menu_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
             f"Режим чата с агентом {action}.\n"
             f"Теперь все сообщения будут отправляться в Qwen Code.",
             parse_mode="Markdown",
-            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id)
+            reply_markup=menu_builder.get_keyboard("main", user_id=query.from_user.id, state=state)
         )
         state.current_menu = "main"