From e5b39edb20b8ce59cad86849b13750f3cc86206e Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:20:26 +0200 Subject: [PATCH] Keep the bot token out of stored image edit jobs VIP image edit jobs stored Telegram file URLs, which carry the bot token, in taskman_jobs, and the same URLs reached Discovery job bodies. - Edit jobs store the pictures' Telegram file IDs (`image_file_ids`); the resolver only checks that the bot can fetch each file. - The VIP edit worker loads the pictures by file ID and hands the editor base64 bytes, so draw-api gets `image_b64` and no token URL. Queued jobs with stored file URLs load the same way. - Migration 211 removes the token from file URLs already stored in taskman_jobs (one-way; the redacted form still resolves). --- crates/openplotva-app/src/dialog_messages.rs | 32 +-- crates/openplotva-app/src/dialog_tools.rs | 138 +++++------ .../src/image_edit_references.rs | 221 +++++++++++++++++- crates/openplotva-app/src/image_jobs.rs | 15 ++ crates/openplotva-app/src/lib.rs | 30 +-- crates/openplotva-taskman/src/lib.rs | 31 +++ ...elegram_file_urls_in_taskman_jobs.down.sql | 3 + ..._telegram_file_urls_in_taskman_jobs.up.sql | 11 + 8 files changed, 362 insertions(+), 119 deletions(-) create mode 100644 migrations/211_redact_telegram_file_urls_in_taskman_jobs.down.sql create mode 100644 migrations/211_redact_telegram_file_urls_in_taskman_jobs.up.sql diff --git a/crates/openplotva-app/src/dialog_messages.rs b/crates/openplotva-app/src/dialog_messages.rs index 9ba43993..1531d12b 100644 --- a/crates/openplotva-app/src/dialog_messages.rs +++ b/crates/openplotva-app/src/dialog_messages.rs @@ -5694,7 +5694,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( crate::dialog_tools::ImageEditFileSelection { photo_file_id: "latest-photo-file".to_owned(), - photo_urls: vec!["https://files.test/photo.png".to_owned()], + photo_file_ids: vec!["latest-photo-file".to_owned()], }, )); let image_scheduler = @@ -5797,10 +5797,8 @@ mod tests { assert_eq!(image.original_text, "contrast"); assert_eq!(image.author, "Ada"); assert_eq!(image.image_file_id, "latest-photo-file"); - assert_eq!( - image.image_urls, - vec!["https://files.test/photo.png".to_owned()] - ); + assert!(image.image_urls.is_empty()); + assert_eq!(image.image_file_ids, vec!["latest-photo-file".to_owned()]); assert!(image.is_image_edit); record.id }; @@ -5849,7 +5847,8 @@ mod tests { thread_id: None, prompt: "contrast".to_owned(), photo_file_id: "latest-photo-file".to_owned(), - photo_urls: vec!["https://files.test/photo.png".to_owned()], + photo_urls: Vec::new(), + photo_file_ids: vec!["latest-photo-file".to_owned()], }] ); // Album delivery: one placeholder frame (single editor) filled with the @@ -5883,14 +5882,14 @@ mod tests { Ok(()) } - struct AlbumFileUrlStub; + struct AlbumFileProbeStub; - impl crate::dialog_tools::ImageEditFileUrlProvider for AlbumFileUrlStub { - fn image_edit_file_url<'a>( + impl crate::dialog_tools::ImageEditFileProbe for AlbumFileProbeStub { + fn image_edit_file_available<'a>( &'a self, - file_id: &'a str, - ) -> crate::dialog_tools::ImageEditFileUrlFuture<'a> { - Box::pin(async move { Ok(Some(format!("https://files.test/{file_id}"))) }) + _file_id: &'a str, + ) -> crate::dialog_tools::ImageEditFileProbeFuture<'a> { + Box::pin(async { Ok(true) }) } } @@ -5918,7 +5917,7 @@ mod tests { let queue = Arc::new(InMemoryTaskQueue::new()); let resolver = crate::dialog_tools::TelegramImageEditFileResolver::new( files.clone(), - AlbumFileUrlStub, + AlbumFileProbeStub, ) .with_media_group_registry(crate::dialog_tools::MediaGroupImageRegistry::new( Duration::ZERO, @@ -5930,7 +5929,7 @@ mod tests { .with_image_edit_file_resolver(Arc::new(resolver)); let scheduler = SchedulerStub::default(); let mut updates = Vec::new(); - let mut expected_urls = Vec::new(); + let mut expected_file_ids = Vec::new(); let mut unique_ids = Vec::new(); for index in 0..count { let id = format!("serial-file-{bot_id}-{case}-{index}"); @@ -5960,7 +5959,7 @@ mod tests { }) .await?; unique_ids.push(id.clone()); - expected_urls.push(format!("https://files.test/{id}")); + expected_file_ids.push(id.clone()); updates.push(update); } for (index, update) in updates.into_iter().enumerate() { @@ -5989,7 +5988,8 @@ mod tests { assert_eq!(records.len(), 1); let data = &records[0].job.data; let image = data.image_data.as_ref().expect("image edit"); - assert_eq!(image.image_urls, expected_urls); + assert!(image.image_urls.is_empty()); + assert_eq!(image.image_file_ids, expected_file_ids); assert_eq!(image.prompt, prompt); assert_eq!(image.original_text, prompt); assert_eq!( diff --git a/crates/openplotva-app/src/dialog_tools.rs b/crates/openplotva-app/src/dialog_tools.rs index 05986ecf..a50c422f 100644 --- a/crates/openplotva-app/src/dialog_tools.rs +++ b/crates/openplotva-app/src/dialog_tools.rs @@ -68,9 +68,9 @@ pub type DrawImageFuture<'a> = /// Boxed future returned by image-edit file resolvers. pub type ImageEditFileResolveFuture<'a> = Pin> + Send + 'a>>; -/// Boxed future returned by image-edit Telegram file URL providers. -pub type ImageEditFileUrlFuture<'a> = - Pin, ToolboxError>> + Send + 'a>>; +/// Boxed future returned by image-edit Telegram file probes. +pub type ImageEditFileProbeFuture<'a> = + Pin> + Send + 'a>>; /// Boxed future returned by VIP status checks for image scheduling. pub type DrawImageVipStatusFuture<'a> = Pin + Send + 'a>>; @@ -232,57 +232,48 @@ pub trait ImageEditFileResolver: Send + Sync { ) -> ImageEditFileResolveFuture<'a>; } -pub trait ImageEditFileUrlProvider: Send + Sync { - /// Return a Telegram file URL for the supplied file ID. - fn image_edit_file_url<'a>(&'a self, file_id: &'a str) -> ImageEditFileUrlFuture<'a>; +pub trait ImageEditFileProbe: Send + Sync { + /// Whether the bot can download the Telegram file with the supplied file ID. + fn image_edit_file_available<'a>(&'a self, file_id: &'a str) -> ImageEditFileProbeFuture<'a>; } #[derive(Clone)] -pub struct TelegramBotFileUrlProvider { +pub struct TelegramBotFileProbe { client: openplotva_telegram::TelegramClient, - bot_key: String, } -impl TelegramBotFileUrlProvider { - /// Build a Telegram file URL provider from the runtime bot client and token. +impl TelegramBotFileProbe { + /// Build a Telegram file probe from the runtime bot client. #[must_use] - pub fn new(client: openplotva_telegram::TelegramClient, bot_key: impl Into) -> Self { - Self { - client, - bot_key: bot_key.into(), - } + pub const fn new(client: openplotva_telegram::TelegramClient) -> Self { + Self { client } } } -impl fmt::Debug for TelegramBotFileUrlProvider { +impl fmt::Debug for TelegramBotFileProbe { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter - .debug_struct("TelegramBotFileUrlProvider") + .debug_struct("TelegramBotFileProbe") .finish_non_exhaustive() } } -impl ImageEditFileUrlProvider for TelegramBotFileUrlProvider { - fn image_edit_file_url<'a>(&'a self, file_id: &'a str) -> ImageEditFileUrlFuture<'a> { +impl ImageEditFileProbe for TelegramBotFileProbe { + fn image_edit_file_available<'a>(&'a self, file_id: &'a str) -> ImageEditFileProbeFuture<'a> { Box::pin(async move { let file_id = file_id.trim(); if file_id.is_empty() { - return Ok(None); + return Ok(false); } let file: carapax::types::File = self .client .execute(carapax::types::GetFile::new(file_id)) .await .map_err(|error| Box::new(error) as ToolboxError)?; - let Some(file_path) = file + Ok(file .file_path .as_deref() - .map(str::trim) - .filter(|path| !path.is_empty()) - else { - return Ok(None); - }; - Ok(Some(telegram_file_url(&self.bot_key, file_path))) + .is_some_and(|path| !path.trim().is_empty())) }) } } @@ -388,20 +379,20 @@ impl MediaGroupImageRegistry { } #[derive(Clone, Debug)] -pub struct TelegramImageEditFileResolver { +pub struct TelegramImageEditFileResolver { files: openplotva_storage::PostgresTelegramFileStore, - urls: UrlProvider, + probe: Probe, media_groups: MediaGroupImageRegistry, pending_albums: Option<(i64, openplotva_storage::PostgresTelegramDeliveryStore)>, } -impl TelegramImageEditFileResolver { +impl TelegramImageEditFileResolver { /// Build a concrete image-edit file resolver. #[must_use] - pub fn new(files: openplotva_storage::PostgresTelegramFileStore, urls: UrlProvider) -> Self { + pub fn new(files: openplotva_storage::PostgresTelegramFileStore, probe: Probe) -> Self { Self { files, - urls, + probe, media_groups: MediaGroupImageRegistry::default(), pending_albums: None, } @@ -425,9 +416,9 @@ impl TelegramImageEditFileResolver { } } -impl ImageEditFileResolver for TelegramImageEditFileResolver +impl ImageEditFileResolver for TelegramImageEditFileResolver where - UrlProvider: ImageEditFileUrlProvider, + Probe: ImageEditFileProbe, { fn capture_image_attachments( &self, @@ -485,13 +476,13 @@ where &rows, attachments, ); - let mut photo_urls = Vec::new(); - for latest_file_id in &selection.url_file_ids { - if let Ok(Some(url)) = self.urls.image_edit_file_url(latest_file_id).await { - photo_urls.push(url); + let mut photo_file_ids = Vec::new(); + for latest_file_id in selection.source_file_ids { + if let Ok(true) = self.probe.image_edit_file_available(&latest_file_id).await { + photo_file_ids.push(latest_file_id); } } - if requested_ids.len() > 1 && photo_urls.len() != requested_ids.len() { + if requested_ids.len() > 1 && photo_file_ids.len() != requested_ids.len() { return Err(Box::new(std::io::Error::other( "an album reference image could not be resolved", )) as ToolboxError); @@ -499,7 +490,7 @@ where Ok(ImageEditFileSelection { photo_file_id: selection.photo_file_id, - photo_urls, + photo_file_ids, }) }) } @@ -528,7 +519,7 @@ impl ImageEditFileResolver for openplotva_storage::PostgresTelegramFileStore { Ok(ImageEditFileSelection { photo_file_id: latest_file_ids.first().cloned().unwrap_or_default(), - photo_urls: Vec::new(), + photo_file_ids: Vec::new(), }) }) } @@ -795,7 +786,7 @@ pub struct DrawImageRateLimitReport { #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct ImageEditFileSelection { pub photo_file_id: String, - pub photo_urls: Vec, + pub photo_file_ids: Vec, } #[derive(Clone, Debug, Default, PartialEq)] @@ -1576,7 +1567,7 @@ impl TaskmanDialogToolAdapter { .await?; let source_count = image_edit_file_unique_ids(&request.attachments).len(); if files.photo_file_id.trim().is_empty() - || (source_count > 1 && files.photo_urls.len() < source_count) + || (source_count > 1 && files.photo_file_ids.len() < source_count) { return Ok(not_scheduled_draw( DrawImageScheduleRejection::EditSourceUnavailable, @@ -1607,7 +1598,8 @@ impl TaskmanDialogToolAdapter { user_full_name: request.user_full_name, prompt, photo_file_id: files.photo_file_id, - photo_urls: files.photo_urls, + photo_urls: Vec::new(), + photo_file_ids: files.photo_file_ids, thread_id: request.thread_id, }, OffsetDateTime::now_utc(), @@ -2148,7 +2140,7 @@ fn image_edit_file_unique_ids(attachments: &[ChatAttachment]) -> Vec { #[derive(Clone, Debug, Default, Eq, PartialEq)] struct ImageEditFileIdSelection { photo_file_id: String, - url_file_ids: Vec, + source_file_ids: Vec, } fn image_edit_file_id_selection( @@ -2163,7 +2155,7 @@ fn image_edit_file_id_selection( latest_file_ids_in_attachment_order(media_group_file_unique_ids, rows, attachments); ImageEditFileIdSelection { photo_file_id: target_file_ids.first().cloned().unwrap_or_default(), - url_file_ids: if media_group_file_ids.is_empty() { + source_file_ids: if media_group_file_ids.is_empty() { target_file_ids } else { media_group_file_ids @@ -2216,14 +2208,6 @@ fn latest_file_ids_in_attachment_order( latest_file_ids } -fn telegram_file_url(bot_key: &str, file_path: &str) -> String { - format!( - "https://api.telegram.org/file/bot{}/{}", - bot_key.trim(), - file_path.trim() - ) -} - fn cleanup_old_media_groups(groups: &mut HashMap, now: OffsetDateTime) { let cutoff = now - time::Duration::minutes(5); groups.retain(|_, entry| entry.created_at >= cutoff); @@ -4471,7 +4455,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "generated-goldfish-file".to_owned(), - photo_urls: vec!["https://files.test/generated-goldfish.png".to_owned()], + photo_file_ids: vec!["generated-goldfish-file".to_owned()], }, )); let scheduler = Arc::new( @@ -4552,10 +4536,8 @@ mod tests { let image = records[0].job.data.image_data.as_ref().expect("image edit"); assert!(image.is_image_edit); assert_eq!(image.image_file_id, "generated-goldfish-file"); - assert_eq!( - image.image_urls, - vec!["https://files.test/generated-goldfish.png"] - ); + assert!(image.image_urls.is_empty()); + assert_eq!(image.image_file_ids, vec!["generated-goldfish-file"]); assert_eq!(image.original_text, caption); assert_eq!(image.prompt, prompt); } @@ -4606,8 +4588,8 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "file-a".to_owned(), - photo_urls: (0..MAX_IMAGE_EDIT_SOURCE_IMAGES) - .map(|index| format!("https://files.test/{index}.png")) + photo_file_ids: (0..MAX_IMAGE_EDIT_SOURCE_IMAGES) + .map(|index| format!("file-{index}")) .collect(), }, )); @@ -4677,8 +4659,8 @@ mod tests { .chain((2..MAX_IMAGE_EDIT_SOURCE_IMAGES).map(|index| format!("photo-{index}"))) .chain(std::iter::once("photo-b".to_owned())) .collect(); - let expected_urls: Vec<_> = (0..MAX_IMAGE_EDIT_SOURCE_IMAGES) - .map(|index| format!("https://files.test/{index}.png")) + let expected_file_ids: Vec<_> = (0..MAX_IMAGE_EDIT_SOURCE_IMAGES) + .map(|index| format!("file-{index}")) .collect(); let result = openplotva_dialog::dispatch_dialog_tool( &toolbox, @@ -4712,8 +4694,8 @@ mod tests { .image_data .as_ref() .expect("image edit") - .image_urls, - expected_urls + .image_file_ids, + expected_file_ids ); input.history.push(openplotva_dialog::HistoryMessage { message_id: input.message.id, @@ -4754,7 +4736,7 @@ mod tests { .image_data .as_ref() .expect("followup edit") - .image_urls + .image_file_ids .len(), MAX_IMAGE_EDIT_SOURCE_IMAGES ); @@ -5866,7 +5848,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "edit-file".to_owned(), - photo_urls: vec!["https://files.test/photo.png".to_owned()], + photo_file_ids: vec!["edit-file".to_owned()], }, )); let adapter = TaskmanDialogToolAdapter::new(Arc::clone(&queue)) @@ -6175,7 +6157,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "latest-photo-file".to_owned(), - photo_urls: Vec::new(), + photo_file_ids: Vec::new(), }, )); let adapter = TaskmanDialogToolAdapter::new(Arc::clone(&queue)) @@ -6499,7 +6481,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "latest-photo-file".to_owned(), - photo_urls: vec!["https://files.test/photo.png".to_owned()], + photo_file_ids: vec!["latest-photo-file".to_owned()], }, )); let adapter = TaskmanDialogToolAdapter::new(Arc::clone(&queue)) @@ -6553,10 +6535,8 @@ mod tests { .expect("image edit job should carry image data"); assert_eq!(image.prompt, "make it brighter"); assert_eq!(image.image_file_id, "latest-photo-file"); - assert_eq!( - image.image_urls, - vec!["https://files.test/photo.png".to_owned()] - ); + assert!(image.image_urls.is_empty()); + assert_eq!(image.image_file_ids, vec!["latest-photo-file".to_owned()]); assert!(image.is_image_edit); Ok(()) } @@ -6593,7 +6573,7 @@ mod tests { image_edit_file_id_selection(&ids, &[], &[], &attachments), ImageEditFileIdSelection { photo_file_id: "history-second".to_owned(), - url_file_ids: vec!["history-second".to_owned(), "history-first".to_owned()], + source_file_ids: vec!["history-second".to_owned(), "history-first".to_owned()], } ); let rows = vec![telegram_file_record("second", "current-second")]; @@ -6628,10 +6608,6 @@ mod tests { latest_file_ids_in_attachment_order(&file_unique_ids, &rows, &[]), vec!["file-a".to_owned(), "file-b".to_owned()] ); - assert_eq!( - telegram_file_url("123:secret", "photos/file-a.jpg"), - "https://api.telegram.org/file/bot123:secret/photos/file-a.jpg" - ); } #[tokio::test] @@ -6642,7 +6618,7 @@ mod tests { .with_image_edit_file_resolver(Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "file-a".to_owned(), - photo_urls: vec!["https://files.test/a.png".to_owned()], + photo_file_ids: vec!["file-a".to_owned()], }, ))); let result = adapter @@ -6676,7 +6652,7 @@ mod tests { image_edit_file_id_selection(&target_unique_ids, &media_group_unique_ids, &rows, &[]), ImageEditFileIdSelection { photo_file_id: "file-b".to_owned(), - url_file_ids: vec!["file-a".to_owned(), "file-b".to_owned()], + source_file_ids: vec!["file-a".to_owned(), "file-b".to_owned()], } ); } @@ -6708,7 +6684,7 @@ mod tests { let resolver = Arc::new(ImageEditFileResolverStub::successful( ImageEditFileSelection { photo_file_id: "latest-photo-file".to_owned(), - photo_urls: Vec::new(), + photo_file_ids: Vec::new(), }, )); let adapter = TaskmanDialogToolAdapter::new(Arc::clone(&queue)) diff --git a/crates/openplotva-app/src/image_edit_references.rs b/crates/openplotva-app/src/image_edit_references.rs index 39400776..60044d2c 100644 --- a/crates/openplotva-app/src/image_edit_references.rs +++ b/crates/openplotva-app/src/image_edit_references.rs @@ -1,7 +1,7 @@ //! Source pictures of an image edit, prepared for the prompt optimizer the way the //! FLUX.3 prompt upsampler sees its references. -use std::{fmt, future::Future, io::Cursor, pin::Pin}; +use std::{fmt, future::Future, io::Cursor, pin::Pin, sync::Arc}; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD}; use futures_util::StreamExt; @@ -11,7 +11,7 @@ use image::{ }; use openplotva_telegram::TelegramClient; -use crate::image_jobs::ImageEditRequest; +use crate::image_jobs::{ImageEditError, ImageEditFuture, ImageEditRequest, ImageEditor}; const MAX_REFERENCE_BYTES: usize = 20 * 1024 * 1024; const MAX_REFERENCE_SIDE: u32 = 16_384; @@ -95,12 +95,21 @@ impl ImageEditReferenceLoader for TelegramImageEditReferenceLoader { .filter(|value| !value.is_empty()) .collect(); if inputs.is_empty() { - let file_id = request.photo_file_id.trim(); - if file_id.is_empty() { - return Ok(Vec::new()); + let mut file_ids: Vec<&str> = request + .photo_file_ids + .iter() + .map(|value| value.trim()) + .filter(|value| !value.is_empty()) + .collect(); + if file_ids.is_empty() && !request.photo_file_id.trim().is_empty() { + file_ids.push(request.photo_file_id.trim()); } - let path = self.file_path(file_id).await?; - return Ok(vec![self.download(&path).await?]); + let mut images = Vec::with_capacity(file_ids.len()); + for file_id in file_ids { + let path = self.file_path(file_id).await?; + images.push(self.download(&path).await?); + } + return Ok(images); } let mut images = Vec::with_capacity(inputs.len()); for input in inputs { @@ -114,6 +123,61 @@ impl ImageEditReferenceLoader for TelegramImageEditReferenceLoader { } } +/// Hands the editor the request's Telegram pictures as bytes, loaded once by +/// `references`, so their file URLs, which carry the bot token, never leave the bot. +pub struct TelegramPictureBytesEditor { + editor: Editor, + references: Arc, +} + +impl TelegramPictureBytesEditor { + #[must_use] + pub fn new(editor: Editor, references: Arc) -> Self { + Self { editor, references } + } +} + +impl ImageEditor for TelegramPictureBytesEditor +where + Editor: ImageEditor + Sync, +{ + fn expected_image_count(&self) -> usize { + self.editor.expected_image_count() + } + + fn image_targets(&self) -> openplotva_media::ImageTargets { + self.editor.image_targets() + } + + fn edit_image<'a>(&'a self, mut request: ImageEditRequest) -> ImageEditFuture<'a> { + Box::pin(async move { + let names_telegram_pictures = request + .photo_file_ids + .iter() + .any(|value| !value.trim().is_empty()) + || request + .photo_urls + .iter() + .any(|value| telegram_file_path(value.trim()).is_some()); + if names_telegram_pictures { + let images = self + .references + .load_edit_references(&request) + .await + .map_err(|error| { + ImageEditError::Provider(format!("load image edit pictures: {error}")) + })?; + request.photo_urls = images + .iter() + .map(|image| BASE64_STANDARD.encode(image)) + .collect(); + request.photo_file_ids.clear(); + } + self.editor.edit_image(request).await + }) + } +} + fn telegram_file_path(url: &str) -> Option<&str> { let rest = url.strip_prefix(TELEGRAM_FILE_URL_PREFIX)?; let (_, path) = rest.split_once('/')?; @@ -264,12 +328,155 @@ mod tests { telegram_file_path("https://api.telegram.org/file/bot123:abc/photos/file_7.jpg"), Some("photos/file_7.jpg") ); + assert_eq!( + telegram_file_path("https://api.telegram.org/file/bot/photos/file_7.jpg"), + Some("photos/file_7.jpg") + ); assert_eq!( telegram_file_path("https://example.com/file/bot1/a.jpg"), None ); } + #[derive(Debug)] + struct LoaderStub { + result: Result>, String>, + requests: std::sync::Mutex>, + } + + impl LoaderStub { + fn new(result: Result>, String>) -> Arc { + Arc::new(Self { + result, + requests: std::sync::Mutex::default(), + }) + } + + fn calls(&self) -> usize { + self.requests.lock().expect("loader requests").len() + } + } + + impl ImageEditReferenceLoader for LoaderStub { + fn load_edit_references<'a>( + &'a self, + request: &'a ImageEditRequest, + ) -> ImageEditReferencesFuture<'a> { + self.requests + .lock() + .expect("loader requests") + .push(request.clone()); + let result = self.result.clone(); + Box::pin(async move { result }) + } + } + + #[derive(Default)] + struct EditorStub { + requests: std::sync::Mutex>, + } + + impl ImageEditor for EditorStub { + fn edit_image<'a>(&'a self, request: ImageEditRequest) -> ImageEditFuture<'a> { + self.requests.lock().expect("editor requests").push(request); + Box::pin(async { Ok(crate::image_jobs::ImageEditResult::default()) }) + } + } + + async fn edited_request( + loader: Arc, + request: ImageEditRequest, + ) -> Result { + let editor = TelegramPictureBytesEditor::new(EditorStub::default(), loader); + editor.edit_image(request).await?; + let requests = editor.editor.requests.lock().expect("editor requests"); + Ok(requests[0].clone()) + } + + #[tokio::test] + async fn file_ids_reach_the_editor_as_picture_bytes() { + let loader = LoaderStub::new(Ok(vec![b"first".to_vec(), b"second".to_vec()])); + let request = edited_request( + Arc::clone(&loader), + ImageEditRequest { + prompt: "merge".to_owned(), + photo_file_id: "file-a".to_owned(), + photo_file_ids: vec!["file-a".to_owned(), "file-b".to_owned()], + ..ImageEditRequest::default() + }, + ) + .await + .expect("edit"); + + assert_eq!( + request.photo_urls, + vec![ + BASE64_STANDARD.encode(b"first"), + BASE64_STANDARD.encode(b"second") + ] + ); + assert!(request.photo_file_ids.is_empty()); + assert_eq!(loader.calls(), 1); + } + + #[tokio::test] + async fn telegram_file_urls_of_queued_jobs_reach_the_editor_as_picture_bytes() { + let loader = LoaderStub::new(Ok(vec![b"photo".to_vec()])); + let request = edited_request( + loader, + ImageEditRequest { + prompt: "brighter".to_owned(), + photo_urls: vec![ + "https://api.telegram.org/file/bot/photos/file_7.jpg".to_owned(), + ], + ..ImageEditRequest::default() + }, + ) + .await + .expect("edit"); + + assert_eq!(request.photo_urls, vec![BASE64_STANDARD.encode(b"photo")]); + } + + #[tokio::test] + async fn pictures_without_telegram_sources_pass_through() { + let loader = LoaderStub::new(Ok(Vec::new())); + let original = ImageEditRequest { + prompt: "brighter".to_owned(), + photo_file_id: "file-a".to_owned(), + photo_urls: vec!["https://img.test/a.png".to_owned()], + ..ImageEditRequest::default() + }; + let request = edited_request(Arc::clone(&loader), original.clone()) + .await + .expect("edit"); + + assert_eq!(request, original); + assert_eq!(loader.calls(), 0); + } + + #[tokio::test] + async fn unloadable_pictures_fail_the_edit() { + let loader = LoaderStub::new(Err("download telegram file failed".to_owned())); + let error = edited_request( + loader, + ImageEditRequest { + prompt: "brighter".to_owned(), + photo_file_ids: vec!["file-a".to_owned()], + ..ImageEditRequest::default() + }, + ) + .await + .expect_err("pictures are required"); + + assert_eq!( + error, + ImageEditError::Provider( + "load image edit pictures: download telegram file failed".to_owned() + ) + ); + } + #[test] fn inline_pictures_decode_from_base64_or_data_urls() { let bytes = png(8, 8, Rgba([1, 2, 3, 255])); diff --git a/crates/openplotva-app/src/image_jobs.rs b/crates/openplotva-app/src/image_jobs.rs index de6bb258..1b57762a 100644 --- a/crates/openplotva-app/src/image_jobs.rs +++ b/crates/openplotva-app/src/image_jobs.rs @@ -192,6 +192,8 @@ pub struct ImageEditRequest { pub prompt: String, pub photo_file_id: String, pub photo_urls: Vec, + /// Telegram file IDs of the pictures to edit, in order; loaded on the bot side. + pub photo_file_ids: Vec, } /// Successful provider result. @@ -3627,6 +3629,7 @@ where prompt: params.prompt.clone(), photo_file_id: params.photo_file_id, photo_urls: params.photo_urls, + photo_file_ids: params.photo_file_ids, }; let cancelled = std::pin::pin!(cancelled); @@ -5921,6 +5924,7 @@ mod tests { prompt: " make it night\"\"\"}ignored ".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, ) @@ -5947,6 +5951,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), }] ); let calls = effects.calls(); @@ -5980,6 +5985,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, ) @@ -6018,6 +6024,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, ) @@ -7004,6 +7011,7 @@ mod tests { prompt: "make it day".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), ..ImageEditRequest::default() }) .await; @@ -7020,6 +7028,7 @@ mod tests { prompt: "make the scene bright daylight".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), ..ImageEditRequest::default() }] ); @@ -7425,6 +7434,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, now, @@ -7482,6 +7492,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, now, @@ -7555,6 +7566,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, now, @@ -7668,6 +7680,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), }, now, @@ -9098,6 +9111,7 @@ mod tests { user_full_name: "Alice".to_owned(), prompt: "make it night".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: Some(9), ..ImageEditJobParams::default() }, @@ -9908,6 +9922,7 @@ mod tests { prompt: "make it night".to_owned(), photo_file_id: "photo-file".to_owned(), photo_urls: vec!["https://telegram.test/original.png".to_owned()], + photo_file_ids: Vec::new(), thread_id: None, }, now, diff --git a/crates/openplotva-app/src/lib.rs b/crates/openplotva-app/src/lib.rs index 359dfef9..93914057 100644 --- a/crates/openplotva-app/src/lib.rs +++ b/crates/openplotva-app/src/lib.rs @@ -12512,10 +12512,7 @@ async fn start_runtime_workers( .with_image_edit_file_resolver(Arc::new( dialog_tools::TelegramImageEditFileResolver::new( PostgresTelegramFileStore::new(service_clients.postgres.clone()), - dialog_tools::TelegramBotFileUrlProvider::new( - telegram.clone(), - bot_key.to_owned(), - ), + dialog_tools::TelegramBotFileProbe::new(telegram.clone()), ) .with_pending_albums( bot_identity.id, @@ -13367,18 +13364,21 @@ async fn start_runtime_workers( ); let vip_image_edit_queue = Arc::clone(&task_queue_for_updates); - let vip_image_edit_provider = image_jobs::OptimizingImageEditor::new( - image_jobs::RoutedImageEditor::new( - image_attempt_walker.clone(), - vision_data_urls.clone(), - image_jobs::aifarm_draw_api_config_from_app_config(config), + let vip_image_edit_references: Arc = + Arc::new(image_edit_references::TelegramImageEditReferenceLoader::new(telegram.clone())); + let vip_image_edit_provider = image_edit_references::TelegramPictureBytesEditor::new( + image_jobs::OptimizingImageEditor::new( + image_jobs::RoutedImageEditor::new( + image_attempt_walker.clone(), + vision_data_urls.clone(), + image_jobs::aifarm_draw_api_config_from_app_config(config), + ) + .with_workflow_key(VIP_IMAGE_EDITOR_WORKFLOW_KEY), + media_prompt_optimizer.clone(), ) - .with_workflow_key(VIP_IMAGE_EDITOR_WORKFLOW_KEY), - media_prompt_optimizer.clone(), - ) - .with_reference_loader(Arc::new( - image_edit_references::TelegramImageEditReferenceLoader::new(telegram.clone()), - )); + .with_reference_loader(Arc::clone(&vip_image_edit_references)), + vip_image_edit_references, + ); let mut vip_image_edit_effects = image_jobs::TelegramImageJobEffects::new(telegram.clone()) .with_last_generation_writer(Arc::new(service_clients.redis.last_generation_store())); { diff --git a/crates/openplotva-taskman/src/lib.rs b/crates/openplotva-taskman/src/lib.rs index 3aafecbf..68ac8fd3 100644 --- a/crates/openplotva-taskman/src/lib.rs +++ b/crates/openplotva-taskman/src/lib.rs @@ -202,6 +202,10 @@ pub struct ImageJobData { pub image_file_id: String, #[serde(skip_serializing_if = "Vec::is_empty")] pub image_urls: Vec, + /// Telegram file IDs of the pictures an edit works on, in order. They stand in for + /// file URLs, which carry the bot token. + #[serde(skip_serializing_if = "Vec::is_empty")] + pub image_file_ids: Vec, #[serde(skip_serializing_if = "is_false")] pub is_image_edit: bool, #[serde(skip_serializing_if = "Vec::is_empty")] @@ -260,6 +264,7 @@ impl Default for ImageJobData { seed: 0, image_file_id: String::new(), image_urls: Vec::new(), + image_file_ids: Vec::new(), is_image_edit: false, prompt_variants: Vec::new(), is_nsfw: false, @@ -470,6 +475,7 @@ pub struct ImageEditJobParams { pub prompt: String, pub photo_file_id: String, pub photo_urls: Vec, + pub photo_file_ids: Vec, pub thread_id: Option, } @@ -2761,6 +2767,7 @@ pub fn new_image_edit_job_at( prompt: clean_prompt, image_file_id: params.photo_file_id, image_urls: params.photo_urls, + image_file_ids: params.photo_file_ids, is_image_edit: true, ..ImageJobData::default() }), @@ -2952,6 +2959,7 @@ pub fn image_edit_job_params_from_payload( prompt: image_data.prompt.clone(), photo_file_id: image_data.image_file_id.clone(), photo_urls: image_data.image_urls.clone(), + photo_file_ids: image_data.image_file_ids.clone(), thread_id: telegram_data.thread_message_id, }) } @@ -4052,6 +4060,7 @@ mod tests { "https://telegram.test/a.png".to_owned(), "https://telegram.test/b.png".to_owned(), ], + photo_file_ids: Vec::new(), thread_id: Some(40), }, created, @@ -4106,12 +4115,34 @@ mod tests { "https://telegram.test/a.png".to_owned(), "https://telegram.test/b.png".to_owned(), ], + photo_file_ids: Vec::new(), thread_id: Some(40), } ); Ok(()) } + #[test] + fn image_edit_job_keeps_file_ids_without_urls() -> Result<(), Box> { + let params = ImageEditJobParams { + chat_id: -100, + message_id: 20, + user_id: 30, + user_full_name: "Alice".to_owned(), + prompt: "make it night".to_owned(), + photo_file_id: "file-a".to_owned(), + photo_file_ids: vec!["file-a".to_owned(), "file-b".to_owned()], + ..ImageEditJobParams::default() + }; + let job = new_image_edit_job_at(params.clone(), OffsetDateTime::UNIX_EPOCH); + + let image_data = serde_json::to_value(&job.data)?["image_data"].clone(); + assert_eq!(image_data["image_file_ids"], json!(["file-a", "file-b"])); + assert!(image_data.get("image_urls").is_none()); + assert_eq!(image_edit_job_params_from_stateless_job(&job)?, params); + Ok(()) + } + #[test] fn music_gen_job_payload_matches_go_shape_for_vip_assignment() -> Result<(), Box> { diff --git a/migrations/211_redact_telegram_file_urls_in_taskman_jobs.down.sql b/migrations/211_redact_telegram_file_urls_in_taskman_jobs.down.sql new file mode 100644 index 00000000..2cd07249 --- /dev/null +++ b/migrations/211_redact_telegram_file_urls_in_taskman_jobs.down.sql @@ -0,0 +1,3 @@ +-- No-op: the up migration removes the bot token from stored Telegram file URLs, +-- and the token must not come back. One-way for data. +SELECT 1; diff --git a/migrations/211_redact_telegram_file_urls_in_taskman_jobs.up.sql b/migrations/211_redact_telegram_file_urls_in_taskman_jobs.up.sql new file mode 100644 index 00000000..fad50a6d --- /dev/null +++ b/migrations/211_redact_telegram_file_urls_in_taskman_jobs.up.sql @@ -0,0 +1,11 @@ +-- Image edit jobs stored Telegram file URLs, which carry the bot token. New jobs +-- store file IDs instead; this removes the token from the URLs already stored. +-- The `bot/` form still resolves for jobs that are queued. +UPDATE taskman_jobs +SET record = regexp_replace( + record::text, + '(api\.telegram\.org/file/bot)[0-9]+:[A-Za-z0-9_-]+', + '\1', + 'g' + )::jsonb +WHERE record::text ~ 'api\.telegram\.org/file/bot[0-9]+:';