From c3f2a5dfe1551743b961abf5033b00139003b03a Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Mon, 24 Aug 2020 23:29:09 +0300
Subject: [PATCH] Add some disconnection handling and a method to set the last
 message ID map

---
 puppet/src/client.js | 21 ++++++++++++++-------
 puppet/src/puppet.js | 33 +++++++++++++++++++++++++++++----
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/puppet/src/client.js b/puppet/src/client.js
index 071ccbe..21da40e 100644
--- a/puppet/src/client.js
+++ b/puppet/src/client.js
@@ -115,14 +115,20 @@ export default class Client {
 	}
 
 	handleStart = async (req) => {
-		if (this.puppet !== null) {
-			return { started: false, is_logged_in: await this.puppet.isLoggedIn() }
+		let started = false
+		if (this.puppet === null) {
+			this.log("Opening new puppeteer for", this.userID)
+			this.puppet = new MessagesPuppeteer(this.userID, this)
+			this.manager.puppets.set(this.userID, this.puppet)
+			await this.puppet.start(!!req.debug)
+			started = true
+		}
+		return {
+			started,
+			is_logged_in: await this.puppet.isLoggedIn(),
+			is_connected: !await this.puppet.isDisconnected(),
+			is_permanently_disconnected: await this.puppet.isPermanentlyDisconnected(),
 		}
-		this.log("Opening new puppeteer for", this.userID)
-		this.puppet = new MessagesPuppeteer(this.userID, this)
-		this.manager.puppets.set(this.userID, this.puppet)
-		await this.puppet.start(!!req.debug)
-		return { started: true, is_logged_in: await this.puppet.isLoggedIn() }
 	}
 
 	handleStop = async () => {
@@ -199,6 +205,7 @@ export default class Client {
 				disconnect: () => this.stop(),
 				login: () => this.puppet.waitForLogin(),
 				send: req => this.puppet.sendMessage(req.chat_id, req.text),
+				set_last_message_ids: req => this.puppet.setLastMessageIDs(req.msg_ids),
 				get_chats: () => this.puppet.getRecentChats(),
 				get_chat: req => this.puppet.getChatInfo(req.chat_id),
 				get_messages: req => this.puppet.getMessages(req.chat_id),
diff --git a/puppet/src/puppet.js b/puppet/src/puppet.js
index e281ba0..c8cd174 100644
--- a/puppet/src/puppet.js
+++ b/puppet/src/puppet.js
@@ -91,9 +91,11 @@ export default class MessagesPuppeteer {
 		// Wait for the page to load (either QR code for login or chat list when already logged in)
 		await Promise.race([
 			this.page.waitForSelector("mw-main-container mws-conversations-list .conv-container",
-				{ visible: true }),
+				{ visible: true, timeout: 60000 }),
 			this.page.waitForSelector("mw-authentication-container mw-qr-code",
-				{ visible: true }),
+				{ visible: true, timeout: 60000 }),
+			this.page.waitForSelector("mw-unable-to-connect-container",
+				{ visible: true, timeout: 60000 }),
 		])
 		this.taskQueue.start()
 		if (await this.isLoggedIn()) {
@@ -153,6 +155,15 @@ export default class MessagesPuppeteer {
 		return await this.page.$("mw-main-container mws-conversations-list") !== null
 	}
 
+	async isPermanentlyDisconnected() {
+		return await this.page.$("mw-unable-to-connect-container") !== null
+	}
+
+	async isDisconnected() {
+		// TODO we should observe this banner appearing and disappearing to notify about disconnects
+		return await this.page.$("mw-main-container mw-error-banner") !== null
+	}
+
 	/**
 	 * Get the IDs of the most recent chats.
 	 *
@@ -196,7 +207,20 @@ export default class MessagesPuppeteer {
 	 * @return {Promise<[MessageData]>} - The messages visible in the chat.
 	 */
 	async getMessages(id) {
-		return this.taskQueue.push(() => this._getMessagesUnsafe(id))
+		return this.taskQueue.push(async () => {
+			const messages = await this._getMessagesUnsafe(id)
+			if (messages.length > 0) {
+				this.mostRecentMessages.set(id, messages[messages.length - 1].id)
+			}
+			return messages
+		})
+	}
+
+	setLastMessageIDs(ids) {
+		for (const [chatID, messageID] of Object.entries(ids)) {
+			this.mostRecentMessages.set(+chatID, messageID)
+		}
+		this.log("Updated most recent message ID map:", this.mostRecentMessages)
 	}
 
 	async startObserving() {
@@ -272,7 +296,8 @@ export default class MessagesPuppeteer {
 		const newFirstID = messages[0].id
 		const newLastID = messages[messages.length - 1].id
 		this.mostRecentMessages.set(id, newLastID)
-		this.log(`Loaded messages in ${id} after ${lastMsgID}: got ${newFirstID}-${newLastID}`)
+		const range = newFirstID === newLastID ? newFirstID : `${newFirstID}-${newLastID}`
+		this.log(`Loaded ${messages.length} messages in ${id} after ${lastMsgID}: got ${range}`)
 
 		if (this.client) {
 			for (const message of messages) {