diff --git a/discord-bot/bot/bot.py b/discord-bot/bot/bot.py index e16504ee..35a3fb23 100644 --- a/discord-bot/bot/bot.py +++ b/discord-bot/bot/bot.py @@ -34,6 +34,9 @@ async def on_starting(event: hikari.StartingEvent): bot.d.oasst_api = OasstApiClient(settings.oasst_api_url, settings.oasst_api_key) + # A set of user id's that are currently doing work. + bot.d.currently_working = set() + @bot.listen() async def on_stopping(event: hikari.StoppingEvent): diff --git a/discord-bot/bot/extensions/work.py b/discord-bot/bot/extensions/work.py index 1d88fb35..767e05ac 100644 --- a/discord-bot/bot/extensions/work.py +++ b/discord-bot/bot/extensions/work.py @@ -34,12 +34,25 @@ MAX_TASK_ACCEPT_TIME = 60 # 1 minute @lightbulb.implements(lightbulb.SlashCommand) async def work(ctx: lightbulb.SlashContext): """Create and handle a task.""" + # make sure the user isn't currently doing a task + currently_working: set[hikari.Snowflakeish] = ctx.bot.d.currently_working + if ctx.author.id in currently_working: + await ctx.respond( + "You are already performing a task. Please complete that one first.", flags=hikari.MessageFlag.EPHEMERAL + ) + return + + currently_working.add(ctx.author.id) + task_type: TaskRequestType = TaskRequestType(ctx.options.type.split(".")[-1]) await ctx.respond("Sending you a task, check your DMs", flags=hikari.MessageFlag.EPHEMERAL) logger.debug(f"Starting task_type: {task_type!r}") - await _handle_task(ctx, task_type) + try: + await _handle_task(ctx, task_type) + finally: + currently_working.remove(ctx.author.id) async def _handle_task(ctx: lightbulb.SlashContext, task_type: TaskRequestType) -> None: