feat: implement punishment process management and integrate with Discord interactions

bun.lock

@@ -4,10 +4,12 @@
     "": {
       "name": "police-discord-bot",
       "dependencies": {
+        "@types/ws": "^8.18.1",
         "hono": "^4.9.4",
         "pino": "^9.9.0",
         "pino-pretty": "^13.1.1",
         "tweetnacl": "^1.0.3",
+        "ws": "^8.18.3",
         "zod": "^4.1.5",
       },
       "devDependencies": {
@@ -22,6 +24,8 @@
 
     "@types/react": ["@types/react@19.1.12", "", { "dependencies": { "csstype": "^3.0.2" } }, "sha512-cMoR+FoAf/Jyq6+Df2/Z41jISvGZZ2eTlnsaJRptmZ76Caldwy1odD4xTr/gNV9VLj0AWgg/nmkevIyUfIIq5w=="],
 
+    "@types/ws": ["@types/ws@8.18.1", "", { "dependencies": { "@types/node": "*" } }, "sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg=="],
+
     "atomic-sleep": ["atomic-sleep@1.0.0", "", {}, "sha512-kNOjDqAh7px0XWNI+4QbzoiR/nTkHAWNud2uvnJquD1/x5a7EQZMJT0AczqK0Qn67oY/TTQ1LbUKajZpp3I9tQ=="],
 
     "bun-types": ["bun-types@1.2.21", "", { "dependencies": { "@types/node": "*" }, "peerDependencies": { "@types/react": "^19" } }, "sha512-sa2Tj77Ijc/NTLS0/Odjq/qngmEPZfbfnOERi0KRUYhT9R8M4VBioWVmMWE5GrYbKMc+5lVybXygLdibHaqVqw=="],
@@ -86,6 +90,8 @@
 
     "wrappy": ["wrappy@1.0.2", "", {}, "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ=="],
 
+    "ws": ["ws@8.18.3", "", { "peerDependencies": { "bufferutil": "^4.0.1", "utf-8-validate": ">=5.0.2" }, "optionalPeers": ["bufferutil", "utf-8-validate"] }, "sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg=="],
+
     "zod": ["zod@4.1.5", "", {}, "sha512-rcUUZqlLJgBC33IT3PNMgsCq6TzLQEG/Ei/KTCU0PedSWRMAXoOUN+4t/0H+Q8bdnLPdqUYnvboJT0bn/229qg=="],
   }
 }

package.json

