From e9da7b493cfab6ef381762025cf972bd6fd4c485 Mon Sep 17 00:00:00 2001
From: Shigma <1700011071@pku.edu.cn>
Date: Thu, 3 Feb 2022 17:19:30 +0800
Subject: [PATCH] fix(core): support bot reconnect

---
 docs/api/core/context.md       |  4 ++--
 packages/core/src/adapter.ts   | 14 +++++++++-----
 packages/core/src/bot.ts       |  9 +++++++--
 packages/koishi/src/adapter.ts |  8 ++++++--
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/docs/api/core/context.md b/docs/api/core/context.md
index 100129dfcf..f5a0d02d18 100644
--- a/docs/api/core/context.md
+++ b/docs/api/core/context.md
@@ -41,9 +41,9 @@ ctx.router.get('/path', (ctx, next) => {
 
 使用 sid 获取机器人实例。
 
-#### ctx.bots.remove(sid)
+#### ctx.bots.remove(id)
 
-- **sid:** `string` 机器人的 sid
+- **sid:** `string` 机器人的 id
 - 返回值: `boolean` 机器人实例是否存在
 
 移除一个机器人实例。
diff --git a/packages/core/src/adapter.ts b/packages/core/src/adapter.ts
index 385d699779..50cf7d620c 100644
--- a/packages/core/src/adapter.ts
+++ b/packages/core/src/adapter.ts
@@ -1,4 +1,4 @@
-import { Logger, paramCase, Dict, Awaitable } from '@koishijs/utils'
+import { Logger, paramCase, Dict, Awaitable, remove } from '@koishijs/utils'
 import { Session } from './session'
 import { App } from './app'
 import { Bot } from './bot'
@@ -51,7 +51,7 @@ export namespace Adapter {
   export type BotConfig<R> = R & { bots?: R[] }
   export type PluginConfig<S = any, R = any> = S & BotConfig<R>
 
-  function join(platform: string, protocol: string) {
+  export function join(platform: string, protocol: string) {
     return protocol ? `${platform}.${protocol}` : platform
   }
 
@@ -73,6 +73,7 @@ export namespace Adapter {
   export function define(platform: string, constructor: Bot.Constructor, ...args: CreatePluginRestParams) {
     const name = platform + '-adapter'
     platform = platform.toLowerCase()
+    Bot.library[platform] = constructor
 
     let BotConfig: Schema
     if (typeof args[0] === 'function') {
@@ -144,7 +145,8 @@ export namespace Adapter {
       return this.find(bot => bot.sid === sid)
     }
 
-    create<T extends Bot>(platform: string, options: any, constructor: new (adapter: Adapter, config: any) => T): T {
+    create<T extends Bot>(platform: string, options: any, constructor?: new (adapter: Adapter, config: any) => T): T {
+      constructor ||= Bot.library[platform] as any
       const adapter = this.resolve(platform, options)
       const bot = new constructor(adapter, options)
       adapter.bots.push(bot)
@@ -153,10 +155,12 @@ export namespace Adapter {
       return bot
     }
 
-    async remove(sid: string) {
-      const index = this.findIndex(bot => bot.sid === sid)
+    async remove(id: string) {
+      const index = this.findIndex(bot => bot.id === id)
       if (index < 0) return
       const [bot] = this.splice(index, 1)
+      remove(bot.adapter.bots, bot)
+      bot.config.disabled = true
       this.app.emit('bot-removed', bot)
       return bot.stop()
     }
diff --git a/packages/core/src/bot.ts b/packages/core/src/bot.ts
index 49adedee1e..d28eca368f 100644
--- a/packages/core/src/bot.ts
+++ b/packages/core/src/bot.ts
@@ -1,4 +1,4 @@
-import { Logger, Random, sleep } from '@koishijs/utils'
+import { Dict, Logger, Random, sleep } from '@koishijs/utils'
 import { Adapter } from './adapter'
 import { App } from './app'
 import { Session } from './session'
@@ -63,8 +63,10 @@ export abstract class Bot<T extends Bot.BaseConfig = Bot.BaseConfig> {
   }
 
   async start() {
+    if (this.config.disabled) return
+    if (['connect', 'reconnect', 'online'].includes(this.status)) return
+    this.status = 'connect'
     try {
-      this.status = 'connect'
       await this.app.parallel('bot-connect', this)
       await this.adapter.connect(this)
     } catch (error) {
@@ -73,6 +75,7 @@ export abstract class Bot<T extends Bot.BaseConfig = Bot.BaseConfig> {
   }
 
   async stop() {
+    if (['disconnect', 'offline'].includes(this.status)) return
     this.status = 'disconnect'
     try {
       await this.app.parallel('bot-disconnect', this)
@@ -124,6 +127,8 @@ export abstract class Bot<T extends Bot.BaseConfig = Bot.BaseConfig> {
 }
 
 export namespace Bot {
+  export const library: Dict<Bot.Constructor> = {}
+
   export interface BaseConfig {
     disabled?: boolean
     protocol?: string
diff --git a/packages/koishi/src/adapter.ts b/packages/koishi/src/adapter.ts
index beb6d0be31..f76d38326e 100644
--- a/packages/koishi/src/adapter.ts
+++ b/packages/koishi/src/adapter.ts
@@ -55,7 +55,7 @@ export namespace InjectedAdapter {
         socket.on('close', (code, reason) => {
           bot.socket = null
           logger.debug(`websocket closed with ${code}`)
-          if (!this.isListening) return bot.status = 'offline'
+          if (!this.isListening || bot.config.disabled) return bot.status = 'offline'
 
           // remove query args to protect privacy
           const message = reason.toString() || `failed to connect to ${url}`
@@ -72,7 +72,7 @@ export namespace InjectedAdapter {
           bot.status = 'reconnect'
           logger.warn(`${message}, will retry in ${Time.formatTimeShort(timeout)}...`)
           setTimeout(() => {
-            if (this.isListening) reconnect()
+            if (this.isListening && !bot.config.disabled) reconnect()
           }, timeout)
         })
 
@@ -87,6 +87,10 @@ export namespace InjectedAdapter {
       reconnect(true)
     }
 
+    disconnect(bot: Bot) {
+      bot.socket?.close()
+    }
+
     start() {
       this.isListening = true
     }