nikita/2026-matrix-csonac
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved net bot implementation

Browse Source
This commit is contained in:
Nikita Tyukalov
2026-08-18 02:21:17 +03:00
parent 8c056c1da3
commit 75388d8693
2 changed files with 106 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
main.py
View File

@@ -5,11 +5,16 @@
import asyncio
import traceback
import signal
from pathlib import Path
import config
import util
import bot
import logic
#import bot
#import logic
from new_bot import MatrixBot
from bot_types import MatrixBotConfig
from datatypes import AppConfig
async def main() -> None:
@@ -29,18 +34,17 @@ async def main() -> None:
if cfg is None:
util.log_error("Could't load config")
return
# start the bot
if not await bot.start(cfg):
util.log_error("Could't start the bot")
return
await logic.setup(bot.get_client())
# wait for stop
matrix_bot_config = MatrixBotConfig(
matrix_homeserver_url=cfg.matrix_homeserver,
matrix_username_localpart=cfg.matrix_user,
storage_directory=Path(cfg.store_dir)
)
bot = MatrixBot(matrix_bot_config)
bot.start()
await util.get_app_stop_event().wait()
# stop
await logic.stop()
await bot.stop()
if __name__ == "__main__":

110
new_bot.py
View File

@@ -2,6 +2,7 @@ import asyncio
import aiofiles
import aioconsole
import traceback
import logging
import time
import json
import os
@@ -93,6 +94,31 @@ class MatrixBot:
localpart = config.matrix_username_localpart
return f"@{localpart}:{homeserver_name}"
@staticmethod
async def _wait_for_task_and_stop_event(payload_task: asyncio.Task, stop_wait_task: asyncio.Task, print_exc: bool = False) -> Any:
"""Waits for `payload_task` to finish. Uses `stop_wait_task` to check if payload task
should be cancelled. Returns value returned by `payload_task` task. Raises exception
raised by `payload_task` task.
If `stop_wait_task` finishes, then `asyncio.CancelledError` is raised.
"""
try:
done, _ = await asyncio.wait(
[payload_task, stop_wait_task],
return_when=asyncio.FIRST_COMPLETED
)
if stop_wait_task in done:
payload_task.cancel()
await payload_task
raise asyncio.CancelledError()
return payload_task.result()
except asyncio.CancelledError:
raise
except:
if print_exc:
traceback.print_exc()
raise
@staticmethod
async def _default_password_callback() -> str:
"""Gets password from `MATRIX_PASSWORD` envvar if it is set. Asks
@@ -101,13 +127,14 @@ class MatrixBot:
return os.environ["MATRIX_PASSWORD"]
print("--- A password is required (btw you might use MATRIX_PASSWORD envvar) ---")
return await aioconsole.ainput("Password: ")
async def _write_next_batch(self, next_batch: str) -> None:
"""Writes `next_batch` value to disk."""
path = self._config.storage_directory / "next_batch"
async with aiofiles.open(path, "w") as f:
await f.write(next_batch)
self._logger.debug("next_batch value is written to the disk")
async def _read_next_batch(self) -> str | None:
"""Reads `next_batch` value from disk. Returns None if file does not exist."""
@@ -126,6 +153,7 @@ class MatrixBot:
}
async with aiofiles.open(path, "w") as f:
await f.write(json.dumps(data, indent=4))
self._logger.debug("Session data is writter to the disk")
async def _read_session_data(self) -> dict[str, Any] | None:
"""Read session data from disk."""
@@ -143,10 +171,11 @@ class MatrixBot:
"""This callback is called when AsyncClient syncs with the server"""
current_time = time.time()
delta_time = current_time - self._last_next_batch_dump
self._last_next_batch = response.next_batch
if delta_time >= self.NEXT_BATCH_DUMP_PERIOD:
self._last_next_batch_dump = current_time
try:
await self._write_next_batch(response.next_batch)
await self._write_next_batch(self._last_next_batch)
except:
traceback.print_exc()
@@ -159,6 +188,7 @@ class MatrixBot:
async def _client_login_session_data(self, session_data: dict[str, Any]) -> None:
"""Login using session data. Raises and exception on failure."""
self._logger.debug("Using stored session data to log in")
# build user id
username = self._build_matrix_username(self._config)
self._client.restore_login(
@@ -167,25 +197,32 @@ class MatrixBot:
)
result = await self._client.whoami()
if type(result) is WhoamiError:
self._logger.error(f"Can't log in using stored session data: '{result.message}'")
raise RuntimeError(result.message)
elif type(result) is not WhoamiResponse:
self._logger.error("Can't log in using stored session data, unknown error")
raise RuntimeError("Unknown response for whoami request")
self._logger.debug("Logged in using stored session data")
async def _client_login_password(self) -> None:
"""Login using password and save result to disk on success.
Raises an exception on failure.
"""
self._logger.debug("Using password to log in")
# get the password
password = await self._cb_password()
result = await self._client.login(password=password)
if type(result) is LoginResponse:
self._logger.debug("Logged in using password")
await self._write_session_data(
access_token=result.access_token,
device_id=result.device_id
)
elif type(result) is LoginError:
self._logger.error(f"Can't log in using password: '{result.message}'")
raise RuntimeError(result.message)
else:
self._logger.error(f"Can't log in using password, unknown error")
raise RuntimeError("Unknown login result")
async def _client_login(self) -> None:
@@ -194,31 +231,68 @@ class MatrixBot:
session_data = await self._read_session_data()
# session data is present, try to log in
if session_data is not None:
self._logger.debug("Some session data found on the disk")
try:
await self._client_login_session_data(session_data)
return
except asyncio.CancelledError:
raise
except:
pass
# no session data - login using password
try:
self._logger.debug("No session data found on the disk OR invalid data")
await self._client_login_password()
return
except asyncio.CancelledError:
raise
except:
pass
# can't login
self._logger.error("Can't log in using available methods")
raise RuntimeError("All login methods have failed, can't continue")
async def _client_destroy(self) -> None:
"""Gracefully destroys the client."""
try:
self._logger.debug("Closing the client")
await self._client.close()
if self._last_next_batch is not None:
self._logger.debug("Saving next_batch")
await self._write_next_batch(self._last_next_batch)
except:
traceback.print_exc()
async def _background_coroutine(self) -> None:
"""This function implements bot lifecycle."""
# we should stop when this task stops
self._logger.debug("_background_coroutine is started")
stop_wait_task = asyncio.create_task(self._stop_event.wait())
self._client = self._build_client(self._config)
self._setup_client_callbacks()
# perform login
login_task = asyncio.create_task(self._client_login())
try:
await self._client_login()
await self._wait_for_task_and_stop_event(login_task, stop_wait_task)
except asyncio.CancelledError:
self._logger.debug("Background task is cancelled during login")
await self._client_destroy()
return
except:
traceback.print_exc()
self._logger.info("Succesfully logged in")
# test
wait_task = asyncio.create_task(asyncio.sleep(100000))
try:
self._logger.info("Bot is not implemented yet, sleeping forever")
await self._wait_for_task_and_stop_event(wait_task, stop_wait_task)
except asyncio.CancelledError:
self._logger.info("Background task is cancelled during eternal sleep")
await self._client_destroy()
return
except:
traceback.print_exc()
#
# PUBLIC
@@ -228,11 +302,15 @@ class MatrixBot:
self._validate_bot_config(config) # may raise an Exception
# save the config
self._config = config
# create the logger
self._logger = logging.getLogger(self._build_matrix_username(config))
self._logger.setLevel(logging.DEBUG)
# prepare some private data
self._background_task: asyncio.Task | None = None
self._client: AsyncClient = None # type: ignore
self._last_next_batch_dump: float = 0.0
self._last_next_batch: str | None = None
self._cb_password = self._default_password_callback
def start(self) -> None:
@@ -249,22 +327,14 @@ class MatrixBot:
self._background_coroutine()
)
def stop(self) -> bool:
"""Stop the bot.
Signals the background task to stop. Returns True on success (the
signal is sent), False on failure (no background task). This
function won't wait for the task to stop. Use `wait_stop` to wait.
"""
if self._background_task is None:
return False
self._stop_event.set()
return True
async def wait_stop(self) -> None:
"""Wait for the background task to stop.
You must call `stop()` by yourself. This function will never return
otherwise. It just waits for the background task stop.
"""
async def stop(self) -> None:
"""Stop the bot and wait for the bot stop."""
if self._background_task is None:
return
await asyncio.wait([self._background_task])
self._stop_event.set()
try:
await self._background_task
except:
traceback.print_exc()
self._stop_event = None
self._background_task = None