From c2d7c8406e2606b40c029319377e0a5fb83f08ee Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Thu, 17 Sep 2026 01:27:05 +0800 Subject: [PATCH] refactor: finish the phase-B seam for the post-send path, funnel settlement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit docs/architecture-refactor.md §3 stopped half-done: `url_media` got an injected `AppContext`, but `send.rs`'s post-send half kept reaching for the process-wide `CHAT_STORE`/`TASK_QUEUE`/`LINK_CACHE` statics, so the whole shell after a successful send (edit-before-forward prompt, channel forward, retry enqueue, cache write) had no test and no way to get one. - `ctx.rs` now owns `AppContext` (sender + the three stores + config) with `from_statics` for production and a `CONTEXT` static for the spawned worker closures; `handlers/urls.rs` drops its private copy and the duplicated assembler, and the queue handler/dead-letter callbacks take the context (main wires them with `CONTEXT`). - `send_media_sequence`/`send_animation`/`forward_messages`/`post_send_actions` take `&AppContext`; the cache write goes through the injected cache. - New `settle_task(ctx, task, Sent|Failed)` is the single place that ends a task: release its keep-alive temp media, and drop the link-cache entry only on failure. All five former call sites funnel through it — the earlier keep-alive leak existed precisely because one of them had to remember. `invalidate_cache`/`invalidate_cache_with` (static + injected pair, the latter only existing because of the former) collapse into one private fn. - `ctx::test_support::TestStores` gives tests a tempdir store set + context; `handlers/urls.rs` tests use it instead of hand-rolled setup. Tests: +5 (post-send forward ok / queued / notified, settle Sent/Failed); the post-send and settle paths were previously untested. fmt/clippy clean, 60 + 69 tests pass. --- crates/xmedia-bot/src/ctx.rs | 114 +++++++ crates/xmedia-bot/src/handlers/callback.rs | 3 +- crates/xmedia-bot/src/handlers/urls.rs | 110 ++----- crates/xmedia-bot/src/main.rs | 10 +- crates/xmedia-bot/src/send.rs | 349 ++++++++++++++++----- 5 files changed, 430 insertions(+), 156 deletions(-) create mode 100644 crates/xmedia-bot/src/ctx.rs diff --git a/crates/xmedia-bot/src/ctx.rs b/crates/xmedia-bot/src/ctx.rs new file mode 100644 index 0000000..e4ef800 --- /dev/null +++ b/crates/xmedia-bot/src/ctx.rs @@ -0,0 +1,114 @@ +//! Runtime context: the collaborators a handler needs, injected as one struct +//! so tests can substitute a scripted sender and tempdir-backed stores. +//! +//! The production context is assembled from the process-wide statics +//! ([`AppContext::from_statics`]); the spawned worker closures hold +//! [`CONTEXT`], which is `'static` for that reason. + +use crate::config::Config; +use crate::handlers::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE}; +use crate::link_cache::LinkCache; +use crate::media_sender::MediaSender; +use crate::queue::PersistentTaskQueue; +use crate::send::BOT; +use crate::state::ChatStore; +use std::sync::LazyLock; + +pub struct AppContext<'a> { + pub sender: &'a dyn MediaSender, + pub chat_store: &'a ChatStore, + pub task_queue: &'a PersistentTaskQueue, + pub link_cache: &'a LinkCache, + pub config: &'a Config, +} + +impl<'a> AppContext<'a> { + /// The stores are the process-wide statics; `sender` is whatever the caller + /// was handed (the dispatcher's `Bot` clone for update handlers, the shared + /// queue `Bot` for the worker loops). Update handlers build their own + /// context from the `Bot` they received so the same code path works with an + /// injected mock in tests. + pub fn from_statics(sender: &'a dyn MediaSender) -> AppContext<'a> { + AppContext { + sender, + chat_store: &CHAT_STORE, + task_queue: &TASK_QUEUE, + link_cache: &LINK_CACHE, + config: &CONFIG, + } + } +} + +/// The URL/queue workers' context: `'static` because `tokio::spawn`ed closures +/// and the queue's handler type require it. +pub static CONTEXT: LazyLock> = + LazyLock::new(|| AppContext::from_statics(&*BOT)); + +/// Test support: a tempdir-backed set of stores plus the context borrowing +/// them, so a handler test needs one line of setup. +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + use std::sync::Arc; + + pub(crate) struct TestStores { + _dir: tempfile::TempDir, + pool: Arc, + chat_store: ChatStore, + task_queue: PersistentTaskQueue, + link_cache: LinkCache, + config: Config, + } + + impl TestStores { + pub(crate) fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let pool = crate::db::open_store(dir.path().join("ctx.db").to_str().unwrap()).unwrap(); + TestStores { + _dir: dir, + chat_store: ChatStore::new(Arc::clone(&pool)), + task_queue: PersistentTaskQueue::new(Arc::clone(&pool)), + link_cache: LinkCache::new(Arc::clone(&pool)), + config: Config::load(), + pool, + } + } + + pub(crate) fn ctx<'a>(&'a self, sender: &'a dyn MediaSender) -> AppContext<'a> { + AppContext { + sender, + chat_store: &self.chat_store, + task_queue: &self.task_queue, + link_cache: &self.link_cache, + config: &self.config, + } + } + + pub(crate) fn link_cache(&self) -> &LinkCache { + &self.link_cache + } + + /// Rows persisted in the task queue: what "queued for retry" looks like + /// from the outside. + pub(crate) async fn queued_tasks(&self) -> i64 { + let pool = Arc::clone(&self.pool); + pool.with_conn(|conn| { + conn.query_row("SELECT COUNT(*) FROM tasks", [], |row| row.get(0)) + }) + .await + .unwrap() + } + + /// The single queued task payload, for asserting what was rescheduled. + pub(crate) async fn queued_payload(&self) -> serde_json::Value { + let pool = Arc::clone(&self.pool); + let payload: String = pool + .with_conn(|conn| { + conn.query_row("SELECT payload FROM tasks LIMIT 1", [], |row| row.get(0)) + }) + .await + .unwrap(); + serde_json::from_str(&payload).unwrap() + } + } +} diff --git a/crates/xmedia-bot/src/handlers/callback.rs b/crates/xmedia-bot/src/handlers/callback.rs index b7f90b7..7e5743e 100644 --- a/crates/xmedia-bot/src/handlers/callback.rs +++ b/crates/xmedia-bot/src/handlers/callback.rs @@ -50,6 +50,7 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() chat_id ); if data == "forward" { + let ctx = crate::ctx::AppContext::from_statics(&bot); match chat_data.forward_channel_id { Some(channel_id) => { let forward_task = Task::ForwardMessages { @@ -59,7 +60,7 @@ pub async fn callback_query_handler(bot: Bot, query: CallbackQuery) -> Result<() notify_chat_id: Some(chat_id), notify_message_id: Some(prompt_message_id), }; - match send::forward_messages(&bot, &forward_task).await { + match send::forward_messages(&ctx, &forward_task).await { Ok(()) => { log::info!( "forwarded {} message(s) to channel {channel_id}", diff --git a/crates/xmedia-bot/src/handlers/urls.rs b/crates/xmedia-bot/src/handlers/urls.rs index 6d6f6d1..ea012a8 100644 --- a/crates/xmedia-bot/src/handlers/urls.rs +++ b/crates/xmedia-bot/src/handlers/urls.rs @@ -1,13 +1,11 @@ //! URL extraction and the per-URL media pipeline: bounded job channel + //! worker pool, link-cache fast path, fetch, task build and send dispatch. -use super::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE, log_key, reply}; -use crate::config::Config; -use crate::link_cache::{CachedMediaKind, CachedPost, LinkCache}; -use crate::media_sender::MediaSender; -use crate::queue::PersistentTaskQueue; +use super::{log_key, reply}; +use crate::ctx::{AppContext, CONTEXT}; +use crate::link_cache::{CachedMediaKind, CachedPost}; use crate::send::{self, MediaItemPayload, Task}; -use crate::state::{ChatData, ChatStore}; +use crate::state::ChatData; use std::collections::HashSet; use std::sync::LazyLock; use teloxide::types::{ChatAction, ChatId, Message, MessageEntityKind, MessageId}; @@ -34,27 +32,6 @@ static URL_WORKER_HANDLES: LazyLock { - pub sender: &'a dyn MediaSender, - pub chat_store: &'a ChatStore, - pub task_queue: &'a PersistentTaskQueue, - pub link_cache: &'a LinkCache, - pub config: &'a Config, -} - -/// Assembles the production context from the process-wide statics. -fn app_context() -> AppContext<'static> { - AppContext { - sender: &*crate::send::BOT, - chat_store: &CHAT_STORE, - task_queue: &TASK_QUEUE, - link_cache: &LINK_CACHE, - config: &CONFIG, - } -} - /// Starts the URL job workers (called once from main after the queue starts). /// teloxide dispatches updates to a per-chat worker that handles them /// sequentially, so a batch-forward of many messages would otherwise be @@ -69,12 +46,11 @@ pub async fn start_url_workers() { for _ in 0..URL_WORKERS { let rx = std::sync::Arc::clone(&rx); handles.push(tokio::spawn(async move { - let ctx = app_context(); while !URL_STOP.load(std::sync::atomic::Ordering::Relaxed) { let job = rx.lock().await.recv().await; match job { Some((message, url)) => { - url_media(&ctx, message.chat.id.0, message.id.0 as i64, &url).await + url_media(&CONTEXT, message.chat.id.0, message.id.0 as i64, &url).await } None => break, } @@ -184,8 +160,8 @@ async fn dispatch_send( url: &str, ) { let result = match task { - Task::SendAnimation { .. } => send::send_animation(ctx.sender, task).await, - Task::SendMediaSequence { .. } => send::send_media_sequence(ctx.sender, task).await, + Task::SendAnimation { .. } => send::send_animation(ctx, task).await, + Task::SendMediaSequence { .. } => send::send_media_sequence(ctx, task).await, Task::ForwardMessages { .. } => unreachable!(), }; match result { @@ -195,9 +171,8 @@ async fn dispatch_send( message_ids.len(), log_key(url) ); - send::post_send_actions(ctx.sender, task, message_ids).await; - // The task settled: drop any keep-alive temp media. - send::release_keep_alive(task); + send::post_send_actions(ctx, task, message_ids).await; + send::settle_task(ctx, task, send::Settled::Sent).await; } Err(send::SendError::Retryable { delay_seconds, @@ -220,8 +195,7 @@ async fn dispatch_send( message: err_message, task, }) => { - send::invalidate_cache_with(ctx.link_cache, &task).await; - send::release_keep_alive(&task); + send::settle_task(ctx, &task, send::Settled::Failed).await; log::error!("send for {url} failed permanently: {err_message}"); let _ = reply( ctx.sender, @@ -434,10 +408,9 @@ async fn url_media(ctx: &AppContext<'_>, chat_id: i64, reply_to_message_id: i64, #[cfg(test)] mod tests { use super::*; - use crate::db; + use crate::ctx::test_support::TestStores; use crate::link_cache::CachedMedia; use crate::media_sender::test_support::{MockSender, Outcome}; - use std::sync::Arc; use std::time::Duration; use teloxide::{ApiError, RequestError}; @@ -465,24 +438,16 @@ mod tests { #[tokio::test] async fn cache_hit_sends_file_ids_and_invalidates_on_permanent_failure() { - let dir = tempfile::tempdir().unwrap(); - let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap(); - let chat_store = ChatStore::new(Arc::clone(&pool)); - let task_queue = PersistentTaskQueue::new(Arc::clone(&pool)); - let link_cache = LinkCache::new(Arc::clone(&pool)); - let config = Config::load(); + let stores = TestStores::new(); let sender = MockSender::scripted( vec![Outcome::GroupErr, Outcome::MessageErr], permanent_error, ); - let ctx = AppContext { - sender: &sender, - chat_store: &chat_store, - task_queue: &task_queue, - link_cache: &link_cache, - config: &config, - }; - link_cache.put("twitter:1", &cached_photo_entry()).await; + let ctx = stores.ctx(&sender); + stores + .link_cache() + .put("twitter:1", &cached_photo_entry()) + .await; url_media(&ctx, 1, 2, "https://x.com/u/status/1").await; @@ -494,7 +459,8 @@ mod tests { ); // The stale cache entry was invalidated so the next request re-fetches. assert!( - link_cache + stores + .link_cache() .get("twitter:1", Duration::from_secs(3600)) .await .is_none() @@ -503,28 +469,21 @@ mod tests { #[tokio::test] async fn cache_hit_success_keeps_the_cache_entry() { - let dir = tempfile::tempdir().unwrap(); - let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap(); - let chat_store = ChatStore::new(Arc::clone(&pool)); - let task_queue = PersistentTaskQueue::new(Arc::clone(&pool)); - let link_cache = LinkCache::new(Arc::clone(&pool)); - let config = Config::load(); + let stores = TestStores::new(); let sender = MockSender::scripted(vec![Outcome::GroupOk], permanent_error); - let ctx = AppContext { - sender: &sender, - chat_store: &chat_store, - task_queue: &task_queue, - link_cache: &link_cache, - config: &config, - }; - link_cache.put("twitter:1", &cached_photo_entry()).await; + let ctx = stores.ctx(&sender); + stores + .link_cache() + .put("twitter:1", &cached_photo_entry()) + .await; url_media(&ctx, 1, 2, "https://x.com/u/status/1").await; assert_eq!(sender.calls(), vec!["send_chat_action", "send_media_group"]); // Success must not evict the entry. assert!( - link_cache + stores + .link_cache() .get("twitter:1", Duration::from_secs(3600)) .await .is_some() @@ -533,20 +492,9 @@ mod tests { #[tokio::test] async fn unsupported_url_is_ignored_silently() { - let dir = tempfile::tempdir().unwrap(); - let pool = db::open_store(dir.path().join("t.db").to_str().unwrap()).unwrap(); - let chat_store = ChatStore::new(Arc::clone(&pool)); - let task_queue = PersistentTaskQueue::new(Arc::clone(&pool)); - let link_cache = LinkCache::new(Arc::clone(&pool)); - let config = Config::load(); + let stores = TestStores::new(); let sender = MockSender::scripted(vec![], permanent_error); - let ctx = AppContext { - sender: &sender, - chat_store: &chat_store, - task_queue: &task_queue, - link_cache: &link_cache, - config: &config, - }; + let ctx = stores.ctx(&sender); // No cache key → the fetch dispatcher returns Ok(None) without any // network; nothing is sent or replied. diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index d4c4556..cfd4651 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -8,6 +8,7 @@ use tokio::sync::watch; use x_media::site; mod config; +mod ctx; mod db; mod handlers; mod link_cache; @@ -18,6 +19,7 @@ mod rate_limit; mod send; mod state; +use ctx::CONTEXT; use handlers::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE}; /// Docker `stop` / `compose down` delivers SIGTERM, which teloxide's ctrlc @@ -61,9 +63,13 @@ async fn main() { ); // Queue worker: handles typed tasks, dead-letters failed sends to the - // task's chat. + // task's chat. Both closures use the shared context (the queue requires + // 'static handlers, and the statics are process-wide anyway). TASK_QUEUE - .start(send::handle_task, send::dead_letter_notify) + .start( + |payload| send::handle_task(&CONTEXT, payload), + |payload, message| send::dead_letter_notify(&CONTEXT, payload, message), + ) .await; log::info!("task queue worker started"); diff --git a/crates/xmedia-bot/src/send.rs b/crates/xmedia-bot/src/send.rs index d798bb0..75c8822 100644 --- a/crates/xmedia-bot/src/send.rs +++ b/crates/xmedia-bot/src/send.rs @@ -3,8 +3,9 @@ //! URL is blocked by hotlink protection; the bot downloads the file itself //! and uploads it via multipart). +use crate::ctx::AppContext; use crate::db::{now_f64, unix_now}; -use crate::handlers::{CHAT_STORE, LINK_CACHE, TASK_QUEUE, log_key}; +use crate::handlers::log_key; use crate::link_cache::{CachedMedia, CachedMediaKind, CachedPost, LinkCache}; use crate::media_sender::MediaSender; use crate::photo::{self, MAX_UPLOAD_BYTES, PhotoPrep}; @@ -223,7 +224,7 @@ fn collect_file_ids(messages: &[Message], batch: &[MediaItemPayload], out: &mut /// Persists a successful send under the post's cache key. Only runs for a /// fresh (non-resumed) task that carried raw cache data with no file ids yet. -async fn cache_sent_task(task: &Task, media: Vec) { +async fn cache_sent_task(ctx: &AppContext<'_>, task: &Task, media: Vec) { let Some(cache_data) = task.cache_data() else { return; }; @@ -233,15 +234,16 @@ async fn cache_sent_task(task: &Task, media: Vec) { let mut post = cache_data.clone(); post.media = media; if let Some(key) = x_media::site::cache_key(&post.url) { - LINK_CACHE.put(&key, &post).await; + ctx.link_cache.put(&key, &post).await; log::debug!("cached send for [key={}]", log_key(&post.url)); } } /// Persists a lone animation send under the post's cache key. -async fn cache_animation_send(task: &Task, message: &Message) { +async fn cache_animation_send(ctx: &AppContext<'_>, task: &Task, message: &Message) { if let Some(file_id) = message.animation().map(|a| a.file.id.to_string()) { cache_sent_task( + ctx, task, vec![CachedMedia { kind: CachedMediaKind::Animation, @@ -252,14 +254,28 @@ async fn cache_animation_send(task: &Task, message: &Message) { } } -/// A cached Telegram file id failed permanently (stale/expired); drop the -/// cache entry so the next request re-fetches instead of repeating it. -pub async fn invalidate_cache(task: &Task) { - invalidate_cache_with(&LINK_CACHE, task).await; +/// How a task ended. The two states differ only in whether a link-cache entry +/// may still be holding the (now unusable) media. +pub enum Settled { + Sent, + Failed, } -/// [`invalidate_cache`] against an injected cache (tests pass a tempdir one). -pub async fn invalidate_cache_with(cache: &LinkCache, task: &Task) { +/// Every path that ends a task's life — sent, permanently failed, or +/// dead-lettered after the last retry — funnels through here, so the cleanup a +/// settled task owes cannot be forgotten by a new path: release the keep-alive +/// temp media (retryable tasks keep it, they will be resent) and drop the +/// link-cache entry that a failed send's stale file ids would keep poisoning. +pub async fn settle_task(ctx: &AppContext<'_>, task: &Task, outcome: Settled) { + if matches!(outcome, Settled::Failed) { + invalidate_cache(ctx.link_cache, task).await; + } + release_keep_alive(task); +} + +/// A cached Telegram file id failed permanently (stale/expired); drop the +/// cache entry so the next request re-fetches instead of repeating it. +async fn invalidate_cache(cache: &LinkCache, task: &Task) { if task.is_cached_send() && let Some(url) = task.source_url() && let Some(key) = x_media::site::cache_key(url) @@ -909,10 +925,7 @@ fn updated_sequence_task(task: &Task, batch_index: usize, sent_message_ids: Vec< /// Sends the media batches starting at `task.batch_index`, extending /// `sent_message_ids`. Returns all sent message ids on full success; on /// failure returns a [`SendError`] whose task carries the resumed state. -pub async fn send_media_sequence( - sender: &dyn MediaSender, - task: &Task, -) -> Result, SendError> { +pub async fn send_media_sequence(ctx: &AppContext<'_>, task: &Task) -> Result, SendError> { let Task::SendMediaSequence { chat_id, reply_to_message_id, @@ -948,7 +961,8 @@ pub async fn send_media_sequence( }); } }; - match sender + match ctx + .sender .send_media_group(ChatId(chat_id), MessageId(reply_to as i32), items) .await { @@ -971,7 +985,7 @@ pub async fn send_media_sequence( .unwrap_or_else(|| "?".into()) ); match send_batch_via_upload( - sender, + ctx.sender, chat_id, reply_to, batch, @@ -997,7 +1011,7 @@ pub async fn send_media_sequence( } } if fresh_send { - cache_sent_task(task, cached_media).await; + cache_sent_task(ctx, task, cached_media).await; } Ok(sent) } @@ -1022,7 +1036,7 @@ async fn send_animation_inner( } /// Sends a lone animation (gif), URL first with the download fallback. -pub async fn send_animation(sender: &dyn MediaSender, task: &Task) -> Result, SendError> { +pub async fn send_animation(ctx: &AppContext<'_>, task: &Task) -> Result, SendError> { let Task::SendAnimation { chat_id, reply_to_message_id, @@ -1052,10 +1066,19 @@ pub async fn send_animation(sender: &dyn MediaSender, task: &Task) -> Result { let id = message.id.0 as i64; - cache_animation_send(task, &message).await; + cache_animation_send(ctx, task, &message).await; Ok(vec![id]) } Err(RequestError::Api(api)) if is_media_fetch_failure(&api) || is_size_error(&api) => { @@ -1079,7 +1102,7 @@ pub async fn send_animation(sender: &dyn MediaSender, task: &Task) -> Result Result { let id = message.id.0 as i64; - cache_animation_send(task, &message).await; + cache_animation_send(ctx, task, &message).await; Ok(vec![id]) } Err(e) => Err(classify_to_send_error( @@ -1113,7 +1136,7 @@ pub async fn send_animation(sender: &dyn MediaSender, task: &Task) -> Result Result<(), SendError> { +pub async fn forward_messages(ctx: &AppContext<'_>, task: &Task) -> Result<(), SendError> { let Task::ForwardMessages { from_chat_id, to_chat_id, @@ -1127,7 +1150,8 @@ pub async fn forward_messages(sender: &dyn MediaSender, task: &Task) -> Result<( .iter() .map(|id| MessageId(*id as i32)) .collect::>(); - match sender + match ctx + .sender .copy_messages( ChatId(*to_chat_id), ChatId(*from_chat_id), @@ -1192,7 +1216,7 @@ pub async fn notify_failure( /// After a successful send: either open the edit-before-forward prompt or /// forward to the configured channel (with retry/queue handling). -pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_ids: Vec) { +pub async fn post_send_actions(ctx: &AppContext<'_>, task: &Task, message_ids: Vec) { let ( chat_id, reply_to, @@ -1234,8 +1258,9 @@ pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_id }; if edit_before_forward { - let keyboard = build_edit_markup(&CHAT_STORE.get(chat_id).await.template); - let prompt = sender + let keyboard = build_edit_markup(&ctx.chat_store.get(chat_id).await.template); + let prompt = ctx + .sender .send_message( ChatId(chat_id), "Reply to edit message.".to_string(), @@ -1252,7 +1277,7 @@ pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_id ); let prompt_id = prompt.id.0 as i64; let source_url = source_url.clone(); - CHAT_STORE + ctx.chat_store .update(chat_id, move |data| { data.edit_message.insert( prompt_id, @@ -1284,17 +1309,17 @@ pub async fn post_send_actions(sender: &dyn MediaSender, task: &Task, message_id notify_chat_id, notify_message_id, }; - match forward_messages(sender, &forward_task).await { + match forward_messages(ctx, &forward_task).await { Ok(()) => {} Err(SendError::Retryable { delay_seconds, task, }) => { - enqueue_retry(&TASK_QUEUE, *task, delay_seconds).await; + enqueue_retry(ctx.task_queue, *task, delay_seconds).await; } Err(SendError::Permanent { message, .. }) => { notify_failure( - sender, + ctx.sender, notify_chat_id, notify_message_id, &format!("Task failed after retries: {message}"), @@ -1318,7 +1343,10 @@ pub async fn enqueue_retry(queue: &PersistentTaskQueue, task: Task, delay_second } /// Queue entry point: parses the stored task and dispatches. -pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { +pub async fn handle_task( + ctx: &AppContext<'_>, + payload: serde_json::Value, +) -> Result<(), QueueError> { let task: Task = match serde_json::from_value(payload.clone()) { Ok(task) => task, Err(e) => { @@ -1328,10 +1356,9 @@ pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { }); } }; - let bot = BOT.clone(); match task { Task::SendMediaSequence { .. } | Task::SendAnimation { .. } => { - let message_ids = match send_media_or_animation(&bot, &task).await { + let message_ids = match send_media_or_animation(ctx, &task).await { Ok(ids) => ids, Err(SendError::Retryable { delay_seconds, @@ -1343,9 +1370,7 @@ pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { }); } Err(SendError::Permanent { message, task }) => { - invalidate_cache(&task).await; - // The task settles here: drop any keep-alive temp media. - release_keep_alive(&task); + settle_task(ctx, &task, Settled::Failed).await; return Err(QueueError::Permanent { message, payload: serde_json::to_value(task).expect("task serializes"), @@ -1359,11 +1384,11 @@ pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { // whole sequence (every batch) completed, so the channel forward // and the edit-before-forward prompt must not be lost just // because the send needed a retry. - post_send_actions(&bot, &task, message_ids).await; - release_keep_alive(&task); + post_send_actions(ctx, &task, message_ids).await; + settle_task(ctx, &task, Settled::Sent).await; Ok(()) } - Task::ForwardMessages { .. } => match forward_messages(&bot, &task).await { + Task::ForwardMessages { .. } => match forward_messages(ctx, &task).await { Ok(()) => Ok(()), Err(SendError::Retryable { delay_seconds, @@ -1373,7 +1398,7 @@ pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { payload: serde_json::to_value(task).expect("task serializes"), }), Err(SendError::Permanent { message, task }) => { - release_keep_alive(&task); + settle_task(ctx, &task, Settled::Failed).await; Err(QueueError::Permanent { message, payload: serde_json::to_value(task).expect("task serializes"), @@ -1383,43 +1408,39 @@ pub async fn handle_task(payload: serde_json::Value) -> Result<(), QueueError> { } } -async fn send_media_or_animation( - sender: &dyn MediaSender, - task: &Task, -) -> Result, SendError> { +async fn send_media_or_animation(ctx: &AppContext<'_>, task: &Task) -> Result, SendError> { match task { - Task::SendMediaSequence { .. } => send_media_sequence(sender, task).await, - Task::SendAnimation { .. } => send_animation(sender, task).await, + Task::SendMediaSequence { .. } => send_media_sequence(ctx, task).await, + Task::SendAnimation { .. } => send_animation(ctx, task).await, Task::ForwardMessages { .. } => unreachable!(), } } -/// Dead-letter callback wired to the queue in main: notifies the task's chat. -pub async fn dead_letter_notify(payload: serde_json::Value, message: String) { - // A dead-lettered task never runs again. The queue dead-letters retry - // exhaustion itself (the handler is not called again), so this is the - // only place that sees the final payload — release the keep-alive temp - // media the fetch pipeline handed over, or it lives until process exit. +/// Dead-letter callback wired to the queue in main: settles the task and +/// notifies its chat. +pub async fn dead_letter_notify(ctx: &AppContext<'_>, payload: serde_json::Value, message: String) { + // A dead-lettered task never runs again, and the queue dead-letters retry + // exhaustion itself (the handler is not called again), so this is the only + // place that sees the final payload. if let Ok(task) = serde_json::from_value::(payload.clone()) { - release_keep_alive(&task); + settle_task(ctx, &task, Settled::Failed).await; } let notify_chat_id = payload.get("notify_chat_id").and_then(|v| v.as_i64()); let notify_message_id = payload.get("notify_message_id").and_then(|v| v.as_i64()); - if notify_chat_id.is_some() { - let bot = BOT.clone(); - notify_failure( - &bot, - notify_chat_id, - notify_message_id, - &format!("Task failed after retries: {message}"), - ) - .await; - } + notify_failure( + ctx.sender, + notify_chat_id, + notify_message_id, + &format!("Task failed after retries: {message}"), + ) + .await; } #[cfg(test)] mod tests { use super::*; + use crate::ctx::test_support::TestStores; + use std::time::Duration; #[test] fn oversized_photo_boundary() { @@ -1734,8 +1755,10 @@ mod tests { vec![Outcome::GroupErr, Outcome::GroupErr], media_fetch_error, ); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); let task = sequence_task(file.to_str().unwrap()); - let result = send_media_sequence(&sender, &task).await; + let result = send_media_sequence(&ctx, &task).await; assert!( matches!(result, Err(SendError::Permanent { .. })), "got {result:?}" @@ -1755,8 +1778,10 @@ mod tests { let sender = MockSender::scripted(vec![Outcome::GroupErr], || { RequestError::RetryAfter(Seconds::from_seconds(7)) }); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); let task = sequence_task(file.to_str().unwrap()); - let result = send_media_sequence(&sender, &task).await; + let result = send_media_sequence(&ctx, &task).await; match result { Err(SendError::Retryable { delay_seconds, .. }) => { assert_eq!(delay_seconds, 7.0) @@ -1791,7 +1816,9 @@ mod tests { notify_message_id: Some(2), cache_data: None, }; - let result = send_animation(&sender, &task).await; + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let result = send_animation(&ctx, &task).await; assert!( matches!(result, Err(SendError::Permanent { .. })), "got {result:?}" @@ -1807,11 +1834,14 @@ mod tests { let file = dir.path().join("media.jpg"); std::fs::write(&file, b"not-a-real-jpeg").unwrap(); let sender = MockSender::scripted(vec![Outcome::GroupOk], media_fetch_error); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); let task = sequence_task(file.to_str().unwrap()); - let result = send_media_sequence(&sender, &task).await; + let result = send_media_sequence(&ctx, &task).await; assert!(result.is_ok(), "got {result:?}"); let sender = MockSender::scripted(vec![Outcome::CopyOk], media_fetch_error); + let ctx = stores.ctx(&sender); let task = Task::ForwardMessages { from_chat_id: 1, to_chat_id: 2, @@ -1819,7 +1849,7 @@ mod tests { notify_chat_id: None, notify_message_id: None, }; - assert!(forward_messages(&sender, &task).await.is_ok()); + assert!(forward_messages(&ctx, &task).await.is_ok()); } #[tokio::test] @@ -1836,7 +1866,9 @@ mod tests { let sender = MockSender::scripted(vec![Outcome::CopyErr], || { RequestError::RetryAfter(Seconds::from_seconds(7)) }); - match forward_messages(&sender, &task).await { + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + match forward_messages(&ctx, &task).await { Err(SendError::Retryable { delay_seconds, .. }) => { assert_eq!(delay_seconds, 7.0) } @@ -1848,8 +1880,9 @@ mod tests { "Bad Request: message is not modified".into(), )) }); + let ctx = stores.ctx(&sender); assert!(matches!( - forward_messages(&sender, &task).await, + forward_messages(&ctx, &task).await, Err(SendError::Permanent { .. }) )); } @@ -1876,12 +1909,184 @@ mod tests { let dir_path = dir.path().to_path_buf(); KEEP_ALIVE.lock().push(dir); + // No chat to notify → the notify path sends nothing (its mock would + // have no scripted outcome left). + let sender = MockSender::scripted(vec![Outcome::MessageErr], media_fetch_error); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); let payload = serde_json::to_value(&task).unwrap(); - dead_letter_notify(payload, "task failed after 2 retries".into()).await; + dead_letter_notify(&ctx, payload, "task failed after 2 retries".into()).await; assert!( !KEEP_ALIVE.lock().iter().any(|dir| dir.path() == dir_path), "dead-lettered task kept its temp media alive" ); } + + #[tokio::test] + async fn post_send_forwards_immediately_when_configured() { + let sender = MockSender::scripted(vec![Outcome::CopyOk], media_fetch_error); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let task = sent_task(Some(2), false); + + post_send_actions(&ctx, &task, vec![10, 11]).await; + + assert_eq!(sender.calls(), vec!["copy_messages"]); + assert_eq!(stores.queued_tasks().await, 0); + } + + #[tokio::test] + async fn post_send_queues_a_retryable_forward() { + use teloxide::types::Seconds; + let sender = MockSender::scripted(vec![Outcome::CopyErr], || { + RequestError::RetryAfter(Seconds::from_seconds(7)) + }); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let task = sent_task(Some(2), false); + + post_send_actions(&ctx, &task, vec![10, 11]).await; + + assert_eq!(sender.calls(), vec!["copy_messages"]); + assert_eq!(stores.queued_tasks().await, 1, "forward retry not queued"); + let payload = stores.queued_payload().await; + assert_eq!(payload["type"], "forward_messages"); + assert_eq!(payload["to_chat_id"], 2); + assert_eq!(payload["message_ids"], serde_json::json!([10, 11])); + } + + #[tokio::test] + async fn post_send_notifies_a_permanent_forward_failure() { + let sender = MockSender::scripted(vec![Outcome::CopyErr, Outcome::MessageErr], || { + RequestError::Api(ApiError::Unknown("Bad Request: chat not found".into())) + }); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let task = sent_task(Some(2), true); + + post_send_actions(&ctx, &task, vec![10, 11]).await; + + // The copy failed permanently → the chat is told, nothing is queued. + assert_eq!(sender.calls(), vec!["copy_messages", "send_message"]); + assert_eq!(stores.queued_tasks().await, 0); + } + + // ── Settlement: the invariant every terminal path owes ────────────── + + /// An already-sent sequence task with the post-send knobs set: the state + /// `post_send_actions` branches on. + fn sent_task(forward_channel_id: Option, notify: bool) -> Task { + Task::SendMediaSequence { + chat_id: 1, + reply_to_message_id: 2, + caption: "cap".into(), + media_batches: vec![vec![MediaItemPayload::Photo { + media: "https://p/1.jpg".into(), + has_spoiler: false, + fallback_url: None, + file_id: false, + }]], + batch_index: 0, + sent_message_ids: vec![], + source_url: "https://x.com/u/status/1".into(), + edit_before_forward: false, + forward_channel_id, + notify_chat_id: notify.then_some(1), + notify_message_id: notify.then_some(2), + cache_data: None, + } + } + + /// A task whose media are cached Telegram file ids (the only kind that can + /// hold a link-cache entry). + fn cached_sequence_task() -> Task { + Task::SendMediaSequence { + chat_id: 1, + reply_to_message_id: 2, + caption: "cap".into(), + media_batches: vec![vec![MediaItemPayload::Photo { + media: "AgAC-file-id".into(), + has_spoiler: false, + fallback_url: None, + file_id: true, + }]], + batch_index: 0, + sent_message_ids: vec![], + source_url: "https://x.com/u/status/1".into(), + edit_before_forward: false, + forward_channel_id: None, + notify_chat_id: None, + notify_message_id: None, + cache_data: Some(CachedPost { + url: "https://x.com/u/status/1".into(), + caption: "cap".into(), + title: "t".into(), + author: "a".into(), + author_url: "au".into(), + tags: String::new(), + sensitive: false, + media: vec![CachedMedia { + kind: CachedMediaKind::Photo, + file_id: "AgAC-file-id".into(), + }], + }), + } + } + + #[tokio::test] + async fn settled_sent_keeps_the_cache_entry() { + let sender = MockSender::scripted(vec![], media_fetch_error); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let task = cached_sequence_task(); + stores + .link_cache() + .put("twitter:1", &cached_sequence_cache_data()) + .await; + + settle_task(&ctx, &task, Settled::Sent).await; + + assert!( + stores + .link_cache() + .get("twitter:1", Duration::from_secs(3600)) + .await + .is_some(), + "a successful send must not drop its own cache entry" + ); + } + + #[tokio::test] + async fn settled_failed_drops_the_cache_entry() { + let sender = MockSender::scripted(vec![], media_fetch_error); + let stores = TestStores::new(); + let ctx = stores.ctx(&sender); + let task = cached_sequence_task(); + stores + .link_cache() + .put("twitter:1", &cached_sequence_cache_data()) + .await; + + settle_task(&ctx, &task, Settled::Failed).await; + + assert!( + stores + .link_cache() + .get("twitter:1", Duration::from_secs(3600)) + .await + .is_none(), + "a permanently failed cached send must drop the entry" + ); + } + + fn cached_sequence_cache_data() -> CachedPost { + match cached_sequence_task() { + Task::SendMediaSequence { + cache_data: Some(post), + .. + } => post, + other => panic!("expected a cached sequence task, got {other:?}"), + } + } }