Coverage for slack_bot / shell_main.py: 0%

1import os 

2import sys 

3import click 

4from slack_bolt import App 

5from slack_bolt.adapter.socket_mode import SocketModeHandler 

6from dotenv import load_dotenv 

7import re 

8import subprocess 

9 

10# Add project root to path 

11sys.path.append(os.getcwd()) 

12 

13from health.utils.logging_config import setup_logger 

14from health.utils.env_loader import load_env_with_extras 

15from slack_bot.dispatcher import MessageDispatcher 

16from slack_bot.shell.help_text import SHELL_HELP_TEXT 

17from health.utils.time_utils import get_current_time_str 

18 

19logger = setup_logger(__name__) 

20 

21# Load environment variables 

22load_env_with_extras() 

23 

24@click.command() 

25@click.option("--model", default=None, help="Gemini model to use (overrides .env)") 

26def main(model: str): 

27 """Butler Slack Bot - Shell Edition.""" 

28 

29 # Guard: require explicit opt-in to prevent accidental startup 

30 if os.environ.get("SHELL_BOT_ENABLED", "").lower() != "true": 

31 logger.error("Shell Bot is disabled on this machine. Set SHELL_BOT_ENABLED=true in .env to enable.") 

32 sys.exit(1) 

33 

34 # Override model in environment only if specified via CLI 

35 if model: 

36 os.environ["GEMINI_MODEL"] = model 

37 logger.info(f"Using model from CLI: {model}") 

38 

39 # Get the actual model being used (from CLI or .env) 

40 actual_model = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash-exp") 

41 logger.info(f"Starting Shell Bot with Gemini model: {actual_model}") 

42 

43 # Initialize Slack App with SHELL-specific tokens 

44 bot_token = os.environ.get("SHELL_SLACK_BOT_TOKEN") 

45 app_token = os.environ.get("SHELL_SLACK_APP_TOKEN") 

46 

47 if not bot_token or not app_token: 

48 logger.error("Missing SHELL_SLACK_BOT_TOKEN or SHELL_SLACK_APP_TOKEN in .env") 

49 sys.exit(1) 

50 

51 app = App(token=bot_token) 

52 

53 # --- SHELL CONFIGURATION --- 

54 MAX_TRACKED_MESSAGES = 1000 

55 

56 # NOTE: We can't easily inject dynamic time into the *static* dispatcher init here  

57 # if we want it to update per-request without recreating the dispatcher. 

58 # However, ShellBot is often single-session.  

59 # Better approach: We'll recreate the dispatcher or update the prompt inside handle_message? 

60 # Actually, MessageDispatcher init takes system_instruction. 

61 # We can just Init it here with STARTUP time, which is better than nothing,  

62 # but ideally we want per-request time.  

63 # The simplest fix for now: Inject it at module level (restart required for day change anyway). 

64 

65 SHELL_SYSTEM_PROMPT = f"""You are a Shell Assistant, an expert in zsh and system administration. 

66Current Time: {get_current_time_str()} 

67 

68Your Goal: Help the user execute commands and understand their system. 

69 

70Capabilities: 

71- You have a persistent zsh session available via the `execute_shell` tool. 

72- You can query files, check processes, and manage the system. 

73 

74CRITICAL ANTI-HALLUCINATION RULES: 

751. **TOOL-FIRST POLICY**: When user asks about CURRENT system state (e.g., "what files are on my desktop", "what processes are running", "system load"), you MUST call `execute_shell` FIRST. NEVER answer based on conversation history or memory. 

762. **ALWAYS PREFER EXECUTION**: If user asks for a command (e.g., "list files"), EXECUTE IT using `execute_shell("ls -la")`. DO NOT just explain it. 

773. **100% OUTPUT FIDELITY**: Report EXACTLY what the terminal returns. If `execute_shell` returns empty string, error, or loading indicator, say "The command returned no output" or "Command failed with error: X". NEVER invent, guess, or assume what the output "should" look like. 

784. **EXPLICIT UNCERTAINTY**: If you haven't executed a command in the current turn, DO NOT provide specific file names, process lists, or system state details. 

79 

80Guidelines: 

815. **Numbered Shortcuts**: Provide numbered lists (1, 2, 3). If user replies with a number, EXECUTE that option immediately. 

826. **Be Concise**: Summarize long outputs, but show relevant parts. 

837. **Safety**: Warn before destructive commands (rm -rf), but respect user's decision. 

84""" 

85 

86 # Only include the shell tool 

87 from slack_bot.tools.registry import TOOLS_SCHEMA 

88 SHELL_TOOLS = [t for t in TOOLS_SCHEMA if t["function"]["name"] == "execute_shell"] 

89 

90 # Initialize Dispatcher with Shell Identity 

91 dispatcher = MessageDispatcher( 

92 bot_token=bot_token, 

93 system_instruction=SHELL_SYSTEM_PROMPT, 

94 tools=SHELL_TOOLS 

95 ) 

96 

97 # Message deduplication: Track processed message IDs 

