From e83d48f1f70e07c36c0b988e025191365fdc7645 Mon Sep 17 00:00:00 2001
From: YoursFunny <admin@yoursfunny.top>
Date: Tue, 4 Aug 2026 22:09:45 +0800
Subject: [PATCH] fix: handle SIGTERM for graceful shutdown on docker stop

---
 AGENTS.md                     |  2 +-
 crates/xmedia-bot/src/main.rs | 47 ++++++++++++++++++++++++++++++-----
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index cdb7844..37ace4f 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -65,7 +65,7 @@ Docker: `docker build -t tgxmb .` then `docker run --rm -d --name tgxmb --env-fi
 
 | File | Why it matters |
 |---|---|
-| `crates/xmedia-bot/src/main.rs` | Startup sequence, webhook vs polling, graceful shutdown (ctrlc → sweep stop → admin msg → queue stop) |
+| `crates/xmedia-bot/src/main.rs` | Startup sequence, webhook vs polling, graceful shutdown (SIGINT via teloxide ctrlc / SIGTERM via `stop_token` for docker, → sweep stop → admin msg → queue stop) |
 | `crates/xmedia-bot/src/handlers.rs` | `CHAT_STORE`/`TASK_QUEUE`/`CONFIG` singletons (open `data/task_queue.db` **relative to CWD**); command dispatch; URL extraction; retry enqueue |
 | `crates/xmedia-bot/src/send.rs` | Constants `MAX_MEDIA_GROUP = 9`, `MAX_UPLOAD_BYTES = 10 MiB`; fallback chain; `classify_request_error` |
 | `crates/x-media/src/site/mod.rs` | Dispatcher, `Fetched`/`FetchError`, shared `CLIENT`, `download_media` (adds `Referer: https://www.pixiv.net/` for `pximg.net` hotlink protection) |
diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs
index aedfa0b..984e7f2 100644
--- a/crates/xmedia-bot/src/main.rs
+++ b/crates/xmedia-bot/src/main.rs
@@ -1,7 +1,8 @@
 use dotenv::dotenv;
 use teloxide::dptree::endpoint;
+use teloxide::stop::StopToken;
 use teloxide::types::{ChatId, InputFile, MessageId};
-use teloxide::update_listeners::webhooks;
+use teloxide::update_listeners::{self, webhooks, UpdateListener};
 use teloxide::prelude::*;
 use tokio::sync::watch;
 use x_media::site;
@@ -14,6 +15,24 @@ mod state;
 
 use handlers::{CHAT_STORE, CONFIG, TASK_QUEUE};
 
+/// Docker `stop` / `compose down` delivers SIGTERM, which teloxide's ctrlc
+/// handler (SIGINT only) never sees — without this the process would die
+/// before the graceful shutdown below (admin notice, queue drain). Stopping
+/// the token unwinds the dispatcher exactly like Ctrl+C does.
+#[cfg(unix)]
+fn spawn_sigterm_handler(stop_token: StopToken) {
+    tokio::spawn(async move {
+        let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+            .expect("failed to install SIGTERM handler");
+        sigterm.recv().await;
+        log::info!("SIGTERM received, stopping the dispatcher");
+        stop_token.stop();
+    });
+}
+
+#[cfg(not(unix))]
+fn spawn_sigterm_handler(_stop_token: StopToken) {}
+
 #[tokio::main]
 async fn main() {
     dotenv().ok();
@@ -108,20 +127,36 @@ async fn main() {
             options = options.secret_token(secret.clone());
         }
 
+        let mut listener = webhooks::axum(bot.clone(), options)
+            .await
+            .expect("Failed to create webhook listener");
+        let stop_token = listener.stop_token();
+        spawn_sigterm_handler(stop_token);
+
         dispatcher
             .dispatch_with_listener(
-                webhooks::axum(bot.clone(), options)
-                    .await
-                    .expect("Failed to create webhook listener"),
+                listener,
                 LoggingErrorHandler::with_custom_text("Error from update listener"),
             )
             .await;
     } else {
         log::info!("running in polling mode");
-        dispatcher.dispatch().await;
+        // Same listener `dispatch()` builds internally — using
+        // `dispatch_with_listener` just exposes its stop token so SIGTERM can
+        // unwind the dispatcher before the graceful shutdown below.
+        let mut listener = update_listeners::polling_default(bot.clone()).await;
+        let stop_token = listener.stop_token();
+        spawn_sigterm_handler(stop_token);
+
+        dispatcher
+            .dispatch_with_listener(
+                listener,
+                LoggingErrorHandler::with_custom_text("Error from update listener"),
+            )
+            .await;
     }
 
-    // Graceful stop (Ctrl+C): stop the sweep, notify the admin, drain the queue.
+    // Graceful stop (Ctrl+C / SIGTERM): stop the sweep, notify the admin, drain the queue.
     log::info!("Stopping bot");
     let _ = stop_tx.send(true);
     if let Some(admin) = CONFIG.admin_ids.first() {