diff --git a/crates/xmedia-bot/src/handlers/mod.rs b/crates/xmedia-bot/src/handlers/mod.rs index 125c8fa..2ddcc03 100644 --- a/crates/xmedia-bot/src/handlers/mod.rs +++ b/crates/xmedia-bot/src/handlers/mod.rs @@ -19,6 +19,7 @@ pub(crate) use inline::prune_idle_states; /// The resolved `$DATA_DIR/task_queue.db` path, for the startup config line. pub(crate) use statics::db_path; pub use statics::{CHAT_STORE, CONFIG, LINK_CACHE, TASK_QUEUE}; +pub(crate) use urls::repair_lost_local_media; pub use urls::{start_url_workers, stop_url_workers}; use crate::ctx::AppContext; diff --git a/crates/xmedia-bot/src/handlers/urls.rs b/crates/xmedia-bot/src/handlers/urls.rs index 6551504..c2a2684 100644 --- a/crates/xmedia-bot/src/handlers/urls.rs +++ b/crates/xmedia-bot/src/handlers/urls.rs @@ -619,6 +619,217 @@ async fn url_media_inner( } } +// ── Startup repair: queued retries whose local media did not survive ─────── + +/// A post's fresh media plus the caption and cache snapshot that go with them: +/// what [`refetch`] hands [`apply_refresh`]. Plain data, so the rewrite below +/// can be tested without a network fetch (which cannot be faked here: +/// [`x_media::site::Fetched`] keeps a private field and is not constructible +/// outside its crate). +struct Refetched { + caption: String, + items: Vec, + cache_data: Option, +} + +/// Whether a queued task should have its post re-fetched, because it still +/// wants a local file (ugoira MP4, a bsky remux, a downloaded temp file) that is +/// gone. Those files live in the system temp dir and the registry that keeps +/// them alive for the retry (`send::KEEP_ALIVE`) is in memory, so a restart +/// takes all of them — a retry that needs one can only dead-letter. +/// +/// A partially delivered album is left alone: its remaining batches cannot be +/// reconciled with a fresh media list without risking a second copy of what the +/// user already received. +fn needs_refetch(task: &Task) -> bool { + if let Task::SendMediaSequence { + batch_index, + sent_message_ids, + .. + } = task + && (*batch_index > 0 || !sent_message_ids.is_empty()) + { + return false; + } + task.local_media_paths().iter().any(|path| !path.exists()) +} + +/// Rebuilds the task from the fresh media, keeping its delivery envelope (chat, +/// reply, forward/edit settings, notify targets): the retry that was queued must +/// still deliver the same way, whoever asked for it. +fn apply_refresh(task: &Task, fresh: &Refetched) -> Option { + let chat_id = task.chat_id()?; + let (edit_before_forward, forward_channel_id) = match task { + Task::SendMediaSequence { + edit_before_forward, + forward_channel_id, + .. + } + | Task::SendAnimation { + edit_before_forward, + forward_channel_id, + .. + } => (*edit_before_forward, *forward_channel_id), + Task::ForwardMessages { .. } => return None, + }; + let reply_to_message_id = match task { + Task::SendMediaSequence { + reply_to_message_id, + .. + } + | Task::SendAnimation { + reply_to_message_id, + .. + } => *reply_to_message_id, + Task::ForwardMessages { .. } => return None, + }; + let (notify_chat_id, notify_message_id) = task.notify_target(); + let items = fresh.items.clone(); + let source_url = task.source_url()?.to_string(); + Some( + if items.len() == 1 && matches!(items[0], MediaItemPayload::Animation { .. }) { + Task::SendAnimation { + chat_id, + reply_to_message_id, + caption: fresh.caption.clone(), + animation: items.into_iter().next().expect("checked len"), + source_url, + edit_before_forward, + forward_channel_id, + notify_chat_id, + notify_message_id, + cache_data: fresh.cache_data.clone(), + } + } else { + Task::SendMediaSequence { + chat_id, + reply_to_message_id, + caption: fresh.caption.clone(), + media_batches: send::chunk_media_items(send::photos_first(items)), + // A fresh delivery: nothing of this payload has been sent. + batch_index: 0, + sent_message_ids: vec![], + source_url, + edit_before_forward, + forward_channel_id, + notify_chat_id, + notify_message_id, + cache_data: fresh.cache_data.clone(), + } + }, + ) +} + +/// Fetches the post again and maps it into [`Refetched`]: the same mapping the +/// fresh-fetch path uses (per-site caption format from the chat, render fields +/// for the link-cache snapshot), so a repaired task looks like a first send. +async fn refetch( + ctx: &AppContext<'_>, + chat_id: i64, + url: &str, +) -> Result, x_media::site::FetchError> { + let Some(mut fetched) = x_media::site::fetch(url).await? else { + return Ok(None); + }; + if fetched.media.is_empty() { + return Ok(None); + } + let chat_data = ctx.chat_store.get(chat_id).await; + let format = chat_data + .message_format + .get(fetched.site_name()) + .cloned() + .unwrap_or_default(); + let caption = fetched.caption_with(&format); + let cache_data = fetched + .render_fields() + .map(|(author, author_url, title, content, tags)| CachedPost { + url: fetched.source_url.clone(), + caption: fetched.caption.clone(), + title: title.to_string(), + content: content.to_string(), + author: author.to_string(), + author_url: author_url.to_string(), + tags: tags.to_string(), + sensitive: fetched.sensitive, + media: vec![], + }); + let items: Vec = fetched + .media + .iter() + .map(|media| media_to_payload(media, fetched.sensitive)) + .collect(); + // The re-fetch may produce a fresh local file (ugoira / bsky remux): hand it + // to the same keep-alive registry the first fetch uses. + if let Some(dir) = fetched.take_keep_alive() { + send::KEEP_ALIVE.lock().push(dir); + } + Ok(Some(Refetched { + caption, + items, + cache_data, + })) +} + +/// Re-fetches every queued task whose local media did not survive the restart, +/// so the user's link is still delivered instead of dead-lettering on a file +/// that cannot come back. Returns how many rows were rewritten. +/// +/// Startup only, before the queue workers start: no worker can lease a row while +/// this writes, which is what lets it replace payloads without the lease-token +/// guard every worker write-back carries. +pub(crate) async fn repair_lost_local_media(ctx: &AppContext<'_>) -> usize { + let mut repaired = 0; + for (id, payload) in ctx.task_queue.runnable_rows().await { + let Ok(task) = serde_json::from_str::(&payload) else { + continue; + }; + if !needs_refetch(&task) { + continue; + } + let (Some(url), Some(chat_id)) = (task.source_url().map(str::to_string), task.chat_id()) + else { + continue; + }; + match refetch(ctx, chat_id, &url).await { + Ok(Some(fresh)) => { + let Some(updated) = apply_refresh(&task, &fresh) else { + continue; + }; + let updated = serde_json::to_value(&updated).expect("task serializes"); + if ctx.task_queue.replace_payload(&id, &updated).await { + repaired += 1; + log::info!( + "startup repair: re-fetched [key={}] for chat={chat_id} (its local media did not survive the restart)", + log_key(&url) + ); + } + } + // The post is gone or withheld now: the retry could not have + // delivered anything either, so say why instead of letting it + // dead-letter on a missing file. + Ok(None) | Err(_) => { + let (notify_chat_id, notify_message_id) = task.notify_target(); + log::warn!( + "startup repair: [key={}] for chat={chat_id} needed a re-fetch and none was possible", + log_key(&url) + ); + send::notify_failure( + ctx.sender, + notify_chat_id, + notify_message_id, + &format!( + "{} — the media held for retry was lost when the bot restarted and the post could not be fetched again. Please send the link again.", + log_key(&url) + ), + ) + .await; + } + } + } + repaired +} + #[cfg(test)] mod tests { use super::*; @@ -818,6 +1029,178 @@ mod tests { assert_eq!(notify_chat_id, Some(1)); } + fn queued_task(media: &str, batch_index: usize, sent: Vec) -> Task { + Task::SendMediaSequence { + chat_id: 1, + reply_to_message_id: 2, + caption: "cap".into(), + media_batches: vec![vec![MediaItemPayload::Photo { + media: media.to_string(), + has_spoiler: false, + fallback_url: None, + file_id: false, + }]], + batch_index, + sent_message_ids: sent, + source_url: "https://x.com/u/status/1".into(), + edit_before_forward: true, + forward_channel_id: Some(2), + notify_chat_id: Some(1), + notify_message_id: Some(2), + cache_data: None, + } + } + + #[test] + fn only_tasks_missing_a_local_file_need_a_refetch() { + // A URL send needs nothing. + assert!(!needs_refetch(&queued_task("https://cdn/1.jpg", 0, vec![]))); + // A local path that is still there (a survived temp file) needs nothing. + let dir = tempfile::tempdir().unwrap(); + let alive = dir.path().join("ugoira.mp4"); + std::fs::write(&alive, b"x").unwrap(); + assert!(!needs_refetch(&queued_task( + alive.to_str().unwrap(), + 0, + vec![] + ))); + // A local path the restart took away does. + assert!(needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 0, + vec![] + ))); + // A partially delivered album is left to its own retry path. + assert!(!needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 1, + vec![7] + ))); + assert!(!needs_refetch(&queued_task( + "/nonexistent-ugoira.mp4", + 0, + vec![7] + ))); + // A channel copy holds no media. + assert!(!needs_refetch(&Task::ForwardMessages { + from_chat_id: 1, + to_chat_id: 2, + message_ids: vec![3], + notify_chat_id: None, + notify_message_id: None, + })); + } + + #[test] + fn apply_refresh_keeps_the_delivery_envelope() { + let task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); + let fresh = Refetched { + caption: "fresh caption".into(), + items: vec![MediaItemPayload::Photo { + media: "https://cdn/fresh.jpg".into(), + has_spoiler: true, + fallback_url: None, + file_id: false, + }], + cache_data: None, + }; + match apply_refresh(&task, &fresh).expect("a repairable task") { + Task::SendMediaSequence { + chat_id, + reply_to_message_id, + caption, + media_batches, + batch_index, + sent_message_ids, + source_url, + edit_before_forward, + forward_channel_id, + notify_chat_id, + notify_message_id, + .. + } => { + // Same delivery: chat, reply, forward/edit settings, notify. + assert_eq!((chat_id, reply_to_message_id), (1, 2)); + assert!(edit_before_forward); + assert_eq!(forward_channel_id, Some(2)); + assert_eq!((notify_chat_id, notify_message_id), (Some(1), Some(2))); + assert_eq!(source_url, "https://x.com/u/status/1"); + // Fresh media, and nothing of it counted as sent yet. + assert_eq!(caption, "fresh caption"); + assert!( + matches!( + &media_batches[0][0], + MediaItemPayload::Photo { media, .. } if media == "https://cdn/fresh.jpg" + ), + "fresh media must replace the lost local file" + ); + assert!(matches!( + media_batches[0][0], + MediaItemPayload::Photo { + has_spoiler: true, + .. + } + )); + assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); + } + other => panic!("expected a media sequence, got {other:?}"), + } + } + + /// The whole repair against a real post: a queued row whose media is a local + /// file the restart took away is re-fetched from its `source_url` and + /// rewritten in place, so the retry can still deliver it. + #[tokio::test] + #[ignore = "live network: requires outbound HTTPS to public.api.bsky.app"] + async fn a_lost_local_media_row_is_refetched_from_its_post() { + let stores = TestStores::new(); + // An empty script: the repair must not need to tell the user anything. + let sender = MockSender::scripted(vec![], permanent_error); + let ctx = stores.ctx(&sender); + let mut task = queued_task("/nonexistent-ugoira.mp4", 0, vec![]); + if let Task::SendMediaSequence { source_url, .. } = &mut task { + *source_url = "https://bsky.app/profile/fu-futa.bsky.social/post/3laoveufjv224".into(); + } + stores + .task_queue() + .enqueue(serde_json::to_value(&task).unwrap(), crate::db::now_f64()) + .await + .unwrap(); + + assert_eq!(repair_lost_local_media(&ctx).await, 1); + + let updated: Task = serde_json::from_value(stores.queued_payload().await).unwrap(); + match updated { + Task::SendMediaSequence { + media_batches, + batch_index, + sent_message_ids, + caption, + .. + } => { + let media: Vec = media_batches + .iter() + .flatten() + .map(|item| match item { + MediaItemPayload::Photo { media, .. } + | MediaItemPayload::Video { media, .. } + | MediaItemPayload::Animation { media, .. } => media.clone(), + }) + .collect(); + assert!(!media.is_empty(), "the fresh fetch yielded no media"); + assert!( + media.iter().all(|m| m.starts_with("http")), + "the retry must be uploadable from URLs again: {media:?}" + ); + assert_eq!((batch_index, sent_message_ids.len()), (0, 0)); + assert!(!caption.is_empty()); + } + other => panic!("expected a repaired media sequence, got {other:?}"), + } + // The post was re-read, not re-delivered: nothing was sent. + assert!(sender.calls().is_empty(), "{:?}", sender.calls()); + } + #[test] fn fetch_errors_map_to_distinct_user_messages() { use x_media::site::FetchError; diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 6f5daa4..44a553c 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -152,6 +152,16 @@ async fn main() { ); log::debug!("config: admin ids {:?}", CONFIG.admin_ids); + // Startup repair, before any worker runs: a queued retry whose media was a + // local file (ugoira MP4, bsky remux, a downloaded temp file) can never + // succeed after a restart — the registry that kept those files alive is in + // memory — so those rows are re-fetched from their post instead of + // dead-lettering the user's link. + let repaired = handlers::repair_lost_local_media(&CONTEXT).await; + if repaired > 0 { + log::info!("startup repair: re-fetched {repaired} queued task(s)"); + } + // Queue worker: handles typed tasks, dead-letters failed sends to the // task's chat. Both closures use the shared context (the queue requires // 'static handlers, and the statics are process-wide anyway). diff --git a/crates/xmedia-bot/src/queue.rs b/crates/xmedia-bot/src/queue.rs index a213e79..c46ddab 100644 --- a/crates/xmedia-bot/src/queue.rs +++ b/crates/xmedia-bot/src/queue.rs @@ -235,6 +235,61 @@ impl PersistentTaskQueue { /// `earliest_run_after`: that one runs on every idle worker cycle and must /// stay a single indexed `MIN`, while the count is only asked for once per /// sweep. + /// `(id, payload)` of every row that can still run (`pending`, + /// `in_progress`). The startup repair reads these before the workers start: + /// with no worker running, no row can be leased while it writes. + pub async fn runnable_rows(&self) -> Vec<(String, String)> { + let result = self + .pool + .with_conn(|conn| { + let mut stmt = conn.prepare( + "SELECT id, payload FROM tasks WHERE status IN ('pending', 'in_progress') ORDER BY run_after", + )?; + let rows = stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?; + rows.collect::>>() + }) + .await; + match result { + Ok(rows) => rows, + Err(e) => { + log::error!("queue row scan failed: {e}"); + Vec::new() + } + } + } + + /// Replaces a runnable row's payload and restarts its attempt budget: the + /// new payload is a fresh delivery of the same task, so the retries it has + /// already spent do not carry over. Startup repair only — a worker's + /// write-back is lease-token guarded instead (`replace_payload` cannot race + /// one: it runs before any worker does). + pub async fn replace_payload(&self, id: &str, payload: &Value) -> bool { + let logged_id = id.to_string(); + let (id, payload) = (id.to_string(), payload.to_string()); + let result = self + .pool + .with_conn(move |conn| { + let affected = conn.execute( + "UPDATE tasks SET payload=?1, attempts=0, run_after=?2, status='pending', locked_until=0 \ + WHERE id=?3 AND status IN ('pending', 'in_progress')", + params![payload, now_f64(), id], + )?; + Ok(affected == 1) + }) + .await; + match result { + Ok(true) => true, + Ok(false) => { + log::warn!("queue: row {logged_id} vanished before its payload could be replaced"); + false + } + Err(e) => { + log::error!("queue payload replace failed for {logged_id}: {e}"); + false + } + } + } + pub async fn pending_backlog(&self) -> Option<(i64, f64)> { let result = self .pool @@ -894,6 +949,51 @@ mod tests { queue.stop().await; } + #[tokio::test] + async fn runnable_rows_and_payload_replacement() { + let (queue, _dir) = new_queue().await; + queue + .enqueue(serde_json::json!({"s": 1}), now_f64()) + .await + .unwrap(); + let rows = queue.runnable_rows().await; + assert_eq!(rows.len(), 1); + let (id, payload) = rows[0].clone(); + assert_eq!(payload, "{\"s\":1}"); + + // A replacement restarts the attempt budget (the new payload is a fresh + // delivery, not the continuation of the old one). + { + let id_owned = id.clone(); + queue + .pool + .with_conn(move |conn| { + conn.execute("UPDATE tasks SET attempts=2 WHERE id=?1", params![id_owned])?; + Ok(()) + }) + .await + .unwrap(); + } + assert!( + queue + .replace_payload(&id, &serde_json::json!({"s": 2})) + .await + ); + let rows = queue.runnable_rows().await; + assert_eq!(rows[0].1, "{\"s\":2}"); + assert_eq!( + queue.pending_backlog().await.map(|(n, _)| n), + Some(1), + "a repaired row is pending work again" + ); + // A row that is gone (or done) is not rewritten. + assert!( + !queue + .replace_payload("task_missing", &serde_json::json!({})) + .await + ); + } + #[tokio::test] async fn permanent_error_dead_letters_immediately() { let (queue, _dir) = new_queue().await; diff --git a/crates/xmedia-bot/src/send/mod.rs b/crates/xmedia-bot/src/send/mod.rs index 172ed6c..6282710 100644 --- a/crates/xmedia-bot/src/send/mod.rs +++ b/crates/xmedia-bot/src/send/mod.rs @@ -29,7 +29,7 @@ use upload::{FallbackError, PreparedItem, prepare_upload_item, send_batch_via_up // parts other modules use so call sites stay `send::x`. pub(crate) use post_send::{ EDIT_PROMPT_EXPIRED_TEXT, KEEP_ALIVE, Settled, dead_letter_notify, enqueue_retry, handle_task, - post_send_actions, settle_task, + notify_failure, post_send_actions, settle_task, }; /// One process-wide Bot for queue workers. Building a fresh Bot (and its HTTP @@ -143,7 +143,7 @@ impl Task { } } - fn source_url(&self) -> Option<&str> { + pub(crate) fn source_url(&self) -> Option<&str> { match self { Task::SendMediaSequence { source_url, .. } | Task::SendAnimation { source_url, .. } => { Some(source_url) @@ -173,9 +173,42 @@ impl Task { } } + /// The chat this task delivers media to (`None` for a channel copy, which + /// names two chats instead). + pub(crate) fn chat_id(&self) -> Option { + match self { + Task::SendMediaSequence { chat_id, .. } | Task::SendAnimation { chat_id, .. } => { + Some(*chat_id) + } + Task::ForwardMessages { .. } => None, + } + } + + /// Where a failure notice for this task goes (both `None` for a copy with + /// nothing to notify). + pub(crate) fn notify_target(&self) -> (Option, Option) { + match self { + Task::SendMediaSequence { + notify_chat_id, + notify_message_id, + .. + } + | Task::SendAnimation { + notify_chat_id, + notify_message_id, + .. + } + | Task::ForwardMessages { + notify_chat_id, + notify_message_id, + .. + } => (*notify_chat_id, *notify_message_id), + } + } + /// Local file paths referenced by this task's media (ugoira / bsky remux /// MP4 and the like); empty for URL or Telegram file-id sends. - fn local_media_paths(&self) -> Vec { + pub(crate) fn local_media_paths(&self) -> Vec { let mut out = Vec::new(); for item in self.media_items() { let is_file_id = match item { diff --git a/crates/xmedia-bot/src/send/post_send.rs b/crates/xmedia-bot/src/send/post_send.rs index df020d4..4597ef2 100644 --- a/crates/xmedia-bot/src/send/post_send.rs +++ b/crates/xmedia-bot/src/send/post_send.rs @@ -175,7 +175,7 @@ pub(super) fn hidden_template_count(templates: &HashMap) -> usiz /// Notifies a chat about a dead-lettered task (skips when `notify_chat_id` is /// absent). -pub(super) async fn notify_failure( +pub(crate) async fn notify_failure( sender: &dyn MediaSender, chat_id: Option, message_id: Option,