Examples¶

"""Minimal TEBHOS bot that answers "!ping" with "pong"."""
import asyncio
import json
import os
import random
import sys

import aiohttp
import websockets

TOKEN = os.environ.get("TEBHOS_BOT_TOKEN")  # "{application_id}.{secret}"
API = os.environ.get("TEBHOS_API", "https://tebhos.com/api/v1")

# Close codes after which reconnecting cannot help.
FATAL_CLOSE_CODES = {4004, 4010, 4011, 4012}
# Close codes after which the old session cannot be resumed.
FRESH_SESSION_CLOSE_CODES = {4001, 4002, 4003, 4005, 4007}


class Bot:
    def __init__(self, http: aiohttp.ClientSession):
        self.http = http
        self.session_id = None
        self.seq = None
        self.user_id = None
        self.acked = True
        self.tasks = set()

    async def api(self, method, path, body=None):
        """REST helper: bot auth header, JSON body, waits and retries on 429."""
        headers = {"Authorization": f"Bot {TOKEN}"}
        while True:
            async with self.http.request(method, API + path, json=body, headers=headers) as res:
                if res.status == 429:
                    try:
                        data = await res.json(content_type=None)
                        seconds = float(data.get("retry_after", 1))
                    except (ValueError, TypeError, AttributeError):
                        seconds = float(res.headers.get("Retry-After", 1))
                    await asyncio.sleep(seconds)
                    continue
                if res.status >= 400:
                    raise RuntimeError(f"{method} {path} failed: {res.status} {await res.text()}")
                return None if res.status == 204 else await res.json()

    async def send(self, ws, op, d):
        await ws.send(json.dumps({"op": op, "d": d}))

    async def heartbeat(self, ws, interval):
        await asyncio.sleep(interval * random.random())
        while True:
            if not self.acked:
                # No ACK since the last heartbeat: the connection is dead. Drop it and resume.
                await ws.close(code=4000)
                return
            self.acked = False
            await self.send(ws, 1, self.seq)
            await asyncio.sleep(interval)

    async def handle_message(self, message):
        author = message.get("author") or {}
        # Never answer yourself or other bots.
        if author.get("id") == self.user_id or author.get("bot"):
            return
        if message.get("content", "").strip() != "!ping":
            return
        await self.api("POST", f"/channels/{message['channel_id']}/messages", {
            "content": "pong",
            "message_reference": {"message_id": message["id"]},
        })

    def spawn(self, coro):
        task = asyncio.create_task(coro)
        self.tasks.add(task)
        task.add_done_callback(self.tasks.discard)

    def handle_dispatch(self, event, data):
        if event == "READY":
            self.session_id = data["session_id"]
            self.user_id = data["user"]["id"]
            print(f"Ready as {data['user']['username']} ({self.user_id})")
        elif event == "RESUMED":
            print("Session resumed")
        elif event == "GUILD_CREATE":
            print(f"Guild available: {data['id']}")
        elif event == "MESSAGE_CREATE":
            self.spawn(self.handle_message(data))

    async def connect_once(self, url):
        """Runs one gateway connection. Returns the close code, if any."""
        heartbeat_task = None
        try:
            async with websockets.connect(url, max_size=None) as ws:
                async for raw in ws:
                    payload = json.loads(raw)
                    op, d = payload["op"], payload.get("d")
                    if payload.get("s") is not None:
                        self.seq = payload["s"]

                    if op == 10:  # Hello
                        self.acked = True
                        heartbeat_task = asyncio.create_task(
                            self.heartbeat(ws, d["heartbeat_interval"] / 1000))
                        if self.session_id and self.seq is not None:
                            await self.send(ws, 6, {  # Resume
                                "token": TOKEN, "session_id": self.session_id, "seq": self.seq})
                        else:
                            await self.send(ws, 2, {  # Identify: the bare token, no "Bot " prefix
                                "token": TOKEN,
                                "properties": {"os": sys.platform, "browser": "example-bot",
                                               "device": "example-bot"},
                            })
                    elif op == 11:  # Heartbeat ACK
                        self.acked = True
                    elif op == 1:  # The gateway asks for a heartbeat now
                        await self.send(ws, 1, self.seq)
                    elif op == 7:  # Reconnect: drop this connection and resume on a new one
                        await ws.close(code=4000)
                    elif op == 9:  # Invalid Session: start over with a fresh Identify
                        self.session_id, self.seq = None, None
                        await asyncio.sleep(random.uniform(1, 5))
                        await ws.close(code=4000)
                    elif op == 0:  # Dispatch
                        self.handle_dispatch(payload.get("t"), d)
                return ws.close_code
        except websockets.ConnectionClosed as closed:
            return closed.rcvd.code if closed.rcvd else None
        except OSError as error:
            print(f"Connection failed: {error}")
            return None
        finally:
            if heartbeat_task:
                heartbeat_task.cancel()

    async def run(self):
        gateway = await self.api("GET", "/gateway/bot")
        url = gateway["url"] + "?v=1&encoding=json"
        while True:
            code = await self.connect_once(url)
            if code in FATAL_CLOSE_CODES:
                sys.exit(f"Gateway closed with {code}; not reconnecting.")
            if code in FRESH_SESSION_CLOSE_CODES:
                self.session_id, self.seq = None, None
            print(f"Disconnected ({code}); reconnecting in 2 s")
            await asyncio.sleep(2)


async def main():
    if not TOKEN:
        sys.exit("Set TEBHOS_BOT_TOKEN first.")
    async with aiohttp.ClientSession() as http:
        await Bot(http).run()


if __name__ == "__main__":
    asyncio.run(main())

async def cmd_ping(message, args):
    return {"content": "pong"}

async def cmd_echo(message, args):
    return {"content": " ".join(args) or "Nothing to echo."}

COMMANDS = {"ping": cmd_ping, "echo": cmd_echo}

# inside class Bot:
async def handle_message(self, message):
    author = message.get("author") or {}
    if author.get("id") == self.user_id or author.get("bot"):
        return
    content = message.get("content", "").strip()
    if not content.startswith("!"):
        return
    name, *args = content[1:].split() or [""]
    command = COMMANDS.get(name.lower())
    if command is None:
        return
    reply = await command(message, args)
    reply["message_reference"] = {"message_id": message["id"]}
    await self.api("POST", f"/channels/{message['channel_id']}/messages", reply)