Separated CLI orchestration from the Crabstero bot object.
This commit is contained in:
@@ -0,0 +1,402 @@
|
||||
# Copyright 2026 Logan Fick
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Command-line interface for the Crabstero Discord bot.
|
||||
|
||||
Provides an argparse CLI with environment variable and systemd credential
|
||||
fallbacks, plus systemd notify/watchdog integration when run as a service.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import uvloop
|
||||
|
||||
from crabstero.bot import Crabstero
|
||||
from crabstero.metrics import MetricsAddress, TcpMetricsAddress, UnixMetricsAddress
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterator
|
||||
|
||||
logger = logging.getLogger("crabstero")
|
||||
|
||||
_KEYBOARD_INTERRUPT_EXIT_CODE = 130
|
||||
_MAX_METRICS_UNIX_SOCKET_MODE = 0o777
|
||||
_HANDLED_SIGNALS = (signal.SIGINT, signal.SIGTERM)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CliConfig:
|
||||
"""Resolved command-line configuration."""
|
||||
|
||||
token: str
|
||||
database_path: str
|
||||
ingest_only: bool
|
||||
metrics_address: MetricsAddress | None
|
||||
|
||||
|
||||
def _sd_notify(state: str) -> None:
|
||||
"""Send a notification to systemd, if running under a systemd service.
|
||||
|
||||
:param state: The systemd notification payload.
|
||||
"""
|
||||
addr = os.environ.get("NOTIFY_SOCKET")
|
||||
if not addr:
|
||||
return
|
||||
|
||||
try:
|
||||
with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
|
||||
if addr[0] == "@":
|
||||
addr = "\0" + addr[1:]
|
||||
sock.sendto(state.encode(), addr)
|
||||
except OSError as exc:
|
||||
logger.debug("Could not send systemd notification: %s", exc)
|
||||
|
||||
|
||||
def _watchdog_interval() -> float | None:
|
||||
"""Return the systemd watchdog ping interval in seconds, if enabled."""
|
||||
usec = os.environ.get("WATCHDOG_USEC")
|
||||
if not usec:
|
||||
return None
|
||||
|
||||
pid = os.environ.get("WATCHDOG_PID")
|
||||
if pid is not None:
|
||||
try:
|
||||
watchdog_pid = int(pid)
|
||||
except ValueError:
|
||||
return None
|
||||
if watchdog_pid != os.getpid():
|
||||
return None
|
||||
|
||||
try:
|
||||
watchdog_usec = int(usec)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
if watchdog_usec <= 0:
|
||||
return None
|
||||
|
||||
return watchdog_usec / 1_000_000 / 2
|
||||
|
||||
|
||||
async def _watchdog_loop(interval: float, stop: asyncio.Event) -> None:
|
||||
"""Send systemd watchdog pings until *stop* is set.
|
||||
|
||||
:param interval: Seconds to wait between watchdog pings.
|
||||
:param stop: Event that stops the watchdog loop.
|
||||
"""
|
||||
while not stop.is_set():
|
||||
_sd_notify("WATCHDOG=1")
|
||||
with contextlib.suppress(TimeoutError):
|
||||
await asyncio.wait_for(stop.wait(), timeout=interval)
|
||||
|
||||
|
||||
def _read_credential(name: str) -> str | None:
|
||||
"""Read a value from a systemd credential file.
|
||||
|
||||
Looks for a file named *name* inside the directory pointed to by the
|
||||
``CREDENTIALS_DIRECTORY`` environment variable (set automatically by
|
||||
systemd when ``LoadCredential=`` or ``SetCredential=`` is used).
|
||||
|
||||
:param name: Credential name to look up.
|
||||
:return: The credential value, or ``None`` if unavailable.
|
||||
"""
|
||||
credentials_dir = os.environ.get("CREDENTIALS_DIRECTORY")
|
||||
if credentials_dir is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return Path(credentials_dir, name).read_text().strip()
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _install_signal_handlers(
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
callback: Callable[[signal.Signals], None],
|
||||
) -> Iterator[None]:
|
||||
"""Install process signal handlers and always remove installed handlers."""
|
||||
registered_signals: list[signal.Signals] = []
|
||||
try:
|
||||
for sig in _HANDLED_SIGNALS:
|
||||
loop.add_signal_handler(sig, callback, sig)
|
||||
registered_signals.append(sig)
|
||||
yield
|
||||
finally:
|
||||
for sig in registered_signals:
|
||||
loop.remove_signal_handler(sig)
|
||||
|
||||
|
||||
def _parse_args(argv: list[str] | None = None) -> _CliConfig:
|
||||
"""Parse command-line arguments with environment variable fallbacks.
|
||||
|
||||
:param argv: Optional argument list (defaults to sys.argv[1:]).
|
||||
:return: Resolved command-line configuration.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="crabstero",
|
||||
description="Crabstero - the simple nonversation Discord bot.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--token",
|
||||
default=os.environ.get("TOKEN") or _read_credential("token"),
|
||||
help=(
|
||||
"Discord bot token"
|
||||
" (default: TOKEN environment variable"
|
||||
" or systemd credential 'token')."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database-path",
|
||||
"--database",
|
||||
default=os.environ.get("DATABASE_PATH", "crabstero.db"),
|
||||
help=(
|
||||
"Path to the SQLite database file"
|
||||
" (default: DATABASE_PATH environment"
|
||||
' variable or "crabstero.db").'
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ingest-only",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Run in ingest-only mode: ingest channel"
|
||||
" history and real-time messages but"
|
||||
" never respond."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--listen-metrics",
|
||||
default=os.environ.get("LISTEN_METRICS"),
|
||||
help=(
|
||||
"Enable Prometheus metrics endpoint on HOST:PORT or unix:/path.sock"
|
||||
" (e.g. 127.0.0.1:9090 or unix:/run/crabstero.sock)."
|
||||
" Disabled by default."
|
||||
" (default: LISTEN_METRICS environment variable)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--metrics-unix-socket-mode",
|
||||
default=os.environ.get("METRICS_UNIX_SOCKET_MODE"),
|
||||
help=(
|
||||
"Octal file mode to set on a metrics Unix socket"
|
||||
" (e.g. 0666). Only valid with --listen-metrics unix:/path.sock."
|
||||
" (default: METRICS_UNIX_SOCKET_MODE environment variable)."
|
||||
),
|
||||
)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
token = cast("str | None", args.token)
|
||||
database_path = cast("str", args.database_path)
|
||||
ingest_only = cast("bool", args.ingest_only)
|
||||
raw_metrics_address = cast("str | None", args.listen_metrics)
|
||||
raw_metrics_unix_socket_mode = cast(
|
||||
"str | None",
|
||||
args.metrics_unix_socket_mode,
|
||||
)
|
||||
|
||||
metrics_unix_socket_mode: int | None = None
|
||||
if raw_metrics_unix_socket_mode is not None:
|
||||
try:
|
||||
metrics_unix_socket_mode = int(raw_metrics_unix_socket_mode, 8)
|
||||
except ValueError:
|
||||
parser.error(
|
||||
"--metrics-unix-socket-mode must be an octal mode"
|
||||
" between 0000 and 0777",
|
||||
)
|
||||
if not 0 <= metrics_unix_socket_mode <= _MAX_METRICS_UNIX_SOCKET_MODE:
|
||||
parser.error(
|
||||
"--metrics-unix-socket-mode must be an octal mode"
|
||||
" between 0000 and 0777",
|
||||
)
|
||||
|
||||
metrics_address: MetricsAddress | None = None
|
||||
if raw_metrics_address is not None:
|
||||
if raw_metrics_address.startswith("unix:"):
|
||||
path = raw_metrics_address.removeprefix("unix:")
|
||||
if not path:
|
||||
parser.error("--listen-metrics Unix socket path must not be empty")
|
||||
metrics_address = UnixMetricsAddress(path, metrics_unix_socket_mode)
|
||||
else:
|
||||
if metrics_unix_socket_mode is not None:
|
||||
parser.error(
|
||||
"--metrics-unix-socket-mode is only valid with"
|
||||
" --listen-metrics unix:/path.sock",
|
||||
)
|
||||
host, sep, port_str = raw_metrics_address.rpartition(":")
|
||||
if not sep or not host:
|
||||
parser.error(
|
||||
"--listen-metrics must be in HOST:PORT or unix:/path.sock format "
|
||||
"(e.g. 127.0.0.1:9090)",
|
||||
)
|
||||
try:
|
||||
metrics_address = TcpMetricsAddress(host, int(port_str))
|
||||
except ValueError:
|
||||
parser.error(
|
||||
f"--listen-metrics port must be an integer, got '{port_str}'",
|
||||
)
|
||||
elif metrics_unix_socket_mode is not None:
|
||||
parser.error(
|
||||
"--metrics-unix-socket-mode is only valid with"
|
||||
" --listen-metrics unix:/path.sock",
|
||||
)
|
||||
|
||||
if token is None:
|
||||
parser.error(
|
||||
"a Discord bot token is required via --token,"
|
||||
" the TOKEN environment variable,"
|
||||
" or a systemd credential named 'token'",
|
||||
)
|
||||
|
||||
return _CliConfig(
|
||||
token=token,
|
||||
database_path=database_path,
|
||||
ingest_only=ingest_only,
|
||||
metrics_address=metrics_address,
|
||||
)
|
||||
|
||||
|
||||
async def _run(config: _CliConfig) -> signal.Signals | None:
|
||||
"""Run the bot under CLI-controlled process orchestration."""
|
||||
bot = Crabstero(
|
||||
database_path=config.database_path,
|
||||
ingest_only=config.ingest_only,
|
||||
metrics_address=config.metrics_address,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
watchdog_stop = asyncio.Event()
|
||||
watchdog_task: asyncio.Task[None] | None = None
|
||||
shutdown_task: asyncio.Task[None] | None = None
|
||||
ready_notified = False
|
||||
stopping_notified = False
|
||||
cleanup_started = False
|
||||
shutdown_signal: signal.Signals | None = None
|
||||
runner_task = asyncio.current_task()
|
||||
if runner_task is None:
|
||||
msg = "Crabstero runner is not running in a task"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
def notify_stopping() -> None:
|
||||
nonlocal stopping_notified
|
||||
if stopping_notified:
|
||||
return
|
||||
stopping_notified = True
|
||||
_sd_notify("STOPPING=1")
|
||||
|
||||
async def shutdown(sig: signal.Signals) -> None:
|
||||
logger.info("Received shutdown signal %s.", sig.name)
|
||||
notify_stopping()
|
||||
await bot.close()
|
||||
|
||||
def request_shutdown(sig: signal.Signals) -> None:
|
||||
nonlocal shutdown_signal, shutdown_task
|
||||
if shutdown_signal is not None:
|
||||
return
|
||||
|
||||
shutdown_signal = sig
|
||||
if not cleanup_started:
|
||||
shutdown_task = asyncio.create_task(
|
||||
shutdown(sig),
|
||||
name="crabstero-shutdown",
|
||||
)
|
||||
runner_task.cancel()
|
||||
|
||||
async def cleanup() -> None:
|
||||
nonlocal cleanup_started
|
||||
if cleanup_started:
|
||||
return
|
||||
|
||||
cleanup_started = True
|
||||
|
||||
if ready_notified:
|
||||
notify_stopping()
|
||||
|
||||
watchdog_stop.set()
|
||||
if watchdog_task is not None:
|
||||
watchdog_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await watchdog_task
|
||||
|
||||
if shutdown_task is not None:
|
||||
await shutdown_task
|
||||
|
||||
with _install_signal_handlers(loop, request_shutdown):
|
||||
try:
|
||||
async with bot:
|
||||
try:
|
||||
await bot.login(config.token)
|
||||
if bot.is_closed() or stopping_notified:
|
||||
return shutdown_signal
|
||||
|
||||
_sd_notify("READY=1")
|
||||
ready_notified = True
|
||||
|
||||
wd_interval = _watchdog_interval()
|
||||
if wd_interval is not None:
|
||||
watchdog_interval = wd_interval
|
||||
logger.info(
|
||||
"systemd watchdog enabled (pinging every %.1fs).",
|
||||
watchdog_interval,
|
||||
)
|
||||
watchdog_task = asyncio.create_task(
|
||||
_watchdog_loop(watchdog_interval, watchdog_stop),
|
||||
name="systemd-watchdog",
|
||||
)
|
||||
|
||||
await bot.connect(reconnect=True)
|
||||
finally:
|
||||
await cleanup()
|
||||
except asyncio.CancelledError:
|
||||
if shutdown_signal is None:
|
||||
raise
|
||||
finally:
|
||||
await cleanup()
|
||||
|
||||
return shutdown_signal
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
"""Run the bot CLI.
|
||||
|
||||
:param argv: Optional argument list (defaults to sys.argv[1:]).
|
||||
:return: Process exit code.
|
||||
"""
|
||||
config = _parse_args(argv)
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s",
|
||||
)
|
||||
|
||||
try:
|
||||
shutdown_signal = uvloop.run(_run(config))
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Interrupted.")
|
||||
return _KEYBOARD_INTERRUPT_EXIT_CODE
|
||||
|
||||
if shutdown_signal == signal.SIGINT:
|
||||
return _KEYBOARD_INTERRUPT_EXIT_CODE
|
||||
return 0
|
||||
Reference in New Issue
Block a user