@@ -5,10 +5,12 @@
     "dev": "bun run --hot src/index.ts"
   },
   "dependencies": {
+    "@types/ws": "^8.18.1",
     "hono": "^4.9.4",
     "pino": "^9.9.0",
     "pino-pretty": "^13.1.1",
     "tweetnacl": "^1.0.3",
+    "ws": "^8.18.3",
     "zod": "^4.1.5"
   },
   "devDependencies": {

src/discord/index.ts

@@ -2,5 +2,6 @@ export * from "./application";
 export * from "./audit-log";
 export * from "./core";
 export * from "./guild";
+export * from "./interaction";
 export * from "./message";
 export * from "./poll";

src/discord/interaction.ts

@@ -0,0 +1,23 @@
+import { fetchDiscord } from "./core";
+import { Config } from "../config";
+
+export const editInteractionResponse = async (token: string, data: any) => {
+  return fetchDiscord(
+    `/webhooks/${Config.APP_ID}/${token}/messages/@original`,
+    {
+      method: "PATCH",
+      body: JSON.stringify(data),
+    },
+  );
+};
+
+export const sendInteractionCallback = async (
+  id: string,
+  token: string,
+  body: any,
+) => {
+  return fetchDiscord(`/interactions/${id}/${token}/callback`, {
+    method: "POST",
+    body: JSON.stringify(body),
+  });
+};

src/discord/message.ts

@@ -217,3 +217,18 @@ export const getReactions = async (
   const res = await fetchDiscord(path);
   return res.json();
 };
+
+export const listReactionUsersMulti = async (
+  channelId: string,
+  messageId: string,
+  emojis: string[],
+) => {
+  const results = await Promise.all(
+    emojis.map((e) => getReactions(channelId, messageId, e)),
+  );
+  const byId = new Map<string, any>();
+  for (const arr of results) {
+    for (const u of arr) byId.set(u.id, u); // de-dupe across variants
+  }
+  return Array.from(byId.values());
+};

src/gateway.ts

@@ -0,0 +1,261 @@
+import { WebSocket } from "ws";
+import { Config } from "./config";
+import { logger } from "./logger";
+import * as Discord from "./discord";
+import {
+  getActiveProcesses,
+  getPunishmentProcess,
+  updatePunishmentProcess,
+  endPunishmentProcess,
+  PunishmentState,
+} from "./store";
+
+let ws: WebSocket;
+let sessionId: string | null = null;
+let sequence: number | null = null;
+let resumeGatewayUrl: string | null = null;
+let heartbeatInterval: number | null = null;
+let botUserId: string | null = null;
+
+const INTENTS =
+  (1 << 0) | // GUILDS
+  (1 << 9) | // GUILD_MESSAGES
+  (1 << 10) | // GUILD_MESSAGE_REACTIONS
+  (1 << 15); // MESSAGE_CONTENT
+
+const send = (op: number, d: any) => {
+  if (ws.readyState === WebSocket.OPEN) {
+    ws.send(JSON.stringify({ op, d }));
+  }
+};
+
+const heartbeat = () => {
+  if (heartbeatInterval) {
+    setInterval(() => {
+      send(1, sequence);
+      logger.info("Sent heartbeat");
+    }, heartbeatInterval);
+  }
+};
+
+const identify = () => {
+  send(2, {
+    token: Config.BOT_TOKEN,
+    intents: INTENTS,
+    properties: {
+      os: "linux",
+      browser: "my_library",
+      device: "my_library",
+    },
+  });
+};
+
+const handleReactionAdd = async (data: any) => {
+  const { message_id, user_id, emoji, channel_id, guild_id } = data;
+  const process = getPunishmentProcess(message_id);
+
+  if (!process || process.chosen || user_id === botUserId) return;
+
+  const RENAME_EMOJIS = ["✏️", "✏"];
+  const ROLE_EMOJIS = ["😭"];
+
+  const isRenameReaction = RENAME_EMOJIS.includes(emoji.name);
+  const isRoleReaction = ROLE_EMOJIS.includes(emoji.name);
+
+  if (!isRenameReaction && !isRoleReaction) return;
+
+  const [renameUsers, roleUsers] = await Promise.all([
+    Discord.listReactionUsersMulti(channel_id, message_id, RENAME_EMOJIS),
+    Discord.listReactionUsersMulti(channel_id, message_id, ROLE_EMOJIS),
+  ]);
+
+  const DECIDER_USER_ID = "311380871901085707";
+  const deciderRename = renameUsers.some((u) => u.id === DECIDER_USER_ID);
+  const deciderRole = roleUsers.some((u) => u.id === DECIDER_USER_ID);
+
+  let chosen: "rename" | "role" | null = null;
+  if (deciderRename) chosen = "rename";
+  else if (deciderRole) chosen = "role";
+  else if (renameUsers.filter((u) => u.id !== botUserId).length >= 2)
+    chosen = "rename";
+  else if (roleUsers.filter((u) => u.id !== botUserId).length >= 2)
+    chosen = "role";
+
+  if (chosen) {
+    updatePunishmentProcess(message_id, { chosen });
+    executePunishment(chosen, process, message_id);
+  }
+};
+
+const handleMessageCreate = async (data: any) => {
+  const { channel_id, author, content, guild_id } = data;
+
+  if (author.bot) return;
+
+  for (const [messageId, process] of Array.from(getActiveProcesses())) {
+    if (
+      process.channelId === channel_id &&
+      process.askMessageId &&
+      !process.chosen
+    ) {
+      if (process.renameDeadline && Date.now() > process.renameDeadline) {
+        Discord.createMessage(channel_id, {
+          content: "หมดเวลาการลงชื่อใหม่ ไม่เปลี่ยนแม่งละชื่อ",
+        });
+        Discord.editInteractionResponse(process.interactionToken, {
+          content: "No decision taken.",
+        });
+        endPunishmentProcess(messageId);
+        continue;
+      }
+
+      const newNick = content.trim().slice(0, 32);
+      if (newNick) {
+        endPunishmentProcess(messageId);
+        try {
+          await Discord.modifyGuildMember(guild_id, process.executorId, {
+            nick: newNick,
+          });
+          await Discord.createMessage(channel_id, {
+            content: `ตั้งชื่อใหม่ให้ <@${process.executorId}> เป็น **${newNick}** แล้วนะคราฟ`,
+          });
+          await Discord.editInteractionResponse(process.interactionToken, {
+            content: "Done.",
+          });
+        } catch (error) {
+          logger.error(error, "Failed to execute rename punishment");
+          await Discord.editInteractionResponse(process.interactionToken, {
+            content: "Failed to execute punishment.",
+          });
+        }
+      }
+    }
+  }
+};
+
+const executePunishment = async (
+  chosen: "rename" | "role",
+  process: PunishmentState,
+  messageId: string,
+) => {
+  if (chosen === "rename") {
+    const ask = await Discord.createMessage(process.channelId, {
+      content: `เลือก: เปลี่ยนชื่อเล่นให้ <@${process.executorId}>\nพิมพ์ชื่อใหม่ (≤32 ตัวอักษร) ภายใน ${Config.RENAME_TIMEOUT_SEC} วินาที`,
+    });
+    updatePunishmentProcess(messageId, {
+      askMessageId: ask.id,
+      renameDeadline: Date.now() + Config.RENAME_TIMEOUT_SEC * 1000,
+    });
+  } else if (chosen === "role") {
+    endPunishmentProcess(messageId);
+    try {
+      if (!Config.PUNISH_ROLE_ID) {
+        await Discord.createMessage(process.channelId, {
+          content: "ยังไม่ได้ตั้งค่า PUNISH_ROLE_ID จึงลบ role ไม่ได้",
+        });
+      } else {
+        await Discord.removeGuildMemberRole(
+          process.guildId,
+          process.executorId,
+          Config.PUNISH_ROLE_ID,
+        );
+        await Discord.createMessage(process.channelId, {
+          content: `ลบ role ออกจาก <@${process.executorId}> แล้วนะคราฟ`,
+        });
+      }
+      await Discord.createMessage(process.channelId, {
+        content: "โดนซะบ้าง 🙂",
+      });
+      await Discord.editInteractionResponse(process.interactionToken, {
+        content: "Done.",
+      });
+    } catch (error) {
+      logger.error(error, "Failed to execute role punishment");
+      await Discord.editInteractionResponse(process.interactionToken, {
+        content: "Failed to execute punishment.",
+      });
+    }
+  }
+};
+
+setInterval(() => {
+  const now = Date.now();
+  for (const [messageId, process] of Array.from(getActiveProcesses())) {
+    if (now > process.deadline && !process.chosen) {
+      Discord.createMessage(process.channelId, {
+        content: "หมดเวลา ไม่มีการเลือกลงโทษ รอดตัวไป",
+      });
+      Discord.editInteractionResponse(process.interactionToken, {
+        content: "No decision taken.",
+      });
+      endPunishmentProcess(messageId);
+    }
+  }
+}, 5000);
+
+const handleEvent = (payload: any) => {
+  const { op, d, s, t } = payload;
+
+  if (s) {
+    sequence = s;
+  }
+
+  switch (op) {
+    case 0: // Dispatch
+      logger.info({ type: t }, "Received dispatch event");
+      if (t === "READY") {
+        botUserId = d.user.id;
+        sessionId = d.session_id;
+        resumeGatewayUrl = d.resume_gateway_url;
+      }
+      if (t === "MESSAGE_REACTION_ADD") {
+        handleReactionAdd(d);
+      }
+      if (t === "MESSAGE_CREATE") {
+        handleMessageCreate(d);
+      }
+      break;
+    case 10: // Hello
+      heartbeatInterval = d.heartbeat_interval;
+      heartbeat();
+      identify();
+      break;
+    case 11: // Heartbeat ACK
+      logger.info("Received heartbeat ACK");
+      break;
+    default:
+      logger.info({ op, d, s, t }, "Received unhandled opcode");
+      break;
+  }
+};
+
+export const connect = async () => {
+  try {
+    const res = await Discord.fetchDiscord("/gateway/bot");
+    const data = await res.json();
+    const gatewayUrl = data.url;
+    resumeGatewayUrl = gatewayUrl;
+
+    ws = new WebSocket(`${gatewayUrl}/?v=10&encoding=json`);
+
+    ws.on("open", () => {
+      logger.info("Connected to Gateway");
+    });
+
+    ws.on("message", (data) => {
+      const payload = JSON.parse(data.toString());
+      handleEvent(payload);
+    });
+
+    ws.on("close", (code) => {
+      logger.warn({ code }, "Gateway connection closed. Reconnecting...");
+      setTimeout(connect, 5000);
+    });
+
+    ws.on("error", (err) => {
+      logger.error(err, "Gateway error");
+    });
+  } catch (error) {
+    logger.error(error, "Failed to connect to Gateway");
+  }
+};

src/index.ts

@@ -4,6 +4,8 @@ import { logger } from "./logger";
 import * as nacl from "tweetnacl";
 import * as Discord from "./discord";
 import { registerCommand } from "./commands";
+import { connect } from "./gateway";
+import { startPunishmentProcess } from "./store";
 
 logger.info(
   { port: Config.PORT, guildId: Config.DEFAULT_GUILD_ID },
@@ -13,27 +15,11 @@ logger.info(
 registerCommand().catch((err) =>
   logger.error(err, "Failed to register commands"),
 );
+connect();
 
 const snowflakeToMs = (id: string) =>
   Number((BigInt(id) >> BigInt(22)) + BigInt(1420070400000));
 
-async function interactionCallback(id: string, token: string, body: any) {
-  return Discord.fetchDiscord(`/interactions/${id}/${token}/callback`, {
-    method: "POST",
-    body: JSON.stringify(body),
-  });
-}
-
-async function editOriginal(token: string, content: string) {
-  return Discord.fetchDiscord(
-    `/webhooks/${Config.APP_ID}/${token}/messages/@original`,
-    {
-      method: "PATCH",
-      body: JSON.stringify({ content }),
-    },
-  );
-}
-
 function findRecentVoicePunish(audit: Discord.AuditLog): {
   entry: Discord.AuditLogEntry;
   kind: "deafen" | "mute";
@@ -52,21 +38,6 @@ function findRecentVoicePunish(audit: Discord.AuditLog): {
   return null;
 }
 
-async function listReactionUsersMulti(
-  channelId: string,
-  messageId: string,
-  emojis: string[],
-) {
-  const results = await Promise.all(
-    emojis.map((e) => Discord.getReactions(channelId, messageId, e)),
-  );
-  const byId = new Map<string, any>();
-  for (const arr of results) {
-    for (const u of arr) byId.set(u.id, u); // de-dupe across variants
-  }
-  return Array.from(byId.values());
-}
-
 const randomSystem = () => {
   const pool = ["คุณเป็นคนพูดจาเกรียนๆ"];
   return pool[Math.floor(Math.random() * pool.length)];
@@ -132,7 +103,7 @@ app.post("/interactions", async (c) => {
   }
 
   const interaction = JSON.parse(raw);
-  logger.info("Received interaction");
+  logger.info({ interaction }, "Received interaction");
   const { type } = interaction;
 
   if (type === 1) {
@@ -145,7 +116,10 @@ app.post("/interactions", async (c) => {
     const commandName = data.name;
 
     if (commandName === "checkaudit") {
-      await interactionCallback(id, token, { type: 5, data: { flags: 64 } });
+      await Discord.sendInteractionCallback(id, token, {
+        type: 5,
+        data: { flags: 64 },
+      });
 
       queueMicrotask(async () => {
         try {
@@ -157,7 +131,9 @@ app.post("/interactions", async (c) => {
 
           const hit = findRecentVoicePunish(auditLogs);
           if (!hit) {
-            await editOriginal(token, "No recent mute/deafen found.");
+            await Discord.editInteractionResponse(token, {
+              content: "No recent mute/deafen found.",
+            });
             return;
           }
 
@@ -179,94 +155,23 @@ app.post("/interactions", async (c) => {
           for (const e of ROLE_EMOJIS)
             await Discord.addReaction(channel_id, sent.id, e);
 
-          const botUserId = String(application_id);
+          startPunishmentProcess(sent.id, {
-          const deadline = Date.now() + Config.REACTION_TIMEOUT_SEC * 1000;
+            interactionToken: token,
-          let chosen: "rename" | "role" | null = null;
+            channelId: channel_id,
+            guildId: guild_id,
+            executorId: executorId,
+            timeout: Config.REACTION_TIMEOUT_SEC,
+          });
 
-          while (Date.now() < deadline && !chosen) {
+          await Discord.editInteractionResponse(token, {
-            const [renameUsers, roleUsers] = await Promise.all([
+            content:
-              listReactionUsersMulti(channel_id, sent.id, RENAME_EMOJIS),
+              "Punishment process initiated. Please react to the message above.",
-              listReactionUsersMulti(channel_id, sent.id, ROLE_EMOJIS),
+          });
-            ]);
-
-            const DECIDER_USER_ID = "311380871901085707";
-            const deciderRename = renameUsers.some(
-              (u) => u.id === DECIDER_USER_ID,
-            );
-            const deciderRole = roleUsers.some((u) => u.id === DECIDER_USER_ID);
-
-            if (deciderRename) chosen = "rename";
-            else if (deciderRole) chosen = "role";
-            else if (renameUsers.filter((u) => u.id !== botUserId).length >= 2)
-              chosen = "rename";
-            else if (roleUsers.filter((u) => u.id !== botUserId).length >= 2)
-              chosen = "role";
-
-            if (!chosen) await new Promise((r) => setTimeout(r, 2000));
-          }
-
-          if (!chosen) {
-            await Discord.createMessage(channel_id, {
-              content: "หมดเวลา ไม่มีการเลือกลงโทษ รอดตัวไป",
-            });
-            await editOriginal(token, "No decision taken.");
-            return;
-          }
-
-          if (chosen === "rename") {
-            const ask = await Discord.createMessage(channel_id, {
-              content: `เลือก: เปลี่ยนชื่อเล่นให้ <@${executorId}>\nพิมพ์ชื่อใหม่ (≤32 ตัวอักษร) ภายใน ${Config.RENAME_TIMEOUT_SEC} วินาที`,
-            });
-            const renameDeadline =
-              Date.now() + Config.RENAME_TIMEOUT_SEC * 1000;
-            let newNick: string | null = null;
-
-            while (Date.now() < renameDeadline && !newNick) {
-              const newMessages = await Discord.getChannelMessages(channel_id, {
-                after: ask.id,
-              });
-              const firstUserMsg = newMessages.find(
-                (m) => !m.author.bot && m.content?.trim().length > 0,
-              );
-              if (firstUserMsg)
-                newNick = firstUserMsg.content.trim().slice(0, 32);
-              if (!newNick) await new Promise((r) => setTimeout(r, 2000));
-            }
-
-            if (!newNick) {
-              await Discord.createMessage(channel_id, {
-                content: "หมดเวลาการลงชื่อใหม่ ไม่เปลี่ยนแม่งละชื่อ",
-              });
-            } else {
-              await Discord.modifyGuildMember(guild_id, executorId, {
-                nick: newNick,
-              });
-              await Discord.createMessage(channel_id, {
-                content: `ตั้งชื่อใหม่ให้ <@${executorId}> เป็น **${newNick}** แล้วนะคราฟ`,
-              });
-            }
-          } else if (chosen === "role") {
-            if (!Config.PUNISH_ROLE_ID) {
-              await Discord.createMessage(channel_id, {
-                content: "ยังไม่ได้ตั้งค่า PUNISH_ROLE_ID จึงลบ role ไม่ได้",
-              });
-            } else {
-              await Discord.removeGuildMemberRole(
-                guild_id,
-                executorId,
-                Config.PUNISH_ROLE_ID,
-              );
-              await Discord.createMessage(channel_id, {
-                content: `ลบ role ออกจาก <@${executorId}> แล้วนะคราฟ`,
-              });
-            }
-          }
-
-          await Discord.createMessage(channel_id, { content: "โดนซะบ้าง 🙂" });
-          await editOriginal(token, "Done.");
         } catch (error) {
           logger.error(error, "Error in checkaudit command");
-          await editOriginal(token, "An error occurred.");
+          await Discord.editInteractionResponse(token, {
+            content: "An error occurred.",
+          });
         }
       });
 
@@ -276,22 +181,27 @@ app.post("/interactions", async (c) => {
     if (commandName === "talk") {
       const prompt =
         data.options?.find((o: any) => o.name === "prompt")?.value || "";
-      await interactionCallback(id, token, { type: 5, data: { flags: 64 } });
+      await Discord.sendInteractionCallback(id, token, {
+        type: 5,
+        data: { flags: 64 },
+      });
 
       queueMicrotask(async () => {
         try {
           const reply = await callOpenRouterChat(prompt);
-          await editOriginal(token, reply);
+          await Discord.editInteractionResponse(token, { content: reply });
         } catch (error) {
           logger.error(error, "Error in talk command");
-          await editOriginal(token, "An error occurred.");
+          await Discord.editInteractionResponse(token, {
+            content: "An error occurred.",
+          });
         }
       });
 
       return c.body(null, 204);
     }
 
-    await interactionCallback(id, token, {
+    await Discord.sendInteractionCallback(id, token, {
       type: 4,
       data: { content: "Unknown command", flags: 64 },
     });

src/store.ts

@@ -0,0 +1,42 @@
+export interface PunishmentState {
+  interactionToken: string;
+  channelId: string;
+  guildId: string;
+  executorId: string;
+  deadline: number;
+  chosen: "rename" | "role" | null;
+  askMessageId?: string;
+  renameDeadline?: number;
+}
+
+const punishmentProcesses = new Map<string, PunishmentState>(); // Key: messageId
+
+export const startPunishmentProcess = (
+  messageId: string,
+  state: Omit<PunishmentState, "chosen" | "deadline"> & { timeout: number },
+) => {
+  const deadline = Date.now() + state.timeout * 1000;
+  punishmentProcesses.set(messageId, { ...state, chosen: null, deadline });
+};
+
+export const getPunishmentProcess = (messageId: string) => {
+  return punishmentProcesses.get(messageId);
+};
+
+export const updatePunishmentProcess = (
+  messageId: string,
+  updates: Partial<PunishmentState>,
+) => {
+  const process = punishmentProcesses.get(messageId);
+  if (process) {
+    punishmentProcesses.set(messageId, { ...process, ...updates });
+  }
+};
+
+export const endPunishmentProcess = (messageId: string) => {
+  punishmentProcesses.delete(messageId);
+};
+
+export const getActiveProcesses = () => {
+  return punishmentProcesses.entries();
+};