use std::collections::HashMap;
use base64::Engine as _;
use uuid::Uuid;
use crate::app::events::AppEvent;
use crate::entities::attachment::{Resolved, handle_number, handle_range, name_is_shared};
use crate::entities::message_image::{MessageImage, infos, resolve_target};
use crate::features::image_command::ImageProgress;
use crate::features::image_fetch::{self, FetchError, FetchedImage};
use crate::features::image_prepare::{PrepareError, PreparedImage, prepare, prepare_rgba};
use crate::shared::api::VisionSupport;
use crate::shared::config::{ImageSettings, ServerMode};
use crate::shared::i18n::Locale;
use super::Orchestrator;
pub(super) struct ImageAttachResult {
pub(super) chat_id: Uuid,
pub(super) vision: VisionSupport,
pub(super) outcome: Result<MessageImage, String>,
}
enum Staging {
Local(ImageSource),
Url(String),
}
enum ImageSource {
File(String),
Clipboard {
image: Box<crate::app::events::ClipboardImage>,
name: String,
},
Downloaded {
fetched: Box<FetchedImage>,
url: String,
},
}
impl Staging {
async fn resolve(
self,
cfg: &ImageSettings,
loc: &'static Locale,
) -> Result<ImageSource, String> {
match self {
Staging::Local(source) => Ok(source),
Staging::Url(url) => {
let fetched = image_fetch::fetch(&url, cfg.max_bytes)
.await
.map_err(|e| localize_fetch(&e, cfg.max_bytes, loc))?;
Ok(ImageSource::Downloaded {
fetched: Box::new(fetched),
url,
})
}
}
}
}
impl ImageSource {
fn prepare(self, cfg: &ImageSettings, loc: &'static Locale) -> Result<MessageImage, String> {
match self {
ImageSource::File(path) => prepare_image(std::path::Path::new(&path), cfg, loc),
ImageSource::Clipboard { image, name } => prepare_clipboard(*image, name, cfg, loc),
ImageSource::Downloaded { fetched, url } => prepare_downloaded(*fetched, url, cfg, loc),
}
}
}
impl Orchestrator {
pub(super) fn handle_image_attach(&mut self, path: String) {
let path = path.trim().to_string();
if path.is_empty() {
return;
}
if image_fetch::looks_like_url(&path) {
self.stage_image(Staging::Url(path));
} else {
self.stage_image(Staging::Local(ImageSource::File(path)));
}
}
pub(super) fn handle_image_paste(&mut self, image: crate::app::events::ClipboardImage) {
let Some(chat_id) = self.active_id else {
self.fail_image(self.ui_locale().t("ui.err.image_no_active_chat"));
return;
};
let name = self.free_clipboard_name(chat_id);
self.stage_image(Staging::Local(ImageSource::Clipboard {
image: Box::new(image),
name,
}));
}
fn stage_image(&mut self, source: Staging) {
let Some(chat_id) = self.active_id else {
self.fail_image(self.ui_locale().t("ui.err.image_no_active_chat"));
return;
};
let cfg = self.config.images;
let staged = self.staged_images.get(&chat_id).map_or(0, Vec::len);
if staged >= cfg.max_count {
let msg = self.ui_locale().tf(
"ui.err.image_too_many",
&[("max", &cfg.max_count.to_string())],
);
self.fail_image(&msg);
return;
}
let loc = self.ui_locale();
let backend = self.engines.backend_if_ready(loc).ok();
let managed = self.config.engine.mode == ServerMode::Managed;
let tx = self.image_tx.clone();
tokio::spawn(async move {
let vision = match &backend {
Some(b) => b.vision().await,
None => VisionSupport::Unknown,
};
if vision == VisionSupport::Unsupported {
let key = if managed {
"ui.err.image_no_vision_managed"
} else {
"ui.err.image_no_vision"
};
let _ = tx.send(ImageAttachResult {
chat_id,
vision,
outcome: Err(loc.t(key).to_string()),
});
return;
}
let outcome = match source.resolve(&cfg, loc).await {
Ok(source) => tokio::task::spawn_blocking(move || source.prepare(&cfg, loc))
.await
.unwrap_or_else(|e| {
Err(loc.tf("ui.err.image_failed", &[("err", &e.to_string())]))
}),
Err(err) => Err(err),
};
let _ = tx.send(ImageAttachResult {
chat_id,
vision,
outcome,
});
});
}
fn free_clipboard_name(&self, chat_id: Uuid) -> String {
let staged = self.staged_images.get(&chat_id);
let taken = |name: &str| {
staged.is_some_and(|v| v.iter().any(|i| i.name.eq_ignore_ascii_case(name)))
};
if !taken(CLIPBOARD_NAME) {
return CLIPBOARD_NAME.to_string();
}
(2..)
.map(|n| format!("clipboard-{n}.png"))
.find(|name| !taken(name))
.unwrap_or_else(|| CLIPBOARD_NAME.to_string())
}
pub(super) fn handle_image_result(&mut self, res: ImageAttachResult) {
let image = match res.outcome {
Ok(i) => i,
Err(err) => {
self.fail_image(&err);
return;
}
};
if !self.chats.iter().any(|c| c.id == res.chat_id) {
return;
}
let staged = self.staged_images.entry(res.chat_id).or_default();
staged.retain(|i| i.source != image.source);
let info = (&image).into();
staged.push(image);
let count = staged.len();
self.emit_image_progress(ImageProgress::Attached {
info,
staged: count,
});
if res.vision == VisionSupport::Unknown && count == 1 {
self.emit_image_progress(ImageProgress::VisionUnknown);
}
self.emit_staged_images();
}
pub(super) fn note_withheld_images(&mut self, chat_id: Uuid, count: usize) {
if count == 0 || !self.images_withheld_noted.insert(chat_id) {
return;
}
let note = self
.ui_locale()
.tf("ui.chat.images_withheld", &[("n", &count.to_string())]);
let _ = self.evt_tx.send(AppEvent::Notice(note));
}
pub(super) fn handle_image_remove(&mut self, target: String) {
let Some(chat_id) = self.active_id else {
self.fail_image(self.ui_locale().t("ui.err.image_no_active_chat"));
return;
};
let loc = self.ui_locale();
let staged = self.staged_images.entry(chat_id).or_default();
let idx = match resolve_target(staged, &target) {
Resolved::One(idx) => idx,
Resolved::Shared(hits) => {
let candidates =
Self::candidate_lines(hits.iter().map(|&i| (i, staged[i].source.as_str())));
let msg = loc.tf(
"ui.err.image_name_shared",
&[("target", target.trim()), ("candidates", &candidates)],
);
self.fail_image(&msg);
return;
}
Resolved::Nothing => {
let msg = match handle_number(&target).filter(|_| !staged.is_empty()) {
Some(n) => loc.tf(
"ui.err.image_no_such_number",
&[
("n", &n.to_string()),
("range", &handle_range(staged.len())),
],
),
None => loc.tf("ui.err.image_not_staged", &[("target", target.trim())]),
};
self.fail_image(&msg);
return;
}
};
let shared = name_is_shared(staged, idx, |i| i.name.as_str());
let removed = staged.remove(idx);
self.emit_image_progress(ImageProgress::Removed {
name: removed.name,
source: shared.then_some(removed.source),
});
self.emit_staged_images();
}
pub(super) fn handle_image_list(&mut self) {
let Some(chat_id) = self.active_id else {
self.fail_image(self.ui_locale().t("ui.err.image_no_active_chat"));
return;
};
let items = self
.staged_images
.get(&chat_id)
.map(|v| infos(v))
.unwrap_or_default();
self.emit_image_progress(ImageProgress::Listed { items });
}
pub(super) fn take_staged_images(&mut self, chat_id: Uuid) -> Vec<MessageImage> {
let taken = self.staged_images.remove(&chat_id).unwrap_or_default();
if !taken.is_empty() {
self.emit_staged_images();
}
taken
}
pub(super) fn forget_staged_images(&mut self, chat_id: Uuid) {
if self.staged_images.remove(&chat_id).is_some() {
self.emit_staged_images();
}
}
pub(super) fn emit_staged_images(&self) {
let items = self
.active_id
.and_then(|id| self.staged_images.get(&id))
.map(|v| infos(v))
.unwrap_or_default();
let _ = self.evt_tx.send(AppEvent::StagedImages(items));
}
fn emit_image_progress(&self, progress: ImageProgress) {
let _ = self.evt_tx.send(AppEvent::ImageProgress(progress));
}
fn fail_image(&self, msg: &str) {
self.emit_image_progress(ImageProgress::Failed(msg.to_string()));
}
}
pub(super) type StagedImages = HashMap<Uuid, Vec<MessageImage>>;
const CLIPBOARD_NAME: &str = "clipboard.png";
fn prepare_clipboard(
image: crate::app::events::ClipboardImage,
name: String,
cfg: &ImageSettings,
loc: &'static Locale,
) -> Result<MessageImage, String> {
let prepared = prepare_rgba(image.width, image.height, &image.rgba, cfg.downscale_px).map_err(
|e| match e {
PrepareError::Undecodable => loc.t("ui.err.image_clipboard_unusable").to_string(),
PrepareError::Failed(err) => loc.tf("ui.err.image_failed", &[("err", &err)]),
},
)?;
if prepared.bytes.len() as u64 > cfg.max_bytes {
return Err(loc.tf(
"ui.err.image_too_big",
&[
(
"size",
&crate::entities::attachment::format_bytes(prepared.bytes.len()),
),
(
"max",
&crate::entities::attachment::format_bytes(cfg.max_bytes as usize),
),
],
));
}
Ok(staged(
prepared,
name,
format!("clipboard:{}", Uuid::new_v4()),
))
}
fn prepare_downloaded(
fetched: FetchedImage,
url: String,
cfg: &ImageSettings,
loc: &'static Locale,
) -> Result<MessageImage, String> {
let prepared = prepare(&fetched.bytes, cfg.downscale_px).map_err(|e| match e {
PrepareError::Undecodable => not_an_image(&fetched.content_type, loc),
PrepareError::Failed(err) => loc.tf("ui.err.image_failed", &[("err", &err)]),
})?;
let name = image_fetch::display_name(&fetched.final_url, extension(prepared.mime));
Ok(staged(prepared, name, url))
}
fn not_an_image(content_type: &str, loc: &'static Locale) -> String {
if content_type.is_empty() || content_type.starts_with("image/") {
loc.t("ui.err.image_undecodable").to_string()
} else {
loc.tf("ui.err.image_url_not_image", &[("type", content_type)])
}
}
fn extension(mime: &str) -> &'static str {
if mime == "image/jpeg" { "jpg" } else { "png" }
}
fn localize_fetch(err: &FetchError, max_bytes: u64, loc: &'static Locale) -> String {
match err {
FetchError::Scheme => loc.t("ui.err.image_url_scheme").to_string(),
FetchError::Malformed => loc.t("ui.err.image_url_malformed").to_string(),
FetchError::TooManyRedirects => loc.tf(
"ui.err.image_url_redirects",
&[("max", &image_fetch::MAX_REDIRECTS.to_string())],
),
FetchError::Request(e) => loc.tf("ui.err.image_url_request", &[("err", e)]),
FetchError::Status(code) => {
loc.tf("ui.err.image_url_status", &[("status", &code.to_string())])
}
FetchError::TooBig => loc.tf(
"ui.err.image_url_too_big",
&[(
"max",
&crate::entities::attachment::format_bytes(max_bytes as usize),
)],
),
FetchError::Empty => loc.t("ui.err.image_url_empty").to_string(),
}
}
fn staged(prepared: PreparedImage, name: String, source: String) -> MessageImage {
let data = base64::engine::general_purpose::STANDARD.encode(&prepared.bytes);
MessageImage::new(
name,
source,
prepared.mime,
prepared.width,
prepared.height,
data,
)
}
fn prepare_image(
path: &std::path::Path,
cfg: &ImageSettings,
loc: &'static Locale,
) -> Result<MessageImage, String> {
let meta = std::fs::metadata(path)
.map_err(|e| loc.tf("ui.err.image_unavailable", &[("err", &e.to_string())]))?;
if !meta.is_file() {
return Err(loc.t("ui.err.image_not_a_file").to_string());
}
if meta.len() > cfg.max_bytes {
return Err(loc.tf(
"ui.err.image_too_big",
&[
(
"size",
&crate::entities::attachment::format_bytes(meta.len() as usize),
),
(
"max",
&crate::entities::attachment::format_bytes(cfg.max_bytes as usize),
),
],
));
}
let bytes = std::fs::read(path)
.map_err(|e| loc.tf("ui.err.image_unavailable", &[("err", &e.to_string())]))?;
let prepared = prepare(&bytes, cfg.downscale_px).map_err(|e| match e {
PrepareError::Undecodable => loc.t("ui.err.image_undecodable").to_string(),
PrepareError::Failed(err) => loc.tf("ui.err.image_failed", &[("err", &err)]),
})?;
let name = path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| path.display().to_string());
Ok(staged(
prepared,
name,
crate::features::rag_ingest::canonical_source(path),
))
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
fn png(width: u32, height: u32) -> Vec<u8> {
let buf = image::ImageBuffer::from_fn(width, height, |x, _| {
image::Rgb([(x % 256) as u8, 40, 90])
});
let mut out = Vec::new();
image::DynamicImage::ImageRgb8(buf)
.write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
.unwrap();
out
}
fn en() -> &'static Locale {
crate::shared::i18n::locale(crate::shared::i18n::Lang::En)
}
#[test]
fn prepares_an_image_and_reports_its_metadata() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("shot.png");
std::fs::write(&path, png(64, 32)).unwrap();
let image = prepare_image(&path, &ImageSettings::default(), en()).unwrap();
assert_eq!(image.name, "shot.png");
assert_eq!(image.mime, "image/png");
assert_eq!((image.width, image.height), (64, 32));
assert!(image.bytes > 0);
let decoded = base64::engine::general_purpose::STANDARD
.decode(&image.data)
.unwrap();
assert_eq!(decoded.len(), image.bytes);
assert_eq!(&decoded[1..4], b"PNG");
}
#[test]
fn an_oversized_file_is_refused_before_it_is_decoded() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("big.png");
std::fs::write(&path, png(200, 200)).unwrap();
let cfg = ImageSettings {
max_bytes: 16,
..ImageSettings::default()
};
let err = prepare_image(&path, &cfg, en()).unwrap_err();
assert!(err.contains("too large"), "{err}");
}
#[test]
fn a_non_image_and_a_missing_file_are_refused_with_their_own_messages() {
let dir = tempfile::tempdir().unwrap();
let text = dir.path().join("not-an-image.txt");
std::fs::write(&text, b"just words").unwrap();
let err = prepare_image(&text, &ImageSettings::default(), en()).unwrap_err();
assert!(
err.contains("png"),
"the refusal must name what works: {err}"
);
let missing = dir.path().join("nope.png");
let err = prepare_image(&missing, &ImageSettings::default(), en()).unwrap_err();
assert!(err.contains("unavailable"), "{err}");
let err = prepare_image(dir.path(), &ImageSettings::default(), en()).unwrap_err();
assert!(err.contains("not a file"), "{err}");
}
#[test]
fn a_large_image_is_downscaled_before_it_is_stored() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("huge.png");
std::fs::write(&path, png(2400, 1200)).unwrap();
let image = prepare_image(&path, &ImageSettings::default(), en()).unwrap();
assert_eq!((image.width, image.height), (1568, 784));
assert!(image.bytes < std::fs::metadata(&path).unwrap().len() as usize);
}
#[test]
fn every_download_failure_names_what_happened() {
let max = ImageSettings::default().max_bytes;
let msg = |e: FetchError| localize_fetch(&e, max, en());
assert!(msg(FetchError::Status(404)).contains("404"));
assert!(msg(FetchError::Request("dns error".into())).contains("dns error"));
assert!(
msg(FetchError::TooManyRedirects).contains(&image_fetch::MAX_REDIRECTS.to_string())
);
assert!(msg(FetchError::TooBig).contains("10"), "the limit in MB");
assert!(msg(FetchError::Scheme).contains("/image attach"));
assert!(msg(FetchError::Malformed).contains("Copy image address"));
}
#[test]
fn an_undecodable_download_is_explained_by_what_the_server_said_it_sent() {
let msg = not_an_image("text/html", en());
assert!(msg.contains("text/html"), "{msg}");
for ct in ["image/heic", ""] {
assert_eq!(not_an_image(ct, en()), en().t("ui.err.image_undecodable"));
}
}
#[test]
fn download_errors_are_localized_for_all_langs() {
for &lang in crate::shared::i18n::Lang::ALL {
let loc = crate::shared::i18n::locale(lang);
let mut msgs: Vec<String> = [
FetchError::Scheme,
FetchError::Malformed,
FetchError::TooManyRedirects,
FetchError::Request("connection refused".into()),
FetchError::Status(404),
FetchError::TooBig,
FetchError::Empty,
]
.into_iter()
.map(|e| localize_fetch(&e, ImageSettings::default().max_bytes, loc))
.collect();
msgs.push(not_an_image("text/html", loc));
for msg in msgs {
assert!(
!msg.contains('{') && !msg.contains('}'),
"unsubstituted placeholder in {lang:?}: {msg}"
);
if lang == crate::shared::i18n::Lang::En {
assert!(
!msg.chars().any(|c| ('\u{0400}'..='\u{04FF}').contains(&c)),
"Cyrillic leaked into the en message: {msg}"
);
}
}
}
}
#[test]
fn preparation_errors_are_localized_for_all_langs() {
let dir = tempfile::tempdir().unwrap();
let text = dir.path().join("not-an-image.txt");
std::fs::write(&text, b"just words").unwrap();
let missing = dir.path().join("nope.png");
let tiny = ImageSettings {
max_bytes: 4,
..ImageSettings::default()
};
for &lang in crate::shared::i18n::Lang::ALL {
let loc = crate::shared::i18n::locale(lang);
let mut msgs = vec![
prepare_image(&text, &ImageSettings::default(), loc).unwrap_err(),
prepare_image(&missing, &ImageSettings::default(), loc).unwrap_err(),
prepare_image(dir.path(), &ImageSettings::default(), loc).unwrap_err(),
prepare_image(&text, &tiny, loc).unwrap_err(),
];
msgs.push(loc.t("ui.err.image_no_vision_managed").to_string());
msgs.push(loc.t("ui.err.image_no_vision").to_string());
msgs.push(loc.tf("ui.err.image_too_many", &[("max", "8")]));
for msg in msgs {
assert!(
!msg.contains('{') && !msg.contains('}'),
"unsubstituted placeholder in {lang:?}: {msg}"
);
if lang == crate::shared::i18n::Lang::En {
assert!(
!msg.chars().any(|c| ('\u{0400}'..='\u{04FF}').contains(&c)),
"Cyrillic leaked into the en message: {msg}"
);
}
}
}
}
}