98 processed_messages = set() 

99 MAX_TRACKED_MESSAGES = 1000 

100 

101 # Message Handler 

102 @app.event("message") 

103 def handle_message(message, say): 

104 try: 

105 # Deduplicate using client_msg_id or event_ts 

106 msg_id = message.get("client_msg_id") or message.get("ts") 

107 

108 if msg_id in processed_messages: 

109 logger.debug(f"Skipping duplicate message: {msg_id}") 

110 return 

111 

112 processed_messages.add(msg_id) 

113 if len(processed_messages) > MAX_TRACKED_MESSAGES: 

114 to_remove = list(processed_messages)[:MAX_TRACKED_MESSAGES // 2] 

115 for old_id in to_remove: 

116 processed_messages.discard(old_id) 

117 

118 text = message.get("text", "") 

119 channel = message["channel"] 

120 user = message["user"] 

121 files = message.get("files", []) 

122 

123 # --- HANDLE HELP COMMAND --- 

124 if text.strip().lower() in ["help", "help?", "帮助", "menu"]: 

125 say(SHELL_HELP_TEXT) 

126 return 

127 

128 # Generate Request ID 

129 import uuid 

130 import time 

131 request_id = str(uuid.uuid4())[:8] 

132 

133 logger.info(f"[{request_id}] 📥 ShellBot received: {text}") 

134 

135 # --- FAILSAFE: CONFIG SWITCHING --- 

136 gmode_match = re.match(r"^### gmode (.*?) ###$", text.strip()) 

137 if gmode_match: 

138 target_config = gmode_match.group(1) 

139 logger.warning(f"[{request_id}] 🚨 Received FAILSAFE cmd: {target_config}") 

140 

141 try: 

142 # Execute switch script 

143 script_path = os.path.join(os.getcwd(), "scripts", "select_gemini.py") 

144 result = subprocess.run( 

145 [sys.executable, script_path, target_config], 

146 capture_output=True, 

147 text=True 

148 ) 

149 

150 if result.returncode == 0: 

151 # Special handling for 'check' and 'check-detailed' commands - show output, don't restart 

152 if target_config in ["check", "check-detailed"]: 

153 output = result.stdout + result.stderr 

154 say(f"🔍 **Health Check Results**:\n```\n{output}\n```") 

155 return 

156 

157 say(f"♻️ **System Reset**: Switched to `{target_config}`. Restarting process...") 

158 # Allow message to be sent 

159 time.sleep(1) 

160 

161 # Self-Restart 

162 logger.warning("RESTARTING PROCESS via os.execv...") 

163 os.execv(sys.executable, [sys.executable] + sys.argv) 

164 else: 

165 say(f"❌ Command Failed:\n```\n{result.stderr}\n```") 

166 return 

167 

168 except Exception as e: 

169 logger.error(f"Failsafe Error: {e}") 

170 say(f"❌ Failsafe Error: {e}") 

171 return 

172 # ---------------------------------- 

173 

174 # --- FAILSAFE: BOT COMMANDS --- 

175 bot_match = re.match(r"^### bot (.*?) ###$", text.strip()) 

176 if bot_match: 

177 bot_args = bot_match.group(1).split() 

178 logger.warning(f"[{request_id}] 🚨 Received FAILSAFE cmd: bot {' '.join(bot_args)}") 

179 

180 try: 

181 # Execute bot_manager.py 

182 script_path = os.path.join(os.getcwd(), "scripts", "bot_manager.py") 

183 result = subprocess.run( 

184 [sys.executable, script_path] + bot_args, 

185 capture_output=True, 

186 text=True 

187 ) 

188 

189 output = result.stdout + "\n" + result.stderr 

190 if result.returncode == 0: 

191 say(f"🤖 **Bot Manager Output**:\n```\n{output.strip()}\n```") 

192 else: 

193 say(f"❌ Command Failed:\n```\n{output.strip()}\n```") 

194 return 

195 

196 except Exception as e: 

197 logger.error(f"Failsafe Error: {e}") 

198 say(f"❌ Failsafe Error: {e}") 

199 return 

200 # ---------------------------------- 

201 

202 # Send temporary "Processing..." message 

203 ack_response = say(f"⏳ ShellBot is thinking...") 

204 ack_ts = ack_response["ts"] 

205 

206 # Dispatch 

207 dispatcher.dispatch(text, channel, user, response_ts=ack_ts, request_id=request_id, files=files) 

208 

209 except Exception as e: 

210 logger.error(f"Error handling message: {e}", exc_info=True) 

211 try: 

212 say(f"🐛 ShellBot Error: {e}") 

213 except Exception: 

214 pass 

215 

216 # Start Socket Mode 

217 try: 

218 handler = SocketModeHandler(app, app_token) 

219 logger.info("⚡️ Shell Bot is connected to Slack!") 

220 handler.start() 

221 except Exception as e: 

222 logger.error(f"Failed to start SocketModeHandler: {e}") 

223 sys.exit(1) 

224 

225if __name__ == "__main__": 

226 main()