Add infinite retry for starting up Slack bot (#540)

backend/danswer/bots/slack/listener.py

@@ -1,5 +1,6 @@
 import logging
 import re
+import time
 from collections.abc import MutableMapping
 from typing import Any
 from typing import cast
@@ -11,9 +12,9 @@ from slack_sdk.socket_mode.response import SocketModeResponse
 
 from danswer.bots.slack.handlers.handle_feedback import handle_slack_feedback
 from danswer.bots.slack.handlers.handle_message import handle_message
+from danswer.bots.slack.tokens import fetch_tokens
 from danswer.bots.slack.utils import decompose_block_id
 from danswer.dynamic_configs.interface import ConfigNotFoundError
-from danswer.server.slack_bot_management import get_tokens
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -22,6 +23,10 @@ logger = setup_logger()
 _CHANNEL_ID = "channel_id"
 
 
+class MissingTokensException(Exception):
+    pass
+
+
 class _ChannelIdAdapter(logging.LoggerAdapter):
     """This is used to add the channel ID to all log messages
     emitted in this file"""
@@ -40,9 +45,9 @@ def _get_socket_client() -> SocketModeClient:
     # For more info on how to set this up, checkout the docs:
     # https://docs.danswer.dev/slack_bot_setup
     try:
-        slack_bot_tokens = get_tokens()
+        slack_bot_tokens = fetch_tokens()
     except ConfigNotFoundError:
-        raise RuntimeError("Slack tokens not found")
+        raise MissingTokensException("Slack tokens not found")
     return SocketModeClient(
         # This app-level token will be used only for establishing a connection
         app_token=slack_bot_tokens.app_token,
@@ -222,14 +227,21 @@ def process_slack_event(client: SocketModeClient, req: SocketModeRequest) -> Non
 # NOTE: we are using Web Sockets so that you can run this from within a firewalled VPC
 # without issue.
 if __name__ == "__main__":
-    socket_client = _get_socket_client()
-    socket_client.socket_mode_request_listeners.append(process_slack_event)  # type: ignore
+    try:
+        socket_client = _get_socket_client()
+        socket_client.socket_mode_request_listeners.append(process_slack_event)  # type: ignore
 
-    # Establish a WebSocket connection to the Socket Mode servers
-    logger.info("Listening for messages from Slack...")
-    socket_client.connect()
+        # Establish a WebSocket connection to the Socket Mode servers
+        logger.info("Listening for messages from Slack...")
+        socket_client.connect()
 
-    # Just not to stop this process
-    from threading import Event
+        # Just not to stop this process
+        from threading import Event
 
-    Event().wait()
+        Event().wait()
+    except MissingTokensException:
+        # try again every 30 seconds. This is needed since the user may add tokens
+        # via the UI at any point in the programs lifecycle - if we just allow it to
+        # fail, then the user will need to restart the containers after adding tokens
+        logger.debug("Missing Slack Bot tokens - waiting 60 seconds and trying again")
+        time.sleep(60)

backend/danswer/bots/slack/tokens.py

@@ -2,7 +2,6 @@ import os
 from typing import cast
 
 from danswer.dynamic_configs import get_dynamic_config_store
-from danswer.dynamic_configs.interface import ConfigNotFoundError
 from danswer.server.models import SlackBotTokens
 
 
@@ -17,12 +16,9 @@ def fetch_tokens() -> SlackBotTokens:
         return SlackBotTokens(app_token=app_token, bot_token=bot_token)
 
     dynamic_config_store = get_dynamic_config_store()
-    try:
-        return SlackBotTokens(
-            **cast(dict, dynamic_config_store.load(key=_SLACK_BOT_TOKENS_CONFIG_KEY))
-        )
-    except ConfigNotFoundError as e:
-        raise ValueError from e
+    return SlackBotTokens(
+        **cast(dict, dynamic_config_store.load(key=_SLACK_BOT_TOKENS_CONFIG_KEY))
+    )
 
 
 def save_tokens(

backend/danswer/server/slack_bot_management.py

@@ -14,6 +14,7 @@ from danswer.db.slack_bot_config import fetch_slack_bot_configs
 from danswer.db.slack_bot_config import insert_slack_bot_config
 from danswer.db.slack_bot_config import remove_slack_bot_config
 from danswer.db.slack_bot_config import update_slack_bot_config
+from danswer.dynamic_configs.interface import ConfigNotFoundError
 from danswer.server.models import DocumentSet
 from danswer.server.models import SlackBotConfig
 from danswer.server.models import SlackBotConfigCreationRequest
@@ -169,5 +170,5 @@ def put_tokens(tokens: SlackBotTokens) -> None:
 def get_tokens() -> SlackBotTokens:
     try:
         return fetch_tokens()
-    except ValueError:
+    except ConfigNotFoundError:
         raise HTTPException(status_code=404, detail="No tokens found")