atlasbot: hard-timeout matrix answers by mode

atlasbot/matrix/bot.py

@@ -175,7 +175,14 @@ class MatrixBot:
                 lambda q, m, h, cid, obs: self._engine.answer(q, mode=m, history=h, observer=obs, conversation_id=cid)
             )
             history = self._history.get(room_id, [])
-            result = await handler(question, mode, history, room_id, observer)
+            timeout_sec = _mode_timeout_sec(self._settings, mode)
+            if timeout_sec > 0:
+                result = await asyncio.wait_for(
+                    handler(question, mode, history, room_id, observer),
+                    timeout=timeout_sec + 1.0,
+                )
+            else:
+                result = await handler(question, mode, history, room_id, observer)
             elapsed = time.monotonic() - started
             await self._client.send_message(token, room_id, result.reply)
             log.info(
@@ -190,6 +197,35 @@ class MatrixBot:
             )
             history.append({"q": question, "a": result.reply})
             self._history[room_id] = history[-4:]
+        except asyncio.TimeoutError:
+            timeout_sec = max(1, int(round(_mode_timeout_sec(self._settings, mode))))
+            if mode in {"quick", "fast"}:
+                msg = (
+                    f"Quick mode hit {timeout_sec}s time budget before finishing. "
+                    "Try atlas-smart for a deeper answer."
+                )
+            elif mode == "smart":
+                msg = (
+                    f"Smart mode hit {timeout_sec}s time budget before finishing. "
+                    "Try atlas-genius or ask a narrower follow-up."
+                )
+            else:
+                msg = "I ran out of time before I could finish this answer."
+            await self._client.send_message(token, room_id, msg)
+            log.warning(
+                "matrix_answer_timeout",
+                extra={"extra": {"mode": mode, "seconds": timeout_sec}},
+            )
+        except Exception as exc:
+            log.warning(
+                "matrix_answer_failed",
+                extra={"extra": {"mode": mode, "error": str(exc)}},
+            )
+            await self._client.send_message(
+                token,
+                room_id,
+                "I hit an internal error while answering. Please retry, or switch to atlas-smart.",
+            )
         finally:
             stop.set()
             task.cancel()
@@ -213,3 +249,11 @@ def _extract_mode(body: str, mentions: tuple[str, ...], default_mode: str) -> tu
                 cleaned = cleaned.replace(tag.capitalize(), "")
             return mode, cleaned.strip()
     return ("", "")
+
+
+def _mode_timeout_sec(settings: Settings, mode: str) -> float:
+    if mode == "genius":
+        return settings.genius_time_budget_sec
+    if mode == "smart":
+        return settings.smart_time_budget_sec
+    return settings.quick_time_budget_sec