From 6cb0428d9a819e187a223b2b55d4f8cd99a1205d Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 22:32:36 +0100 Subject: [PATCH 1/6] feat(imessage): support inbound image attachments --- ARCHITECTURE.md | 6 +- Cargo.lock | 1 + Cargo.toml | 1 + docs/channels/imessage.md | 35 ++- docs/getting-started.md | 5 +- docs/security.md | 16 +- docs/services.md | 9 +- src/channel.rs | 25 +- src/doctor.rs | 11 +- src/gateway/tests.rs | 199 +++++++++++++++ src/imessage/attachments.rs | 480 ++++++++++++++++++++++++++++++++++++ src/imessage/mod.rs | 2 + src/imessage/poller.rs | 158 +++++++++++- 13 files changed, 923 insertions(+), 25 deletions(-) create mode 100644 src/imessage/attachments.rs diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 7af73be..b13ffe9 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -425,7 +425,11 @@ Files: [`src/imessage/`](src/imessage/) - Uses local row IDs as the monotonic cursor. - Sends through `osascript`. - Keeps legacy unprefixed route and session aliases for migration. -- Supports text only. +- Joins ordered attachment metadata during polling without opening files. +- Opens images only after acceptance, confines canonical paths to the Messages + attachment directory, and converts HEIC or HEIF locally with `sips`. +- Applies provider-neutral image limits and temporary-file cleanup before + passing images to Claude Code, Codex, or Pi. ### Telegram diff --git a/Cargo.lock b/Cargo.lock index ed3545b..54b9629 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -805,6 +805,7 @@ dependencies = [ "chrono-tz", "futures-util", "humantime", + "libc", "pulldown-cmark", "reqwest", "rusqlite", diff --git a/Cargo.toml b/Cargo.toml index ed5850f..4f6adf7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ chrono = "0.4" tracing = "0.1" tracing-subscriber = "0.3" pulldown-cmark = { version = "0.13", default-features = false } +libc = "0.2" [profile.release] strip = true diff --git a/docs/channels/imessage.md b/docs/channels/imessage.md index b0bf894..addf539 100644 --- a/docs/channels/imessage.md +++ b/docs/channels/imessage.md @@ -8,8 +8,9 @@ iMessage API or expose a network service. - macOS with Messages signed in - Full Disk Access for the terminal or service process running Push -- access to `~/Library/Messages/chat.db` +- access to `~/Library/Messages/chat.db` and `~/Library/Messages/Attachments` - `osascript` on `PATH` +- `/usr/bin/sips` for HEIC and HEIF image conversion Run `push doctor` from the same user and environment as the long-running service. A successful interactive check does not prove that a separate service @@ -46,11 +47,37 @@ sensitivity. Treat every allowed handle as an operator of the configured backend. A sender can ask the agent to use any capability allowed by that agent's configuration. +## Image messages + +Send up to four JPEG, PNG, WebP, HEIC, or HEIF images in one accepted +conversation, with or without message text. Their combined prepared size must +be at most 6 MiB. Images work with Claude Code, Codex, and Pi. + +Polling reads attachment paths, byte-size hints, and MIME type hints from +`chat.db`; it does not open the attachment files. After the direct-message, +sender, reply-marker, and message checks pass, the worker canonicalizes each +path and requires it to remain under `~/Library/Messages/Attachments` (or the +`Attachments` directory beside a custom `imessage.db_path`). Missing files, +directories, escaping symlinks, and unsupported documents are rejected with a +safe retry reply before an agent starts. + +JPEG, PNG, and WebP files go through the shared byte limit and signature +validation directly. Push converts HEIC and HEIF locally with macOS `sips` in +an owner-only temporary directory, validates the resulting JPEG against the +same shared limit, and removes the conversion file immediately. The prepared +agent handoff files are also owner-only and are removed after the turn. +Conversation history retains only the message text or an image placeholder, +not image bytes or local attachment paths. Review the configured backend's +image and data controls before using this feature. + +Sending generated images back through iMessage is not supported. + ## What Push ignores - group chats - tapbacks and Messages system rows -- blank messages +- blank messages without an image attachment +- stickers, videos, and Live Photo video components - messages from handles outside the allowlist - Push's own replies containing the built-in Push reply marker @@ -98,6 +125,10 @@ process, and rerun: push doctor ``` +Image messages also require that the same process can read +`~/Library/Messages/Attachments`. Recheck Full Disk Access if text works but +images fail. + ### Messages are ignored Confirm the conversation is one-to-one and that its sender or chat identifier diff --git a/docs/getting-started.md b/docs/getting-started.md index c155cad..b5c9841 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -116,8 +116,9 @@ practical structure for identity, context, shared skills, jobs, and evals. === "iMessage" - Give the terminal or service host Full Disk Access in macOS System - Settings, then edit `$PUSH_HOME/config.toml`: + Give the terminal or service host Full Disk Access to the Messages database + and attachment files in macOS System Settings, then edit + `$PUSH_HOME/config.toml`: ```toml channel = "imessage" diff --git a/docs/security.md b/docs/security.md index c5f1d22..e8a7d6e 100644 --- a/docs/security.md +++ b/docs/security.md @@ -83,11 +83,17 @@ variable or move the config outside. When `voice.openai_api_key` is configured, chmod 600 "${PUSH_HOME:-$HOME/.push}/config.toml" ``` -Accepted Telegram and Slack images are briefly written under `$PUSH_HOME/cache` -with owner-only permissions, passed to the selected agent backend, and -removed when the turn ends. Conversation history retains only the caption or an -image placeholder. Protect the cache directory while Push is running and -review the configured model provider's image data controls. +Accepted iMessage, Telegram, and Slack images are briefly written under +`$PUSH_HOME/cache` with owner-only permissions, passed to the selected agent +backend, and removed when the turn ends. Conversation history retains only the +message text or an image placeholder. Protect the cache directory while Push +is running and review the configured model provider's image data controls. + +iMessage polling stores only attachment paths and safe metadata in memory. It +opens files only after conversation and sender acceptance, and only when their +canonical paths remain under the Messages attachment directory. HEIC and HEIF +conversion uses macOS `sips` in a separate owner-only temporary directory that +is removed immediately after conversion. Slack's recovery inbox persists file IDs and safe size and MIME type hints, but not private download URLs or file bytes. Push resolves and downloads a Slack diff --git a/docs/services.md b/docs/services.md index 5bd6a20..63b1c9c 100644 --- a/docs/services.md +++ b/docs/services.md @@ -3,9 +3,10 @@ This guide covers running `push` continuously under a process manager. The iMessage channel is macOS-only because it reads -`~/Library/Messages/chat.db` and sends replies with `osascript`. Telegram uses -outbound HTTPS long polling. Slack uses outbound Socket Mode. Both can run -under `systemd` on Linux or a VM. +`~/Library/Messages/chat.db`, opens accepted files under +`~/Library/Messages/Attachments`, and sends replies with `osascript`. Telegram +uses outbound HTTPS long polling. Slack uses outbound Socket Mode. Both can +run under `systemd` on Linux or a VM. ## Before Installing a Service @@ -31,7 +32,7 @@ Set one absolute `PUSH_HOME` in the service definition. It defaults to - agent write access to `assistant_root/jobs/` when jobs should be created from chat - access to the selected `claude`, `codex`, or `pi` executable on `PATH` - backend login, tokens, settings, MCP config, and project credentials -- for iMessage on macOS, Full Disk Access and `osascript` +- for iMessage on macOS, Full Disk Access, `osascript`, and `sips` - for Telegram, a token in the private config and network access to `api.telegram.org` - for Slack, app and bot tokens in the private config or service environment, diff --git a/src/channel.rs b/src/channel.rs index 1474a49..a112e35 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -379,7 +379,23 @@ impl ChannelContract for IMessageChannel { is_group: message.is_group, text: message.text, voice: None, - images: Vec::new(), + images: message + .attachments + .into_iter() + .map(|attachment| InboundImage { + locator: attachment.locator.clone(), + file_size: if crate::imessage::needs_conversion( + &attachment.locator, + attachment.mime_type.as_deref(), + ) { + None + } else { + attachment.file_size + }, + mime_type: attachment.mime_type, + data: None, + }) + .collect(), is_from_me: message.is_from_me, is_supported: true, thread_id: None, @@ -476,12 +492,7 @@ impl ChannelContract for IMessageChannel { } async fn download_image(&self, image: &InboundImage) -> Result { - let Some(bytes) = &image.data else { - bail!("iMessage image attachments are not supported yet"); - }; - Ok(DownloadedImage { - bytes: bytes.clone(), - }) + self.poller.download_image(image).await } } diff --git a/src/doctor.rs b/src/doctor.rs index cc92381..54387fb 100644 --- a/src/doctor.rs +++ b/src/doctor.rs @@ -387,6 +387,7 @@ fn check_bins_with( .is_ok_and(|channels| channels.contains(&config::ChannelKind::IMessage)) { bins.push("osascript"); + bins.push("/usr/bin/sips"); } bins.sort_unstable(); bins.dedup(); @@ -641,6 +642,9 @@ claude_tools = [] assert!(checks.iter().any(|check| { check.name == "binary osascript" && matches!(check.status, CheckStatus::Fail) })); + assert!(checks.iter().any(|check| { + check.name == "binary /usr/bin/sips" && matches!(check.status, CheckStatus::Fail) + })); } #[test] @@ -651,7 +655,8 @@ claude_tools = [] let mut checks = Vec::new(); check_bins_with(&cfg, &mut checks, |bin| { - (bin == "/custom/pi" || bin == "osascript").then(|| PathBuf::from(bin)) + (bin == "/custom/pi" || bin == "osascript" || bin == "/usr/bin/sips") + .then(|| PathBuf::from(bin)) }); assert!(checks.iter().any(|check| { @@ -680,7 +685,8 @@ claude_tools = [] let mut checks = Vec::new(); check_bins_with(&cfg, &mut checks, |bin| { - (bin == "codex" || bin == "osascript").then(|| PathBuf::from(bin)) + (bin == "codex" || bin == "osascript" || bin == "/usr/bin/sips") + .then(|| PathBuf::from(bin)) }); assert!(checks.iter().any(|check| { @@ -713,6 +719,7 @@ claude_tools = [] .iter() .any(|check| check.name == "binary /fake/claude")); assert!(!checks.iter().any(|check| check.name.contains("osascript"))); + assert!(!checks.iter().any(|check| check.name.contains("sips"))); } #[test] diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 598ccfa..fbb4e80 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -1869,6 +1869,205 @@ async fn slack_images_reach_every_agent_backend_and_are_removed_after_each_turn( } } +#[tokio::test(flavor = "current_thread")] +async fn imessage_images_reach_every_agent_backend_and_are_removed_after_each_turn() { + for (backend, name) in [ + (AgentBackend::Claude, "claude"), + (AgentBackend::Codex, "codex"), + (AgentBackend::Pi, "pi"), + ] { + let state_path = temp_state_path(); + let sessions_dir = temp_path(&format!("imessage-{name}-image-sessions")); + let assistant_dir = temp_path(&format!("imessage-{name}-image-assistant")); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut cfg = test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + ); + cfg.agent = name.to_string(); + let mut gateway = Gateway::new(cfg).unwrap(); + gateway.ctx.runners = Arc::new(HashMap::from([( + backend, + Runner::Fake(FakeRunner { + backend, + session_id: "fake-session".to_string(), + calls: calls.clone(), + before_return: None, + wait_for_release: None, + failure: None, + resume_missing_once: None, + }), + )])); + let mut inbound = message(1, "+15551234567", "+15551234567", false, ""); + inbound.images.push(InboundImage { + locator: "inline.png".to_string(), + file_size: Some(12), + mime_type: Some("image/png".to_string()), + data: Some(valid_png()), + }); + + run_messages(&mut gateway, vec![inbound]).await; + + let calls = calls.lock().unwrap(); + assert_eq!(calls.len(), 1, "{name}"); + assert_eq!( + crate::prompt::current_message(&calls[0].prompt).as_deref(), + Some("[Image attachment]"), + "{name}" + ); + assert_eq!(calls[0].images.len(), 1, "{name}"); + assert!(!calls[0].images[0].exists(), "{name}"); + drop(calls); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(format!("{state_path}.cache")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); + } +} + +#[tokio::test(flavor = "current_thread")] +async fn imessage_worker_reads_ordered_local_images_and_rejects_bad_attachments() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-local-image-sessions"); + let assistant_dir = temp_path("imessage-local-image-assistant"); + let messages_dir = temp_path("imessage-local-messages"); + let attachments = messages_dir.join("Attachments"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + std::fs::create_dir_all(&attachments).unwrap(); + let png = attachments.join("first.png"); + let jpeg = attachments.join("second.jpg"); + let webp = attachments.join("third.webp"); + let pdf = attachments.join("document.pdf"); + std::fs::write(&png, valid_png()).unwrap(); + std::fs::write(&jpeg, b"\xff\xd8\xffbody").unwrap(); + std::fs::write(&webp, b"RIFF\x04\x00\x00\x00WEBPbody").unwrap(); + std::fs::write(&pdf, b"%PDF").unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut cfg = test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + ); + cfg.db_path = messages_dir.join("chat.db").to_string_lossy().to_string(); + let mut gateway = Gateway::new(cfg).unwrap(); + gateway.ctx.runners = Arc::new(fake_runners(calls.clone())); + let mut valid = message(1, "+15551234567", "+15551234567", false, "compare these"); + valid.images = vec![ + InboundImage { + locator: "~/Library/Messages/Attachments/first.png".to_string(), + file_size: Some(valid_png().len()), + mime_type: Some("image/png".to_string()), + data: None, + }, + InboundImage { + locator: jpeg.to_string_lossy().to_string(), + file_size: Some(7), + mime_type: Some("image/jpeg".to_string()), + data: None, + }, + InboundImage { + locator: "third.webp".to_string(), + file_size: Some(16), + mime_type: Some("image/webp".to_string()), + data: None, + }, + ]; + let mut missing = message(2, "+15551234567", "+15551234567", false, "missing"); + missing.images.push(InboundImage { + locator: "missing.png".to_string(), + file_size: None, + mime_type: Some("image/png".to_string()), + data: None, + }); + let mut unsupported = message(3, "+15551234567", "+15551234567", false, "unsupported"); + unsupported.images.push(InboundImage { + locator: pdf.to_string_lossy().to_string(), + file_size: Some(4), + mime_type: Some("application/pdf".to_string()), + data: None, + }); + + run_messages(&mut gateway, vec![valid, missing, unsupported]).await; + + let calls = calls.lock().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].images.len(), 3); + assert_eq!(calls[0].images[0].extension().unwrap(), "png"); + assert_eq!(calls[0].images[1].extension().unwrap(), "jpg"); + assert_eq!(calls[0].images[2].extension().unwrap(), "webp"); + assert!(calls[0].images.iter().all(|path| !path.exists())); + drop(calls); + let replies = gateway.ctx.sent_replies.lock().unwrap(); + assert_eq!(replies.len(), 3); + assert!(replies[1].1.contains("JPEG, PNG, or WebP")); + assert!(replies[2].1.contains("JPEG, PNG, or WebP")); + drop(replies); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 3); + let audit = std::fs::read_to_string(format!("{state_path}.audit.jsonl")).unwrap(); + assert!(!audit.contains(messages_dir.to_string_lossy().as_ref())); + assert!(!audit.contains("missing.png")); + assert!(!audit.contains("document.pdf")); + let history = gateway + .ctx + .history + .lock() + .unwrap() + .recent_messages_before("imessage", "imessage:dm:+15551234567", i64::MAX, 10) + .unwrap(); + assert!(history.iter().all(|message| { + !message.content.contains("Attachments") + && !message.content.contains("missing.png") + && !message.content.contains("document.pdf") + })); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(format!("{state_path}.cache")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); + let _ = std::fs::remove_dir_all(messages_dir); +} + +#[tokio::test(flavor = "current_thread")] +async fn rejected_imessage_sender_does_not_open_local_attachments() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-rejected-image-sessions"); + let assistant_dir = temp_path("imessage-rejected-image-assistant"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut gateway = Gateway::new(test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + )) + .unwrap(); + gateway.ctx.runners = Arc::new(fake_runners(calls.clone())); + let mut inbound = message(1, "+15550000000", "+15550000000", false, "inspect"); + inbound.images.push(InboundImage { + locator: "/private/outside/missing.png".to_string(), + file_size: None, + mime_type: Some("image/png".to_string()), + data: None, + }); + + run_messages(&mut gateway, vec![inbound]).await; + + assert!(calls.lock().unwrap().is_empty()); + assert!(gateway.ctx.sent_replies.lock().unwrap().is_empty()); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 1); + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); +} + #[tokio::test(flavor = "current_thread")] async fn slack_download_failure_replies_without_running_an_agent() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/src/imessage/attachments.rs b/src/imessage/attachments.rs new file mode 100644 index 0000000..90bba83 --- /dev/null +++ b/src/imessage/attachments.rs @@ -0,0 +1,480 @@ +//! Safe, worker-time loading of local Messages image attachments. + +use std::ffi::CString; +use std::fs::{DirBuilder, File, OpenOptions}; +use std::io::{Read, Write}; +use std::os::fd::{AsRawFd, FromRawFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::{DirBuilderExt, MetadataExt, OpenOptionsExt, PermissionsExt}; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::time::Duration; + +use anyhow::{bail, Context, Result}; +use tokio::process::Command; +use uuid::Uuid; + +use crate::channel::InboundImage; +use crate::image::{DownloadedImage, MAX_IMAGE_BYTES}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Attachment { + pub locator: String, + pub file_size: Option, + pub mime_type: Option, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum ImageKind { + Direct, + Heic, +} + +pub fn needs_conversion(locator: &str, mime_type: Option<&str>) -> bool { + image_kind(Path::new(locator), mime_type).ok() == Some(ImageKind::Heic) +} + +pub async fn download(attachment_root: &Path, image: &InboundImage) -> Result { + if let Some(bytes) = &image.data { + return Ok(DownloadedImage { + bytes: bytes.clone(), + }); + } + download_with_converter(attachment_root, image, Path::new("/usr/bin/sips")).await +} + +async fn download_with_converter( + attachment_root: &Path, + image: &InboundImage, + converter: &Path, +) -> Result { + let (file, path) = open_attachment(attachment_root, &image.locator)?; + match image_kind(&path, image.mime_type.as_deref())? { + ImageKind::Direct => read_bounded(file), + ImageKind::Heic => convert_heic(file, converter).await, + } +} + +fn open_attachment(attachment_root: &Path, locator: &str) -> Result<(File, PathBuf)> { + open_attachment_with_hook(attachment_root, locator, || {}) +} + +fn open_attachment_with_hook( + attachment_root: &Path, + locator: &str, + before_open: impl FnOnce(), +) -> Result<(File, PathBuf)> { + if locator.trim().is_empty() { + bail!("iMessage attachment omitted its local path"); + } + let root = + std::fs::canonicalize(attachment_root).context("open the Messages attachment directory")?; + let root_metadata = root + .metadata() + .context("inspect the Messages attachment directory")?; + if !root_metadata.is_dir() { + bail!("Messages attachment path is not a directory"); + } + let root_identity = (root_metadata.dev(), root_metadata.ino()); + let locator = Path::new(locator); + let candidate = if locator.is_absolute() { + locator.to_path_buf() + } else if let Ok(relative) = locator.strip_prefix("~/Library/Messages/Attachments") { + root.join(relative) + } else { + root.join(locator) + }; + let path = std::fs::canonicalize(candidate).context("open the iMessage attachment")?; + if !path.starts_with(&root) { + bail!("iMessage attachment is outside the Messages attachment directory"); + } + let relative = path + .strip_prefix(&root) + .context("verify the iMessage attachment path")?; + before_open(); + let file = open_beneath(&root, root_identity, relative)?; + Ok((file, path)) +} + +fn open_beneath(root: &Path, root_identity: (u64, u64), relative: &Path) -> Result { + let parts = relative + .components() + .map(|component| match component { + std::path::Component::Normal(value) => CString::new(value.as_bytes()) + .context("iMessage attachment path contains an invalid name"), + _ => bail!("iMessage attachment path is not relative to the attachment directory"), + }) + .collect::>>()?; + if parts.is_empty() { + bail!("iMessage attachment is not a regular file"); + } + + let mut parent = open_canonical_directory(root)?; + let opened_metadata = parent + .metadata() + .context("inspect the opened Messages attachment directory")?; + if (opened_metadata.dev(), opened_metadata.ino()) != root_identity { + bail!("Messages attachment directory changed during validation"); + } + for (index, part) in parts.iter().enumerate() { + let is_last = index + 1 == parts.len(); + let child = open_component(&parent, part, !is_last)?; + if is_last { + if !child + .metadata() + .context("inspect the iMessage attachment")? + .is_file() + { + bail!("iMessage attachment is not a regular file"); + } + return Ok(child); + } + parent = child; + } + unreachable!("empty attachment paths are rejected above") +} + +fn open_canonical_directory(path: &Path) -> Result { + if !path.is_absolute() { + bail!("Messages attachment directory is not absolute"); + } + let mut directory = OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_DIRECTORY | libc::O_NOFOLLOW) + .open("/") + .context("open the filesystem root")?; + for component in path.components() { + match component { + std::path::Component::RootDir => {} + std::path::Component::Normal(value) => { + let part = CString::new(value.as_bytes()) + .context("Messages attachment path contains an invalid name")?; + directory = open_component(&directory, &part, true)?; + } + _ => bail!("Messages attachment directory is not canonical"), + } + } + Ok(directory) +} + +fn open_component(parent: &File, part: &CString, directory: bool) -> Result { + let mut flags = libc::O_RDONLY | libc::O_CLOEXEC | libc::O_NOFOLLOW; + if directory { + flags |= libc::O_DIRECTORY; + } + // SAFETY: parent is an open directory descriptor, part is a NUL-terminated + // single path component, and a successful descriptor is immediately owned. + let descriptor = unsafe { libc::openat(parent.as_raw_fd(), part.as_ptr(), flags) }; + if descriptor < 0 { + return Err(std::io::Error::last_os_error()).context("open the iMessage attachment safely"); + } + // SAFETY: openat returned a new owned descriptor that is not used elsewhere. + Ok(unsafe { File::from_raw_fd(descriptor) }) +} + +fn image_kind(path: &Path, mime_type: Option<&str>) -> Result { + let mime_type = mime_type.map(|value| value.trim().to_ascii_lowercase()); + match mime_type.as_deref() { + Some("image/jpeg" | "image/jpg" | "image/png" | "image/webp") => { + return Ok(ImageKind::Direct); + } + Some("image/heic" | "image/heif") => return Ok(ImageKind::Heic), + Some("") | Some("application/octet-stream") | None => {} + Some(_) => bail!("iMessage attachment is not a supported image type"), + } + match path + .extension() + .and_then(|value| value.to_str()) + .map(str::to_ascii_lowercase) + .as_deref() + { + Some("jpg" | "jpeg" | "png" | "webp") => Ok(ImageKind::Direct), + Some("heic" | "heif") => Ok(ImageKind::Heic), + _ => bail!("iMessage attachment is not a supported image type"), + } +} + +fn read_bounded(file: File) -> Result { + if file.metadata().context("inspect the iMessage image")?.len() > MAX_IMAGE_BYTES as u64 { + bail!("iMessage image exceeds the 6 MiB limit"); + } + let mut bytes = Vec::new(); + file.take(MAX_IMAGE_BYTES as u64 + 1) + .read_to_end(&mut bytes) + .context("read the iMessage image")?; + if bytes.len() > MAX_IMAGE_BYTES { + bail!("iMessage image exceeds the 6 MiB limit"); + } + Ok(DownloadedImage { bytes }) +} + +async fn convert_heic(mut source: File, converter: &Path) -> Result { + let directory = std::env::temp_dir().join(format!("push-heic-{}", Uuid::new_v4())); + DirBuilder::new() + .mode(0o700) + .create(&directory) + .context("create a private iMessage conversion directory")?; + let mut cleanup = ConversionCleanup { + directory, + input: None, + output: None, + }; + let input = cleanup.directory.join("image.heic"); + let mut private_source = OpenOptions::new() + .write(true) + .create_new(true) + .mode(0o600) + .open(&input) + .context("create a private HEIC conversion input")?; + cleanup.input = Some(input.clone()); + std::io::copy(&mut source, &mut private_source) + .context("copy the HEIC image into the private conversion directory")?; + private_source + .flush() + .context("finish the private HEIC conversion input")?; + drop(private_source); + let output = cleanup.directory.join("image.jpg"); + cleanup.output = Some(output.clone()); + + let mut command = Command::new(converter); + command + .arg("-s") + .arg("format") + .arg("jpeg") + .arg(&input) + .arg("--out") + .arg(&output) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .kill_on_drop(true); + let status = tokio::time::timeout(Duration::from_secs(60), command.status()) + .await + .context("macOS image conversion timed out")? + .context("start macOS image conversion")?; + if !status.success() { + bail!("macOS could not convert the HEIC or HEIF image"); + } + let mut permissions = std::fs::metadata(&output) + .context("inspect the converted iMessage image")? + .permissions(); + permissions.set_mode(0o600); + std::fs::set_permissions(&output, permissions) + .context("protect the converted iMessage image")?; + let converted = OpenOptions::new() + .read(true) + .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(&output) + .context("open the converted iMessage image")?; + read_bounded(converted) +} + +struct ConversionCleanup { + directory: PathBuf, + input: Option, + output: Option, +} + +impl Drop for ConversionCleanup { + fn drop(&mut self) { + if let Some(output) = &self.output { + let _ = std::fs::remove_file(output); + } + if let Some(input) = &self.input { + let _ = std::fs::remove_file(input); + } + let _ = std::fs::remove_dir(&self.directory); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::{sh_arg, temp_dir, FakeCli}; + + fn image(locator: &Path, mime_type: &str) -> InboundImage { + InboundImage { + locator: locator.to_string_lossy().to_string(), + file_size: None, + mime_type: Some(mime_type.to_string()), + data: None, + } + } + + #[tokio::test] + async fn reads_supported_files_only_from_the_attachment_root() { + let root = temp_dir("imessage-attachments"); + let image_path = root.join("nested/image.png"); + std::fs::create_dir_all(image_path.parent().unwrap()).unwrap(); + std::fs::write(&image_path, b"\x89PNG\r\n\x1a\nbody").unwrap(); + + let downloaded = download(&root, &image(&image_path, "image/png")) + .await + .unwrap(); + assert_eq!(downloaded.bytes, b"\x89PNG\r\n\x1a\nbody"); + + let outside = root.parent().unwrap().join("outside.png"); + std::fs::write(&outside, b"\x89PNG\r\n\x1a\nbody").unwrap(); + assert!(download(&root, &image(&outside, "image/png")) + .await + .unwrap_err() + .to_string() + .contains("outside")); + + let symlink = root.join("escaping.png"); + std::os::unix::fs::symlink(&outside, &symlink).unwrap(); + assert!(download(&root, &image(&symlink, "image/png")) + .await + .unwrap_err() + .to_string() + .contains("outside")); + let _ = std::fs::remove_dir_all(root); + let _ = std::fs::remove_file(outside); + } + + #[test] + fn rejects_a_symlink_swap_between_validation_and_open() { + let root = temp_dir("imessage-attachment-race"); + let image_path = root.join("image.png"); + std::fs::write(&image_path, b"trusted").unwrap(); + let outside = root.parent().unwrap().join("raced-outside.png"); + std::fs::write(&outside, b"outside").unwrap(); + + let error = open_attachment_with_hook(&root, image_path.to_str().unwrap(), || { + std::fs::remove_file(&image_path).unwrap(); + std::os::unix::fs::symlink(&outside, &image_path).unwrap(); + }) + .unwrap_err(); + + assert!(error + .to_string() + .contains("open the iMessage attachment safely")); + let _ = std::fs::remove_dir_all(root); + let _ = std::fs::remove_file(outside); + } + + #[test] + fn rejects_an_attachment_root_parent_swap_after_validation() { + let sandbox = temp_dir("imessage-attachment-root-race"); + let messages = sandbox.join("Messages"); + let root = messages.join("Attachments"); + std::fs::create_dir_all(&root).unwrap(); + let image_path = root.join("image.png"); + std::fs::write(&image_path, b"trusted").unwrap(); + + let replacement = sandbox.join("replacement"); + let replacement_root = replacement.join("Attachments"); + std::fs::create_dir_all(&replacement_root).unwrap(); + std::fs::write(replacement_root.join("image.png"), b"outside").unwrap(); + let original_messages = sandbox.join("Messages-original"); + + let error = open_attachment_with_hook(&root, image_path.to_str().unwrap(), || { + std::fs::rename(&messages, &original_messages).unwrap(); + std::os::unix::fs::symlink(&replacement, &messages).unwrap(); + }) + .unwrap_err(); + + assert!(error + .to_string() + .contains("open the iMessage attachment safely")); + std::fs::remove_file(&messages).unwrap(); + std::fs::rename(&original_messages, &messages).unwrap(); + let _ = std::fs::remove_dir_all(sandbox); + } + + #[tokio::test] + async fn rejects_missing_directories_unsupported_types_and_oversized_files() { + let root = temp_dir("imessage-rejected-attachments"); + let missing = root.join("missing.png"); + assert!(download(&root, &image(&missing, "image/png")) + .await + .is_err()); + assert!(download(&root, &image(&root, "image/png")).await.is_err()); + + let pdf = root.join("document.pdf"); + std::fs::write(&pdf, b"%PDF").unwrap(); + assert!(download(&root, &image(&pdf, "application/pdf")) + .await + .is_err()); + + let oversized = root.join("large.png"); + let file = File::create(&oversized).unwrap(); + file.set_len(MAX_IMAGE_BYTES as u64 + 1).unwrap(); + assert!(download(&root, &image(&oversized, "image/png")) + .await + .is_err()); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn converts_heic_to_private_jpeg_and_removes_conversion_files() { + let root = temp_dir("imessage-heic-attachments"); + let source = root.join("photo.heic"); + std::fs::write(&source, b"heic source").unwrap(); + let marker_dir = temp_dir("imessage-converter-marker"); + let marker = marker_dir.join("conversion-paths"); + let converter = FakeCli::new( + "sips", + &format!( + "#!/bin/sh\ninput=\"$4\"\noutput=\"$6\"\nprintf '%s\\n%s' \"$input\" \"$output\" > {}\nprintf '\\377\\330\\377body' > \"$output\"\n", + sh_arg(&marker) + ), + ); + + let downloaded = download_with_converter( + &root, + &image(&source, "image/heic"), + Path::new(&converter.bin()), + ) + .await + .unwrap(); + + assert!(downloaded.bytes.starts_with(&[0xff, 0xd8, 0xff])); + let paths = std::fs::read_to_string(&marker).unwrap(); + let paths = paths.lines().map(Path::new).collect::>(); + assert_eq!(paths.len(), 2); + assert_ne!(paths[0], source); + assert!(!paths[0].exists()); + assert!(!paths[1].exists()); + assert!(!paths[0].parent().unwrap().exists()); + let _ = std::fs::remove_dir_all(root); + let _ = std::fs::remove_dir_all(marker_dir); + } + + #[cfg(target_os = "macos")] + #[tokio::test] + async fn converts_a_real_heic_with_macos_sips() { + use base64::Engine; + + let root = temp_dir("imessage-real-heic"); + let png = root.join("source.png"); + let heic = root.join("photo.heic"); + let bytes = base64::engine::general_purpose::STANDARD + .decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Y9Zl1sAAAAASUVORK5CYII=") + .unwrap(); + std::fs::write(&png, bytes).unwrap(); + let status = Command::new("/usr/bin/sips") + .arg("-s") + .arg("format") + .arg("heic") + .arg(&png) + .arg("--out") + .arg(&heic) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .await + .unwrap(); + assert!(status.success()); + + let downloaded = download_with_converter( + &root, + &image(&heic, "image/heic"), + Path::new("/usr/bin/sips"), + ) + .await + .unwrap(); + + assert!(downloaded.bytes.starts_with(&[0xff, 0xd8, 0xff])); + let _ = std::fs::remove_dir_all(root); + } +} diff --git a/src/imessage/mod.rs b/src/imessage/mod.rs index 7b4e55c..3558a34 100644 --- a/src/imessage/mod.rs +++ b/src/imessage/mod.rs @@ -1,8 +1,10 @@ //! Reading from and sending to the macOS Messages app. +mod attachments; mod attributed_body; mod poller; mod sender; +pub use attachments::{download as download_image, needs_conversion, Attachment}; pub use poller::Poller; pub use sender::Sender; diff --git a/src/imessage/poller.rs b/src/imessage/poller.rs index 06eacef..007a338 100644 --- a/src/imessage/poller.rs +++ b/src/imessage/poller.rs @@ -1,9 +1,12 @@ //! Reads new messages directly from the macOS Messages SQLite database. +use std::path::{Path, PathBuf}; + use anyhow::{Context, Result}; use rusqlite::{Connection, OpenFlags}; use super::attributed_body; +use super::Attachment; /// One inbound iMessage row relevant to the gateway. #[derive(Debug, Clone)] @@ -17,12 +20,14 @@ pub struct Message { pub is_group: bool, pub text: String, pub is_from_me: bool, + pub attachments: Vec, } /// Poller reads `chat.db` read-only via rusqlite. #[derive(Clone)] pub struct Poller { db_path: String, + attachment_root: PathBuf, } // Only real text messages: exclude tapbacks/reactions (associated_message_type) @@ -38,9 +43,27 @@ LEFT JOIN chat c ON c.ROWID = cmj.chat_id \ WHERE m.ROWID > ?1 AND m.associated_message_type = 0 AND m.item_type = 0 \ ORDER BY m.ROWID ASC"; +const SELECT_ATTACHMENTS: &str = "\ +SELECT COALESCE(a.filename, ''), a.total_bytes, a.mime_type \ +FROM message_attachment_join maj \ +JOIN attachment a ON a.ROWID = maj.attachment_id \ +WHERE maj.message_id = ?1 \ + AND COALESCE(a.is_sticker, 0) = 0 \ + AND COALESCE(a.hide_attachment, 0) = 0 \ + AND COALESCE(a.mime_type, '') NOT LIKE 'video/%' \ + AND COALESCE(a.uti, '') NOT IN ('public.movie', 'com.apple.quicktime-movie') \ +ORDER BY maj.ROWID ASC, a.ROWID ASC"; + impl Poller { pub fn new(db_path: String) -> Self { - Self { db_path } + let attachment_root = Path::new(&db_path) + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("Attachments"); + Self { + db_path, + attachment_root, + } } fn open(&self) -> Result { @@ -69,12 +92,26 @@ impl Poller { chat_identifier: row.get(5)?, is_group: participant_count > 1, is_from_me: is_from_me == 1, + attachments: Vec::new(), }) })?; let mut out = Vec::new(); for r in rows { out.push(r?); } + let mut attachments = conn.prepare(SELECT_ATTACHMENTS)?; + for message in &mut out { + message.attachments = attachments + .query_map([message.row_id], |row| { + let bytes: Option = row.get(1)?; + Ok(Attachment { + locator: row.get(0)?, + file_size: bytes.and_then(|value| usize::try_from(value).ok()), + mime_type: row.get(2)?, + }) + })? + .collect::, _>>()?; + } Ok(out) } @@ -87,12 +124,19 @@ impl Poller { })?; Ok(v) } + + pub async fn download_image( + &self, + image: &crate::channel::InboundImage, + ) -> Result { + super::download_image(&self.attachment_root, image).await + } } #[cfg(test)] mod tests { use super::*; - use crate::test_support::temp_path; + use crate::test_support::{temp_dir, temp_path}; use rusqlite::Connection; #[test] @@ -157,6 +201,82 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn poll_joins_ordered_attachment_metadata_without_opening_files() { + let root = temp_dir("chat-db-attachments"); + let path = root.join("chat.db"); + let conn = Connection::open(&path).unwrap(); + create_schema(&conn); + insert_handle(&conn, 1, "+15551234567"); + insert_chat(&conn, 1, "+15551234567"); + insert_chat_handle(&conn, 1, 1); + insert_message(&conn, 1, Some(""), None, 0, 0, 0, 1, 1); + insert_attachment( + &conn, + 1, + 1, + "missing-first.png", + Some("image/png"), + Some(12), + ); + insert_attachment( + &conn, + 2, + 1, + "missing-second.heic", + Some("image/heic"), + Some(24), + ); + conn.execute( + "INSERT INTO attachment ( + ROWID, filename, uti, mime_type, total_bytes, is_sticker + ) VALUES (3, 'sticker.png', 'public.png', 'image/png', 10, 1)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (1, 3)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO attachment ( + ROWID, filename, uti, mime_type, total_bytes + ) VALUES (4, 'live.mov', 'com.apple.quicktime-movie', 'video/quicktime', 10)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (1, 4)", + [], + ) + .unwrap(); + drop(conn); + + let got = Poller::new(path.to_string_lossy().to_string()) + .poll(0) + .unwrap(); + + assert_eq!(got.len(), 1); + assert!(got[0].text.is_empty()); + assert_eq!( + got[0].attachments, + vec![ + Attachment { + locator: "missing-first.png".to_string(), + file_size: Some(12), + mime_type: Some("image/png".to_string()), + }, + Attachment { + locator: "missing-second.heic".to_string(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + }, + ] + ); + let _ = std::fs::remove_dir_all(root); + } + fn create_schema(conn: &Connection) { conn.execute_batch( " @@ -173,6 +293,19 @@ mod tests { CREATE TABLE chat (ROWID INTEGER PRIMARY KEY, chat_identifier TEXT NOT NULL); CREATE TABLE chat_message_join (message_id INTEGER NOT NULL, chat_id INTEGER NOT NULL); CREATE TABLE chat_handle_join (chat_id INTEGER NOT NULL, handle_id INTEGER NOT NULL); + CREATE TABLE attachment ( + ROWID INTEGER PRIMARY KEY, + filename TEXT, + uti TEXT, + mime_type TEXT, + total_bytes INTEGER, + is_sticker INTEGER DEFAULT 0, + hide_attachment INTEGER DEFAULT 0 + ); + CREATE TABLE message_attachment_join ( + message_id INTEGER NOT NULL, + attachment_id INTEGER NOT NULL + ); ", ) .unwrap(); @@ -202,6 +335,27 @@ mod tests { .unwrap(); } + fn insert_attachment( + conn: &Connection, + row_id: i64, + message_id: i64, + filename: &str, + mime_type: Option<&str>, + total_bytes: Option, + ) { + conn.execute( + "INSERT INTO attachment (ROWID, filename, mime_type, total_bytes) + VALUES (?1, ?2, ?3, ?4)", + (row_id, filename, mime_type, total_bytes), + ) + .unwrap(); + conn.execute( + "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (?1, ?2)", + (message_id, row_id), + ) + .unwrap(); + } + #[allow(clippy::too_many_arguments)] fn insert_message( conn: &Connection, From f58c140f9d5e4f525139adb85c1c7d52d8fa88c1 Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 22:42:22 +0100 Subject: [PATCH 2/6] fix(imessage): bound HEIC conversion input --- docs/channels/imessage.md | 3 ++- src/imessage/attachments.rs | 38 ++++++++++++++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/docs/channels/imessage.md b/docs/channels/imessage.md index addf539..1fbf225 100644 --- a/docs/channels/imessage.md +++ b/docs/channels/imessage.md @@ -51,7 +51,8 @@ can ask the agent to use any capability allowed by that agent's configuration. Send up to four JPEG, PNG, WebP, HEIC, or HEIF images in one accepted conversation, with or without message text. Their combined prepared size must -be at most 6 MiB. Images work with Claude Code, Codex, and Pi. +be at most 6 MiB. Each HEIC or HEIF source must be at most 32 MiB before local +conversion. Images work with Claude Code, Codex, and Pi. Polling reads attachment paths, byte-size hints, and MIME type hints from `chat.db`; it does not open the attachment files. After the direct-message, diff --git a/src/imessage/attachments.rs b/src/imessage/attachments.rs index 90bba83..35b02c6 100644 --- a/src/imessage/attachments.rs +++ b/src/imessage/attachments.rs @@ -17,6 +17,8 @@ use uuid::Uuid; use crate::channel::InboundImage; use crate::image::{DownloadedImage, MAX_IMAGE_BYTES}; +const MAX_HEIC_INPUT_BYTES: u64 = 32 * 1024 * 1024; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct Attachment { pub locator: String, @@ -208,7 +210,7 @@ fn read_bounded(file: File) -> Result { Ok(DownloadedImage { bytes }) } -async fn convert_heic(mut source: File, converter: &Path) -> Result { +async fn convert_heic(source: File, converter: &Path) -> Result { let directory = std::env::temp_dir().join(format!("push-heic-{}", Uuid::new_v4())); DirBuilder::new() .mode(0o700) @@ -227,8 +229,14 @@ async fn convert_heic(mut source: File, converter: &Path) -> Result MAX_HEIC_INPUT_BYTES { + bail!("HEIC or HEIF image exceeds the 32 MiB conversion input limit"); + } private_source .flush() .context("finish the private HEIC conversion input")?; @@ -440,6 +448,30 @@ mod tests { let _ = std::fs::remove_dir_all(marker_dir); } + #[tokio::test] + async fn bounds_heic_input_before_starting_conversion() { + let root = temp_dir("imessage-large-heic"); + let source = root.join("large.heic"); + let source_file = File::create(&source).unwrap(); + source_file.set_len(MAX_HEIC_INPUT_BYTES + 1).unwrap(); + let marker_dir = temp_dir("imessage-large-heic-marker"); + let marker = marker_dir.join("converter-ran"); + let converter = FakeCli::new("sips", &format!("#!/bin/sh\ntouch {}\n", sh_arg(&marker))); + + let error = download_with_converter( + &root, + &image(&source, "image/heic"), + Path::new(&converter.bin()), + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("32 MiB")); + assert!(!marker.exists()); + let _ = std::fs::remove_dir_all(root); + let _ = std::fs::remove_dir_all(marker_dir); + } + #[cfg(target_os = "macos")] #[tokio::test] async fn converts_a_real_heic_with_macos_sips() { From bcce44324aa1c2a22783b6fa01a0cb9465e26464 Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 22:55:06 +0100 Subject: [PATCH 3/6] fix(imessage): bound pending attachment deferral --- docs/channels/imessage.md | 6 +- src/imessage/poller.rs | 161 +++++++++++++++++++++++++++++++++++--- 2 files changed, 154 insertions(+), 13 deletions(-) diff --git a/docs/channels/imessage.md b/docs/channels/imessage.md index 1fbf225..6abcefb 100644 --- a/docs/channels/imessage.md +++ b/docs/channels/imessage.md @@ -55,7 +55,11 @@ be at most 6 MiB. Each HEIC or HEIF source must be at most 32 MiB before local conversion. Images work with Claude Code, Codex, and Pi. Polling reads attachment paths, byte-size hints, and MIME type hints from -`chat.db`; it does not open the attachment files. After the direct-message, +`chat.db`; it does not open the attachment files. If Messages has not populated +an attachment filename yet, Push gives it a three-poll grace period and defers +that message and later rows so the cursor cannot skip the image. If the filename +is still blank after that grace period, the worker treats it as missing, sends +the safe fallback, and lets later messages continue. After the direct-message, sender, reply-marker, and message checks pass, the worker canonicalizes each path and requires it to remain under `~/Library/Messages/Attachments` (or the `Attachments` directory beside a custom `imessage.db_path`). Missing files, diff --git a/src/imessage/poller.rs b/src/imessage/poller.rs index 007a338..a8b6d68 100644 --- a/src/imessage/poller.rs +++ b/src/imessage/poller.rs @@ -1,6 +1,8 @@ //! Reads new messages directly from the macOS Messages SQLite database. +use std::collections::HashMap; use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; use anyhow::{Context, Result}; use rusqlite::{Connection, OpenFlags}; @@ -28,8 +30,11 @@ pub struct Message { pub struct Poller { db_path: String, attachment_root: PathBuf, + pending_filename_polls: Arc>>, } +const PENDING_FILENAME_POLL_LIMIT: u8 = 3; + // Only real text messages: exclude tapbacks/reactions (associated_message_type) // and system rows like group renames or joins (item_type). const SELECT_NEW: &str = "\ @@ -44,7 +49,7 @@ WHERE m.ROWID > ?1 AND m.associated_message_type = 0 AND m.item_type = 0 \ ORDER BY m.ROWID ASC"; const SELECT_ATTACHMENTS: &str = "\ -SELECT COALESCE(a.filename, ''), a.total_bytes, a.mime_type \ +SELECT a.filename, a.total_bytes, a.mime_type \ FROM message_attachment_join maj \ JOIN attachment a ON a.ROWID = maj.attachment_id \ WHERE maj.message_id = ?1 \ @@ -63,6 +68,7 @@ impl Poller { Self { db_path, attachment_root, + pending_filename_polls: Arc::new(Mutex::new(HashMap::new())), } } @@ -75,6 +81,7 @@ impl Poller { /// `attributedBody` when `text` is NULL. This is blocking; call it from a /// blocking context. pub fn poll(&self, since: i64) -> Result> { + self.clear_completed_pending_filenames(since)?; let conn = self.open()?; let mut stmt = conn.prepare(SELECT_NEW)?; let rows = stmt.query_map([since], |row| { @@ -100,21 +107,69 @@ impl Poller { out.push(r?); } let mut attachments = conn.prepare(SELECT_ATTACHMENTS)?; - for message in &mut out { - message.attachments = attachments - .query_map([message.row_id], |row| { - let bytes: Option = row.get(1)?; - Ok(Attachment { - locator: row.get(0)?, - file_size: bytes.and_then(|value| usize::try_from(value).ok()), - mime_type: row.get(2)?, - }) - })? - .collect::, _>>()?; + let mut ready_count = out.len(); + for (index, message) in out.iter_mut().enumerate() { + let rows = attachments.query_map([message.row_id], |row| { + let locator: Option = row.get(0)?; + let bytes: Option = row.get(1)?; + let mime_type: Option = row.get(2)?; + Ok((locator, bytes, mime_type)) + })?; + let mut message_attachments = Vec::new(); + let mut has_pending_filename = false; + for row in rows { + let (locator, bytes, mime_type) = row?; + let locator = locator.unwrap_or_default(); + has_pending_filename |= locator.trim().is_empty(); + message_attachments.push(Attachment { + locator, + file_size: bytes.and_then(|value| usize::try_from(value).ok()), + mime_type, + }); + } + if has_pending_filename { + if self.should_defer_filename(message.row_id)? { + ready_count = index; + break; + } + } else { + self.clear_pending_filename(message.row_id)?; + } + message.attachments = message_attachments; } + out.truncate(ready_count); Ok(out) } + fn should_defer_filename(&self, row_id: i64) -> Result { + let mut pending = self + .pending_filename_polls + .lock() + .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))?; + let polls = pending.entry(row_id).or_default(); + if *polls < PENDING_FILENAME_POLL_LIMIT { + *polls += 1; + return Ok(true); + } + Ok(false) + } + + fn clear_completed_pending_filenames(&self, since: i64) -> Result<()> { + self.pending_filename_polls + .lock() + .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))? + .retain(|row_id, _| *row_id > since); + Ok(()) + } + + fn clear_pending_filename(&self, row_id: i64) -> Result<()> { + self.pending_filename_polls + .lock() + .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))? + .remove(&row_id); + Ok(()) + } + /// Highest message ROWID in the database, or 0 when empty. Used to skip /// backlog on first run. pub fn max_row_id(&self) -> Result { @@ -277,6 +332,88 @@ mod tests { let _ = std::fs::remove_dir_all(root); } + #[test] + fn poll_defers_a_message_and_later_rows_until_attachment_filename_is_ready() { + let root = temp_dir("chat-db-pending-attachment"); + let path = root.join("chat.db"); + let conn = Connection::open(&path).unwrap(); + create_schema(&conn); + insert_handle(&conn, 1, "+15551234567"); + insert_chat(&conn, 1, "+15551234567"); + insert_chat_handle(&conn, 1, 1); + insert_message(&conn, 1, Some("photo"), None, 0, 0, 0, 1, 1); + insert_message(&conn, 2, Some("later"), None, 0, 0, 0, 1, 1); + conn.execute( + "INSERT INTO attachment (ROWID, filename, mime_type, total_bytes) + VALUES (1, NULL, 'image/heic', 24)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (1, 1)", + [], + ) + .unwrap(); + drop(conn); + + let poller = Poller::new(path.to_string_lossy().to_string()); + assert!(poller.poll(0).unwrap().is_empty()); + + let conn = Connection::open(&path).unwrap(); + conn.execute( + "UPDATE attachment SET filename = 'missing-photo.heic' WHERE ROWID = 1", + [], + ) + .unwrap(); + drop(conn); + let ready = poller.poll(0).unwrap(); + assert_eq!(ready.len(), 2); + assert_eq!(ready[0].row_id, 1); + assert_eq!(ready[0].attachments[0].locator, "missing-photo.heic"); + assert_eq!(ready[1].row_id, 2); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn poll_bounds_pending_filename_deferral_so_later_rows_can_progress() { + let root = temp_dir("chat-db-abandoned-attachment"); + let path = root.join("chat.db"); + let conn = Connection::open(&path).unwrap(); + create_schema(&conn); + insert_handle(&conn, 1, "+15551234567"); + insert_chat(&conn, 1, "+15551234567"); + insert_chat_handle(&conn, 1, 1); + insert_message(&conn, 1, Some("failed photo"), None, 0, 0, 0, 1, 1); + insert_message(&conn, 2, Some("later"), None, 0, 0, 0, 1, 1); + conn.execute( + "INSERT INTO attachment (ROWID, filename, mime_type, total_bytes) + VALUES (1, NULL, 'image/heic', 24)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (1, 1)", + [], + ) + .unwrap(); + drop(conn); + + let poller = Poller::new(path.to_string_lossy().to_string()); + for _ in 0..PENDING_FILENAME_POLL_LIMIT { + assert!(poller.poll(0).unwrap().is_empty()); + } + let ready = poller.poll(0).unwrap(); + assert_eq!(ready.len(), 2); + assert_eq!(ready[0].row_id, 1); + assert_eq!(ready[0].attachments[0].locator, ""); + assert_eq!(ready[1].row_id, 2); + assert_eq!(poller.poll(0).unwrap().len(), 2); + let after_failed = poller.poll(1).unwrap(); + assert_eq!(after_failed.len(), 1); + assert_eq!(after_failed[0].row_id, 2); + let _ = std::fs::remove_dir_all(root); + } + fn create_schema(conn: &Connection) { conn.execute_batch( " From 639892c8f07d1e2c17547084784440f3abbeb09c Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 23:16:13 +0100 Subject: [PATCH 4/6] fix(imessage): defer pending files after acceptance --- docs/channels/imessage.md | 16 ++-- src/channel.rs | 36 ++++++++ src/gateway/mod.rs | 84 +++++++++++++++++-- src/gateway/tests.rs | 172 +++++++++++++++++++++++++++++++++++++- src/imessage/poller.rs | 60 +++++++------ 5 files changed, 328 insertions(+), 40 deletions(-) diff --git a/docs/channels/imessage.md b/docs/channels/imessage.md index 6abcefb..545ed48 100644 --- a/docs/channels/imessage.md +++ b/docs/channels/imessage.md @@ -55,14 +55,14 @@ be at most 6 MiB. Each HEIC or HEIF source must be at most 32 MiB before local conversion. Images work with Claude Code, Codex, and Pi. Polling reads attachment paths, byte-size hints, and MIME type hints from -`chat.db`; it does not open the attachment files. If Messages has not populated -an attachment filename yet, Push gives it a three-poll grace period and defers -that message and later rows so the cursor cannot skip the image. If the filename -is still blank after that grace period, the worker treats it as missing, sends -the safe fallback, and lets later messages continue. After the direct-message, -sender, reply-marker, and message checks pass, the worker canonicalizes each -path and requires it to remain under `~/Library/Messages/Attachments` (or the -`Attachments` directory beside a custom `imessage.db_path`). Missing files, +`chat.db`; it does not open the attachment files. After the direct-message, +sender, reply-marker, and message checks pass, an accepted attachment with no +filename gets a three-poll grace period. Push leaves that row unacknowledged so +the cursor cannot skip the image, while rejected and later ready messages can +continue. If the filename is still blank after the grace period, the worker +treats it as missing and sends the safe fallback. The worker canonicalizes each +ready path and requires it to remain under `~/Library/Messages/Attachments` (or +the `Attachments` directory beside a custom `imessage.db_path`). Missing files, directories, escaping symlinks, and unsupported documents are rejected with a safe retry reply before an agent starts. diff --git a/src/channel.rs b/src/channel.rs index a112e35..1370853 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -97,8 +97,14 @@ trait ChannelContract { fn id(&self) -> &'static str; fn primary_target(&self, configured: &str) -> Result; async fn poll(&self, since: i64) -> Result>; + fn poll_is_complete_snapshot(&self) -> bool { + false + } async fn latest_cursor(&self) -> Result; fn accept(&self, message: &RawMessage) -> Option<(String, String)>; + fn should_defer(&self, _message: &RawMessage) -> Result { + Ok(false) + } fn reject_reason(&self, message: &RawMessage) -> &'static str; fn approval_origin(&self, message: &RawMessage, thread: &str) -> AnswerOrigin; fn route_thread_groups(&self, thread: &str) -> Vec>; @@ -222,6 +228,14 @@ impl Channel { } } + pub fn poll_is_complete_snapshot(&self) -> bool { + match self { + Self::IMessage(channel) => ChannelContract::poll_is_complete_snapshot(channel), + Self::Telegram(channel) => ChannelContract::poll_is_complete_snapshot(channel), + Self::Slack(channel) => ChannelContract::poll_is_complete_snapshot(channel), + } + } + /// Returns `(thread_key, reply_target)` for an accepted message. pub fn accept(&self, message: &RawMessage) -> Option<(String, String)> { match self { @@ -231,6 +245,14 @@ impl Channel { } } + pub fn should_defer(&self, message: &RawMessage) -> Result { + match self { + Self::IMessage(channel) => ChannelContract::should_defer(channel, message), + Self::Telegram(channel) => ChannelContract::should_defer(channel, message), + Self::Slack(channel) => ChannelContract::should_defer(channel, message), + } + } + pub fn reject_reason(&self, message: &RawMessage) -> &'static str { match self { Self::IMessage(channel) => ChannelContract::reject_reason(channel, message), @@ -408,6 +430,10 @@ impl ChannelContract for IMessageChannel { tokio::task::spawn_blocking(move || poller.max_row_id()).await? } + fn poll_is_complete_snapshot(&self) -> bool { + true + } + fn accept(&self, message: &RawMessage) -> Option<(String, String)> { if common_reject_reason(message).is_some() || (!self.reply_marker.is_empty() && message.text.contains(&self.reply_marker)) @@ -430,6 +456,16 @@ impl ChannelContract for IMessageChannel { None } + fn should_defer(&self, message: &RawMessage) -> Result { + self.poller.should_defer_pending_filename( + message.row_id, + message + .images + .iter() + .any(|image| image.locator.trim().is_empty()), + ) + } + fn reject_reason(&self, message: &RawMessage) -> &'static str { common_reject_reason(message).unwrap_or_else(|| { if !self.reply_marker.is_empty() && message.text.contains(&self.reply_marker) { diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 8b872de..9129950 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -99,6 +99,7 @@ pub struct PrimaryDestination { #[derive(Default)] struct AckState { in_flight: BTreeSet, + deferred: BTreeSet, completed: BTreeSet, } @@ -570,18 +571,24 @@ impl Gateway { } }; - self.process_messages(msgs).await; + let complete_snapshot = self.channel.poll_is_complete_snapshot(); + self.process_messages(msgs, complete_snapshot).await; } #[cfg(test)] async fn tick_fake(&mut self, msgs: Vec) { - self.process_messages(msgs).await; + self.process_messages(msgs, false).await; } - async fn process_messages(&mut self, msgs: Vec) { + #[cfg(test)] + async fn tick_fake_complete(&mut self, msgs: Vec) { + self.process_messages(msgs, true).await; + } + + async fn process_messages(&mut self, msgs: Vec, complete_snapshot: bool) { self.recover_closed_workers(); persist_cursor(&self.store, &self.ack, self.channel.id()); - let since = match self.store.lock().unwrap().cursor(self.channel.id()) { + let mut since = match self.store.lock().unwrap().cursor(self.channel.id()) { Ok(cursor) => cursor, Err(error) => { error!( @@ -591,6 +598,46 @@ impl Gateway { return; } }; + if complete_snapshot { + let visible = msgs + .iter() + .filter(|message| message.row_id > since) + .map(|message| message.row_id) + .collect::>(); + let missing = self + .ack + .lock() + .unwrap() + .deferred + .iter() + .filter(|row_id| !visible.contains(row_id)) + .copied() + .collect::>(); + if !missing.is_empty() { + { + let mut ack = self.ack.lock().unwrap(); + for row_id in &missing { + ack.deferred.remove(row_id); + ack.completed.insert(*row_id); + warn!( + "deferred {} row {row_id} disappeared before processing", + self.channel.id() + ); + } + } + persist_cursor(&self.store, &self.ack, self.channel.id()); + since = match self.store.lock().unwrap().cursor(self.channel.id()) { + Ok(cursor) => cursor, + Err(error) => { + error!( + "{} cursor read error after deferred reconciliation: {error:#}", + self.channel.id() + ); + return; + } + }; + } + } for m in &msgs { if m.row_id <= since { continue; @@ -599,6 +646,25 @@ impl Gateway { continue; } if let Some((thread, target)) = self.channel.accept(m) { + let deferred = match self.channel.should_defer(m) { + Ok(deferred) => deferred, + Err(error) => { + error!("[{thread}] attachment readiness check failed: {error:#}"); + self.audit(self.ctx.audit.failed( + "message_defer_failed", + m.row_id, + &thread, + None, + error.to_string(), + )); + return; + } + }; + if deferred { + self.ack.lock().unwrap().deferred.insert(m.row_id); + info!("[{thread}] waiting for iMessage attachment filename"); + continue; + } let reply_with_voice = m.voice.is_some(); let message_text = if reply_with_voice { "[Voice message]".to_string() @@ -1236,6 +1302,7 @@ fn complete_row(store: &Arc>, ack: &Arc>, channel: { let mut ack = ack.lock().unwrap(); ack.in_flight.remove(&row_id); + ack.deferred.remove(&row_id); ack.completed.insert(row_id); } persist_cursor(store, ack, channel); @@ -1265,7 +1332,14 @@ impl AckState { } fn next_cursor(&self) -> Option { - let limit = self.in_flight.first().copied().unwrap_or(i64::MAX); + let limit = self + .in_flight + .first() + .into_iter() + .chain(self.deferred.first()) + .copied() + .min() + .unwrap_or(i64::MAX); self.completed .iter() .copied() diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index fbb4e80..70296af 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -2068,6 +2068,170 @@ async fn rejected_imessage_sender_does_not_open_local_attachments() { let _ = std::fs::remove_dir_all(assistant_dir); } +#[tokio::test(flavor = "current_thread")] +async fn imessage_pending_filename_defers_only_accepted_rows_and_cannot_block_forever() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-pending-filename-sessions"); + let assistant_dir = temp_path("imessage-pending-filename-assistant"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut gateway = Gateway::new(test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + )) + .unwrap(); + gateway.ctx.runners = Arc::new(fake_runners(calls.clone())); + + let mut rejected = message(1, "+15550000000", "+15550000000", false, "ignored"); + rejected.images.push(InboundImage { + locator: String::new(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + data: None, + }); + let accepted = message( + 2, + "+15551234567", + "+15551234567", + false, + "continues immediately", + ); + run_messages(&mut gateway, vec![rejected, accepted]).await; + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 2); + + let mut pending = message(3, "+15551234567", "+15551234567", false, "photo"); + pending.images.push(InboundImage { + locator: String::new(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + data: None, + }); + let later = message( + 4, + "+15551234567", + "+15551234567", + false, + "runs while photo waits", + ); + let batch = vec![pending, later]; + for _ in 0..3 { + run_messages(&mut gateway, batch.clone()).await; + } + assert_eq!(calls.lock().unwrap().len(), 2); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 2); + + run_messages(&mut gateway, batch).await; + assert_eq!(calls.lock().unwrap().len(), 2); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 4); + assert!(gateway + .ctx + .sent_replies + .lock() + .unwrap() + .last() + .unwrap() + .1 + .contains("JPEG, PNG, or WebP")); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(format!("{state_path}.cache")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); +} + +#[tokio::test(flavor = "current_thread")] +async fn imessage_deferred_barrier_survives_preworker_failure() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-deferred-failure-sessions"); + let assistant_dir = temp_path("imessage-deferred-failure-assistant"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut gateway = Gateway::new(test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + )) + .unwrap(); + gateway.ctx.runners = Arc::new(fake_runners(calls.clone())); + + let mut pending = message(1, "+15551234567", "+15551234567", false, "photo"); + pending.images.push(InboundImage { + locator: String::new(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + data: None, + }); + let later = message(2, "+15551234567", "+15551234567", false, "later"); + let batch = vec![pending, later]; + for _ in 0..3 { + run_messages(&mut gateway, batch.clone()).await; + } + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 0); + + gateway + .ctx + .history + .lock() + .unwrap() + .execute_batch_for_test("DROP TABLE messages"); + run_messages(&mut gateway, batch).await; + + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 0); + assert!(gateway.ack.lock().unwrap().deferred.contains(&1)); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); +} + +#[tokio::test(flavor = "current_thread")] +async fn imessage_complete_poll_reconciles_a_deleted_deferred_row() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-deferred-deleted-sessions"); + let assistant_dir = temp_path("imessage-deferred-deleted-assistant"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let calls = Arc::new(Mutex::new(Vec::::new())); + let mut gateway = Gateway::new(test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + )) + .unwrap(); + gateway.ctx.runners = Arc::new(fake_runners(calls.clone())); + + let mut pending = message(1, "+15551234567", "+15551234567", false, "photo"); + pending.images.push(InboundImage { + locator: String::new(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + data: None, + }); + let later = message(2, "+15551234567", "+15551234567", false, "later"); + run_complete_snapshot(&mut gateway, vec![pending, later.clone()]).await; + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 0); + + run_complete_snapshot(&mut gateway, vec![later]).await; + + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(gateway.store.lock().unwrap().cursor("imessage").unwrap(), 2); + assert!(gateway.ack.lock().unwrap().deferred.is_empty()); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); +} + #[tokio::test(flavor = "current_thread")] async fn slack_download_failure_replies_without_running_an_agent() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -2512,7 +2676,7 @@ async fn closed_worker_queue_is_recovered_without_another_message() { }, ); - gateway.process_messages(Vec::new()).await; + gateway.process_messages(Vec::new(), false).await; gateway.queues.clear(); gateway.drain_workers().await; @@ -3739,6 +3903,12 @@ async fn run_messages(gateway: &mut Gateway, messages: Vec) { gateway.drain_workers().await; } +async fn run_complete_snapshot(gateway: &mut Gateway, messages: Vec) { + gateway.tick_fake_complete(messages).await; + gateway.queues.clear(); + gateway.drain_workers().await; +} + fn approval_question( channel: &str, thread: &str, diff --git a/src/imessage/poller.rs b/src/imessage/poller.rs index a8b6d68..353d396 100644 --- a/src/imessage/poller.rs +++ b/src/imessage/poller.rs @@ -107,8 +107,7 @@ impl Poller { out.push(r?); } let mut attachments = conn.prepare(SELECT_ATTACHMENTS)?; - let mut ready_count = out.len(); - for (index, message) in out.iter_mut().enumerate() { + for message in &mut out { let rows = attachments.query_map([message.row_id], |row| { let locator: Option = row.get(0)?; let bytes: Option = row.get(1)?; @@ -116,31 +115,29 @@ impl Poller { Ok((locator, bytes, mime_type)) })?; let mut message_attachments = Vec::new(); - let mut has_pending_filename = false; for row in rows { let (locator, bytes, mime_type) = row?; let locator = locator.unwrap_or_default(); - has_pending_filename |= locator.trim().is_empty(); message_attachments.push(Attachment { locator, file_size: bytes.and_then(|value| usize::try_from(value).ok()), mime_type, }); } - if has_pending_filename { - if self.should_defer_filename(message.row_id)? { - ready_count = index; - break; - } - } else { - self.clear_pending_filename(message.row_id)?; - } message.attachments = message_attachments; } - out.truncate(ready_count); Ok(out) } + pub fn should_defer_pending_filename(&self, row_id: i64, pending: bool) -> Result { + if pending { + self.should_defer_filename(row_id) + } else { + self.clear_pending_filename(row_id)?; + Ok(false) + } + } + fn should_defer_filename(&self, row_id: i64) -> Result { let mut pending = self .pending_filename_polls @@ -333,7 +330,7 @@ mod tests { } #[test] - fn poll_defers_a_message_and_later_rows_until_attachment_filename_is_ready() { + fn pending_filename_deferral_ends_when_attachment_filename_is_ready() { let root = temp_dir("chat-db-pending-attachment"); let path = root.join("chat.db"); let conn = Connection::open(&path).unwrap(); @@ -357,7 +354,12 @@ mod tests { drop(conn); let poller = Poller::new(path.to_string_lossy().to_string()); - assert!(poller.poll(0).unwrap().is_empty()); + let pending = poller.poll(0).unwrap(); + assert_eq!(pending.len(), 2); + assert_eq!(pending[0].attachments[0].locator, ""); + assert!(poller + .should_defer_pending_filename(pending[0].row_id, true) + .unwrap()); let conn = Connection::open(&path).unwrap(); conn.execute( @@ -370,12 +372,15 @@ mod tests { assert_eq!(ready.len(), 2); assert_eq!(ready[0].row_id, 1); assert_eq!(ready[0].attachments[0].locator, "missing-photo.heic"); + assert!(!poller + .should_defer_pending_filename(ready[0].row_id, false) + .unwrap()); assert_eq!(ready[1].row_id, 2); let _ = std::fs::remove_dir_all(root); } #[test] - fn poll_bounds_pending_filename_deferral_so_later_rows_can_progress() { + fn pending_filename_deferral_is_bounded_and_stable_until_cursor_progress() { let root = temp_dir("chat-db-abandoned-attachment"); let path = root.join("chat.db"); let conn = Connection::open(&path).unwrap(); @@ -399,18 +404,21 @@ mod tests { drop(conn); let poller = Poller::new(path.to_string_lossy().to_string()); + let messages = poller.poll(0).unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].attachments[0].locator, ""); for _ in 0..PENDING_FILENAME_POLL_LIMIT { - assert!(poller.poll(0).unwrap().is_empty()); + assert!(poller + .should_defer_pending_filename(messages[0].row_id, true) + .unwrap()); } - let ready = poller.poll(0).unwrap(); - assert_eq!(ready.len(), 2); - assert_eq!(ready[0].row_id, 1); - assert_eq!(ready[0].attachments[0].locator, ""); - assert_eq!(ready[1].row_id, 2); - assert_eq!(poller.poll(0).unwrap().len(), 2); - let after_failed = poller.poll(1).unwrap(); - assert_eq!(after_failed.len(), 1); - assert_eq!(after_failed[0].row_id, 2); + assert!(!poller + .should_defer_pending_filename(messages[0].row_id, true) + .unwrap()); + assert!(!poller + .should_defer_pending_filename(messages[0].row_id, true) + .unwrap()); + assert_eq!(poller.poll(1).unwrap().len(), 1); let _ = std::fs::remove_dir_all(root); } From d83e8c9bac855bed26ebd8cd167d4e1e51c49774 Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 23:33:20 +0100 Subject: [PATCH 5/6] fix(imessage): persist deferred row completion --- src/gateway/mod.rs | 70 +++++++++++++++++++++- src/gateway/tests.rs | 104 +++++++++++++++++++++++++++++++++ src/history.rs | 14 ++++- src/store.rs | 134 ++++++++++++++++++++++++++++++++++++++++--- 4 files changed, 311 insertions(+), 11 deletions(-) diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 9129950..109826f 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -100,6 +100,7 @@ pub struct PrimaryDestination { struct AckState { in_flight: BTreeSet, deferred: BTreeSet, + persisting: BTreeSet, completed: BTreeSet, } @@ -587,6 +588,7 @@ impl Gateway { async fn process_messages(&mut self, msgs: Vec, complete_snapshot: bool) { self.recover_closed_workers(); + retry_completion_persistence(&self.store, &self.ack, self.channel.id()); persist_cursor(&self.store, &self.ack, self.channel.id()); let mut since = match self.store.lock().unwrap().cursor(self.channel.id()) { Ok(cursor) => cursor, @@ -638,10 +640,34 @@ impl Gateway { }; } } + let durable_completed = if complete_snapshot { + match self + .store + .lock() + .unwrap() + .completed_rows_after(self.channel.id(), since) + { + Ok(rows) => rows.into_iter().collect::>(), + Err(error) => { + error!( + "{} completed row read error; polling paused: {error:#}", + self.channel.id() + ); + return; + } + } + } else { + BTreeSet::new() + }; for m in &msgs { if m.row_id <= since { continue; } + if durable_completed.contains(&m.row_id) { + self.ack.lock().unwrap().completed.insert(m.row_id); + persist_cursor(&self.store, &self.ack, self.channel.id()); + continue; + } if self.ack.lock().unwrap().is_known(m.row_id) { continue; } @@ -1299,15 +1325,52 @@ async fn send_scheduled_chunk( } fn complete_row(store: &Arc>, ack: &Arc>, channel: &str, row_id: i64) { + let persisted = match store.lock().unwrap().mark_row_completed(channel, row_id) { + Ok(()) => true, + Err(error) => { + error!("persist {channel} completed row {row_id}: {error:#}"); + false + } + }; { let mut ack = ack.lock().unwrap(); ack.in_flight.remove(&row_id); ack.deferred.remove(&row_id); - ack.completed.insert(row_id); + if persisted { + ack.persisting.remove(&row_id); + ack.completed.insert(row_id); + } else { + ack.persisting.insert(row_id); + } } persist_cursor(store, ack, channel); } +fn retry_completion_persistence( + store: &Arc>, + ack: &Arc>, + channel: &str, +) { + let pending = ack + .lock() + .unwrap() + .persisting + .iter() + .copied() + .collect::>(); + for row_id in pending { + let persisted = store.lock().unwrap().mark_row_completed(channel, row_id); + match persisted { + Ok(()) => { + let mut ack = ack.lock().unwrap(); + ack.persisting.remove(&row_id); + ack.completed.insert(row_id); + } + Err(error) => error!("retry {channel} completed row {row_id}: {error:#}"), + } + } +} + fn persist_cursor(store: &Arc>, ack: &Arc>, channel: &str) { let mut ack = ack.lock().unwrap(); let Some(row_id) = ack.next_cursor() else { @@ -1328,7 +1391,9 @@ fn runners(cfg: &Config) -> HashMap { impl AckState { fn is_known(&self, row_id: i64) -> bool { - self.in_flight.contains(&row_id) || self.completed.contains(&row_id) + self.in_flight.contains(&row_id) + || self.persisting.contains(&row_id) + || self.completed.contains(&row_id) } fn next_cursor(&self) -> Option { @@ -1337,6 +1402,7 @@ impl AckState { .first() .into_iter() .chain(self.deferred.first()) + .chain(self.persisting.first()) .copied() .min() .unwrap_or(i64::MAX); diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 70296af..54a7f70 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -543,6 +543,51 @@ fn setup_failure_completion_unblocks_later_completed_rows() { let _ = std::fs::remove_file(path); } +#[test] +fn completed_row_marker_failure_keeps_a_retryable_cursor_barrier() { + let path = temp_state_path(); + let store = Arc::new(Mutex::new( + Store::open_at(format!("{path}.db"), &path).unwrap(), + )); + store + .lock() + .unwrap() + .fail_next_completed_row_save_for_test(); + let ack = Arc::new(Mutex::new(AckState::default())); + { + let mut ack = ack.lock().unwrap(); + ack.in_flight.insert(10); + ack.completed.insert(11); + } + + complete_row(&store, &ack, "imessage", 10); + + assert_eq!(store.lock().unwrap().last_row(), 0); + assert!(store + .lock() + .unwrap() + .completed_rows_after("imessage", 0) + .unwrap() + .is_empty()); + { + let ack = ack.lock().unwrap(); + assert!(!ack.in_flight.contains(&10)); + assert!(ack.persisting.contains(&10)); + assert!(!ack.completed.contains(&10)); + assert!(ack.completed.contains(&11)); + } + + retry_completion_persistence(&store, &ack, "imessage"); + persist_cursor(&store, &ack, "imessage"); + + assert_eq!(store.lock().unwrap().last_row(), 11); + let ack = ack.lock().unwrap(); + assert!(ack.persisting.is_empty()); + assert!(ack.completed.is_empty()); + + let _ = std::fs::remove_file(path); +} + #[tokio::test] async fn session_lookup_failure_completes_in_flight_row() { let state_path = temp_state_path(); @@ -2143,6 +2188,65 @@ async fn imessage_pending_filename_defers_only_accepted_rows_and_cannot_block_fo let _ = std::fs::remove_dir_all(assistant_dir); } +#[tokio::test(flavor = "current_thread")] +async fn imessage_completed_row_after_deferred_barrier_is_not_rerun_after_restart() { + let state_path = temp_state_path(); + let sessions_dir = temp_path("imessage-deferred-restart-sessions"); + let assistant_dir = temp_path("imessage-deferred-restart-assistant"); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let cfg = test_config( + &state_path, + sessions_dir.to_str().unwrap(), + assistant_dir.to_str().unwrap(), + ); + let calls = Arc::new(Mutex::new(Vec::::new())); + + let mut pending = message(1, "+15551234567", "+15551234567", false, "photo"); + pending.images.push(InboundImage { + locator: String::new(), + file_size: Some(24), + mime_type: Some("image/heic".to_string()), + data: None, + }); + let later = message(2, "+15551234567", "+15551234567", false, "later"); + let batch = vec![pending, later]; + + let mut first = Gateway::new(cfg.clone()).unwrap(); + first.ctx.runners = Arc::new(fake_runners(calls.clone())); + run_complete_snapshot(&mut first, batch.clone()).await; + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!(first.store.lock().unwrap().cursor("imessage").unwrap(), 0); + assert_eq!( + first + .store + .lock() + .unwrap() + .completed_rows_after("imessage", 0) + .unwrap(), + vec![2] + ); + drop(first); + + let mut restarted = Gateway::new(cfg).unwrap(); + restarted.ctx.runners = Arc::new(fake_runners(calls.clone())); + run_complete_snapshot(&mut restarted, batch).await; + + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!( + restarted.store.lock().unwrap().cursor("imessage").unwrap(), + 0 + ); + assert!(restarted.ack.lock().unwrap().deferred.contains(&1)); + assert!(restarted.ack.lock().unwrap().completed.contains(&2)); + + let _ = std::fs::remove_file(&state_path); + let _ = std::fs::remove_file(format!("{state_path}.db")); + let _ = std::fs::remove_file(format!("{state_path}.audit.jsonl")); + let _ = std::fs::remove_dir_all(format!("{state_path}.cache")); + let _ = std::fs::remove_dir_all(sessions_dir); + let _ = std::fs::remove_dir_all(assistant_dir); +} + #[tokio::test(flavor = "current_thread")] async fn imessage_deferred_barrier_survives_preworker_failure() { let state_path = temp_state_path(); diff --git a/src/history.rs b/src/history.rs index 3acb0c5..a6e975a 100644 --- a/src/history.rs +++ b/src/history.rs @@ -10,7 +10,7 @@ use crate::approval::{parse_answer, AnswerOrigin, AnswerOutcome, NormalizedAnswe #[cfg(test)] use crate::approval::{DeliveryStatus as ApprovalDeliveryStatus, Question, QuestionState}; -const SCHEMA_VERSION: i64 = 12; +const SCHEMA_VERSION: i64 = 13; const RETIRED_JOB_APPROVAL_ERROR: &str = "job approval was removed; request direct job creation"; const MAX_HISTORY_READ_BYTES: usize = 8 * 1024; const READ_TRUNCATED: &str = "\n[truncated by push while reading history]"; @@ -939,6 +939,18 @@ fn migrate(conn: &Connection) -> Result<()> { )?; conn.execute_batch("PRAGMA user_version = 12;")?; } + if version <= 12 { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS channel_completed_rows ( + channel TEXT NOT NULL, + row_id INTEGER NOT NULL, + completed_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + PRIMARY KEY(channel, row_id), + CHECK(length(trim(channel)) > 0) + ); + PRAGMA user_version = 13;", + )?; + } conn.execute_batch("COMMIT;")?; Ok(()) } diff --git a/src/store.rs b/src/store.rs index dcab2a8..fd19f85 100644 --- a/src/store.rs +++ b/src/store.rs @@ -44,6 +44,8 @@ pub struct Store { #[cfg(test)] cursor_save_failures_remaining: usize, #[cfg(test)] + completed_row_save_failures_remaining: usize, + #[cfg(test)] session_save_failures_remaining: usize, } @@ -93,6 +95,8 @@ impl Store { #[cfg(test)] cursor_save_failures_remaining: 0, #[cfg(test)] + completed_row_save_failures_remaining: 0, + #[cfg(test)] session_save_failures_remaining: 0, }; store.migrate_legacy_state(fail_migration_before_commit)?; @@ -145,25 +149,98 @@ impl Store { self.cursor_save_failures_remaining -= 1; return Err(anyhow::anyhow!("injected cursor save failure")); } + let tx = self + .conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .with_context(|| { + format!( + "begin {channel} cursor transaction in {}", + self.database_path.display() + ) + })?; + insert_monotonic_cursor(&tx, channel, id).with_context(|| { + format!( + "advance {channel} cursor transactionally in {}", + self.database_path.display() + ) + })?; + tx.execute( + "DELETE FROM channel_completed_rows + WHERE channel = ?1 + AND row_id <= COALESCE( + (SELECT cursor FROM channel_cursors WHERE channel = ?1), + 0 + )", + [channel], + ) + .with_context(|| { + format!( + "prune {channel} completed rows in {}", + self.database_path.display() + ) + })?; + tx.commit().with_context(|| { + format!( + "commit {channel} cursor transaction in {}", + self.database_path.display() + ) + })?; + Ok(()) + } + + pub fn mark_row_completed(&mut self, channel: &str, row_id: i64) -> Result<()> { + validate_channel(channel)?; + #[cfg(test)] + if self.completed_row_save_failures_remaining > 0 { + self.completed_row_save_failures_remaining -= 1; + return Err(anyhow::anyhow!("injected completed row save failure")); + } self.conn .execute( - "INSERT INTO channel_cursors (channel, cursor) - VALUES (?1, ?2) - ON CONFLICT(channel) DO UPDATE SET - cursor = excluded.cursor, - updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') - WHERE excluded.cursor > channel_cursors.cursor", - params![channel, id], + "INSERT OR IGNORE INTO channel_completed_rows (channel, row_id) + VALUES (?1, ?2)", + params![channel, row_id], ) .with_context(|| { format!( - "advance {channel} cursor transactionally in {}", + "persist {channel} completed row {row_id} in {}", self.database_path.display() ) })?; Ok(()) } + pub fn completed_rows_after(&self, channel: &str, cursor: i64) -> Result> { + validate_channel(channel)?; + let mut statement = self + .conn + .prepare( + "SELECT row_id FROM channel_completed_rows + WHERE channel = ?1 AND row_id > ?2 + ORDER BY row_id", + ) + .with_context(|| { + format!( + "prepare {channel} completed row read from {}", + self.database_path.display() + ) + })?; + let rows = statement + .query_map(params![channel, cursor], |row| row.get(0)) + .with_context(|| { + format!( + "read {channel} completed rows from {}", + self.database_path.display() + ) + })?; + rows.collect::>>().with_context(|| { + format!( + "decode {channel} completed rows from {}", + self.database_path.display() + ) + }) + } + /// Returns the agent session id for a thread, creating one if needed. The /// second value is true when the backend has not started that session yet. pub fn session_for( @@ -411,6 +488,11 @@ impl Store { self.cursor_save_failures_remaining += 1; } + #[cfg(test)] + pub fn fail_next_completed_row_save_for_test(&mut self) { + self.completed_row_save_failures_remaining += 1; + } + #[cfg(test)] pub fn fail_next_session_save_for_test(&mut self) { self.session_save_failures_remaining += 1; @@ -704,6 +786,42 @@ mod tests { cleanup(&database_path, &state_path); } + #[test] + fn completed_rows_survive_reopen_and_are_pruned_by_cursor() { + let (database_path, state_path) = temp_paths(); + let mut store = open(&database_path, &state_path); + store.mark_row_completed("imessage", 12).unwrap(); + store.mark_row_completed("imessage", 14).unwrap(); + store.mark_row_completed("telegram", 13).unwrap(); + drop(store); + + let mut reopened = open(&database_path, &state_path); + assert_eq!( + reopened.completed_rows_after("imessage", 0).unwrap(), + vec![12, 14] + ); + assert_eq!( + reopened.completed_rows_after("telegram", 0).unwrap(), + vec![13] + ); + + reopened.set_cursor("imessage", 12).unwrap(); + assert_eq!( + reopened.completed_rows_after("imessage", 0).unwrap(), + vec![14] + ); + reopened.set_cursor("imessage", 14).unwrap(); + assert!(reopened + .completed_rows_after("imessage", 0) + .unwrap() + .is_empty()); + assert_eq!( + reopened.completed_rows_after("telegram", 0).unwrap(), + vec![13] + ); + cleanup(&database_path, &state_path); + } + #[test] fn interrupted_migration_rolls_back_and_retries_cleanly() { let (database_path, state_path) = temp_paths(); From f99ef04cedb1e331953ed948204beccd2924a79d Mon Sep 17 00:00:00 2001 From: Owain Lewis Date: Tue, 4 Aug 2026 23:45:38 +0100 Subject: [PATCH 6/6] fix(imessage): persist filename deferral limit --- src/channel.rs | 35 +++++++----- src/gateway/mod.rs | 6 +- src/gateway/tests.rs | 43 ++++++++++++++- src/history.rs | 15 ++++- src/imessage/poller.rs | 96 +------------------------------- src/store.rs | 122 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 205 insertions(+), 112 deletions(-) diff --git a/src/channel.rs b/src/channel.rs index 1370853..4dc8720 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -10,10 +10,12 @@ use crate::config::{ChannelKind, Config}; use crate::image::DownloadedImage; use crate::imessage::{Poller as IMessagePoller, Sender as IMessageSender}; use crate::slack::{parse_message_target, Slack}; +use crate::store::Store; use crate::telegram::Telegram; use crate::voice::AudioClip; pub(crate) const REPLY_MARKER: &str = "\n\n-- sent by push"; +const IMESSAGE_PENDING_FILENAME_POLL_LIMIT: u8 = 3; #[derive(Debug, Clone)] pub struct InboundVoice { @@ -102,7 +104,7 @@ trait ChannelContract { } async fn latest_cursor(&self) -> Result; fn accept(&self, message: &RawMessage) -> Option<(String, String)>; - fn should_defer(&self, _message: &RawMessage) -> Result { + fn should_defer(&self, _message: &RawMessage, _store: &mut Store) -> Result { Ok(false) } fn reject_reason(&self, message: &RawMessage) -> &'static str; @@ -245,11 +247,11 @@ impl Channel { } } - pub fn should_defer(&self, message: &RawMessage) -> Result { + pub fn should_defer(&self, message: &RawMessage, store: &mut Store) -> Result { match self { - Self::IMessage(channel) => ChannelContract::should_defer(channel, message), - Self::Telegram(channel) => ChannelContract::should_defer(channel, message), - Self::Slack(channel) => ChannelContract::should_defer(channel, message), + Self::IMessage(channel) => ChannelContract::should_defer(channel, message, store), + Self::Telegram(channel) => ChannelContract::should_defer(channel, message, store), + Self::Slack(channel) => ChannelContract::should_defer(channel, message, store), } } @@ -456,14 +458,21 @@ impl ChannelContract for IMessageChannel { None } - fn should_defer(&self, message: &RawMessage) -> Result { - self.poller.should_defer_pending_filename( - message.row_id, - message - .images - .iter() - .any(|image| image.locator.trim().is_empty()), - ) + fn should_defer(&self, message: &RawMessage, store: &mut Store) -> Result { + let pending = message + .images + .iter() + .any(|image| image.locator.trim().is_empty()); + if pending { + store.should_defer_pending_filename( + self.id(), + message.row_id, + IMESSAGE_PENDING_FILENAME_POLL_LIMIT, + ) + } else { + store.clear_pending_filename(self.id(), message.row_id)?; + Ok(false) + } } fn reject_reason(&self, message: &RawMessage) -> &'static str { diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 109826f..e02b687 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -672,7 +672,11 @@ impl Gateway { continue; } if let Some((thread, target)) = self.channel.accept(m) { - let deferred = match self.channel.should_defer(m) { + let deferred = { + let mut store = self.store.lock().unwrap(); + self.channel.should_defer(m, &mut store) + }; + let deferred = match deferred { Ok(deferred) => deferred, Err(error) => { error!("[{thread}] attachment readiness check failed: {error:#}"); diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 54a7f70..c7a0bd0 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -2227,9 +2227,9 @@ async fn imessage_completed_row_after_deferred_barrier_is_not_rerun_after_restar ); drop(first); - let mut restarted = Gateway::new(cfg).unwrap(); + let mut restarted = Gateway::new(cfg.clone()).unwrap(); restarted.ctx.runners = Arc::new(fake_runners(calls.clone())); - run_complete_snapshot(&mut restarted, batch).await; + run_complete_snapshot(&mut restarted, batch.clone()).await; assert_eq!(calls.lock().unwrap().len(), 1); assert_eq!( @@ -2238,6 +2238,45 @@ async fn imessage_completed_row_after_deferred_barrier_is_not_rerun_after_restar ); assert!(restarted.ack.lock().unwrap().deferred.contains(&1)); assert!(restarted.ack.lock().unwrap().completed.contains(&2)); + drop(restarted); + + let mut restarted_again = Gateway::new(cfg.clone()).unwrap(); + restarted_again.ctx.runners = Arc::new(fake_runners(calls.clone())); + run_complete_snapshot(&mut restarted_again, batch.clone()).await; + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!( + restarted_again + .store + .lock() + .unwrap() + .cursor("imessage") + .unwrap(), + 0 + ); + drop(restarted_again); + + let mut final_restart = Gateway::new(cfg).unwrap(); + final_restart.ctx.runners = Arc::new(fake_runners(calls.clone())); + run_complete_snapshot(&mut final_restart, batch).await; + assert_eq!(calls.lock().unwrap().len(), 1); + assert_eq!( + final_restart + .store + .lock() + .unwrap() + .cursor("imessage") + .unwrap(), + 2 + ); + assert!(final_restart + .ctx + .sent_replies + .lock() + .unwrap() + .last() + .unwrap() + .1 + .contains("JPEG, PNG, or WebP")); let _ = std::fs::remove_file(&state_path); let _ = std::fs::remove_file(format!("{state_path}.db")); diff --git a/src/history.rs b/src/history.rs index a6e975a..f69c492 100644 --- a/src/history.rs +++ b/src/history.rs @@ -10,7 +10,7 @@ use crate::approval::{parse_answer, AnswerOrigin, AnswerOutcome, NormalizedAnswe #[cfg(test)] use crate::approval::{DeliveryStatus as ApprovalDeliveryStatus, Question, QuestionState}; -const SCHEMA_VERSION: i64 = 13; +const SCHEMA_VERSION: i64 = 14; const RETIRED_JOB_APPROVAL_ERROR: &str = "job approval was removed; request direct job creation"; const MAX_HISTORY_READ_BYTES: usize = 8 * 1024; const READ_TRUNCATED: &str = "\n[truncated by push while reading history]"; @@ -951,6 +951,19 @@ fn migrate(conn: &Connection) -> Result<()> { PRAGMA user_version = 13;", )?; } + if version <= 13 { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS channel_pending_filename_polls ( + channel TEXT NOT NULL, + row_id INTEGER NOT NULL, + polls INTEGER NOT NULL CHECK(polls > 0), + updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + PRIMARY KEY(channel, row_id), + CHECK(length(trim(channel)) > 0) + ); + PRAGMA user_version = 14;", + )?; + } conn.execute_batch("COMMIT;")?; Ok(()) } diff --git a/src/imessage/poller.rs b/src/imessage/poller.rs index 353d396..c51c9c3 100644 --- a/src/imessage/poller.rs +++ b/src/imessage/poller.rs @@ -1,8 +1,6 @@ //! Reads new messages directly from the macOS Messages SQLite database. -use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::{Arc, Mutex}; use anyhow::{Context, Result}; use rusqlite::{Connection, OpenFlags}; @@ -30,11 +28,8 @@ pub struct Message { pub struct Poller { db_path: String, attachment_root: PathBuf, - pending_filename_polls: Arc>>, } -const PENDING_FILENAME_POLL_LIMIT: u8 = 3; - // Only real text messages: exclude tapbacks/reactions (associated_message_type) // and system rows like group renames or joins (item_type). const SELECT_NEW: &str = "\ @@ -68,7 +63,6 @@ impl Poller { Self { db_path, attachment_root, - pending_filename_polls: Arc::new(Mutex::new(HashMap::new())), } } @@ -81,7 +75,6 @@ impl Poller { /// `attributedBody` when `text` is NULL. This is blocking; call it from a /// blocking context. pub fn poll(&self, since: i64) -> Result> { - self.clear_completed_pending_filenames(since)?; let conn = self.open()?; let mut stmt = conn.prepare(SELECT_NEW)?; let rows = stmt.query_map([since], |row| { @@ -129,44 +122,6 @@ impl Poller { Ok(out) } - pub fn should_defer_pending_filename(&self, row_id: i64, pending: bool) -> Result { - if pending { - self.should_defer_filename(row_id) - } else { - self.clear_pending_filename(row_id)?; - Ok(false) - } - } - - fn should_defer_filename(&self, row_id: i64) -> Result { - let mut pending = self - .pending_filename_polls - .lock() - .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))?; - let polls = pending.entry(row_id).or_default(); - if *polls < PENDING_FILENAME_POLL_LIMIT { - *polls += 1; - return Ok(true); - } - Ok(false) - } - - fn clear_completed_pending_filenames(&self, since: i64) -> Result<()> { - self.pending_filename_polls - .lock() - .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))? - .retain(|row_id, _| *row_id > since); - Ok(()) - } - - fn clear_pending_filename(&self, row_id: i64) -> Result<()> { - self.pending_filename_polls - .lock() - .map_err(|_| anyhow::anyhow!("iMessage pending attachment state lock poisoned"))? - .remove(&row_id); - Ok(()) - } - /// Highest message ROWID in the database, or 0 when empty. Used to skip /// backlog on first run. pub fn max_row_id(&self) -> Result { @@ -330,7 +285,7 @@ mod tests { } #[test] - fn pending_filename_deferral_ends_when_attachment_filename_is_ready() { + fn poll_observes_when_attachment_filename_becomes_ready() { let root = temp_dir("chat-db-pending-attachment"); let path = root.join("chat.db"); let conn = Connection::open(&path).unwrap(); @@ -357,9 +312,6 @@ mod tests { let pending = poller.poll(0).unwrap(); assert_eq!(pending.len(), 2); assert_eq!(pending[0].attachments[0].locator, ""); - assert!(poller - .should_defer_pending_filename(pending[0].row_id, true) - .unwrap()); let conn = Connection::open(&path).unwrap(); conn.execute( @@ -372,56 +324,10 @@ mod tests { assert_eq!(ready.len(), 2); assert_eq!(ready[0].row_id, 1); assert_eq!(ready[0].attachments[0].locator, "missing-photo.heic"); - assert!(!poller - .should_defer_pending_filename(ready[0].row_id, false) - .unwrap()); assert_eq!(ready[1].row_id, 2); let _ = std::fs::remove_dir_all(root); } - #[test] - fn pending_filename_deferral_is_bounded_and_stable_until_cursor_progress() { - let root = temp_dir("chat-db-abandoned-attachment"); - let path = root.join("chat.db"); - let conn = Connection::open(&path).unwrap(); - create_schema(&conn); - insert_handle(&conn, 1, "+15551234567"); - insert_chat(&conn, 1, "+15551234567"); - insert_chat_handle(&conn, 1, 1); - insert_message(&conn, 1, Some("failed photo"), None, 0, 0, 0, 1, 1); - insert_message(&conn, 2, Some("later"), None, 0, 0, 0, 1, 1); - conn.execute( - "INSERT INTO attachment (ROWID, filename, mime_type, total_bytes) - VALUES (1, NULL, 'image/heic', 24)", - [], - ) - .unwrap(); - conn.execute( - "INSERT INTO message_attachment_join (message_id, attachment_id) VALUES (1, 1)", - [], - ) - .unwrap(); - drop(conn); - - let poller = Poller::new(path.to_string_lossy().to_string()); - let messages = poller.poll(0).unwrap(); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0].attachments[0].locator, ""); - for _ in 0..PENDING_FILENAME_POLL_LIMIT { - assert!(poller - .should_defer_pending_filename(messages[0].row_id, true) - .unwrap()); - } - assert!(!poller - .should_defer_pending_filename(messages[0].row_id, true) - .unwrap()); - assert!(!poller - .should_defer_pending_filename(messages[0].row_id, true) - .unwrap()); - assert_eq!(poller.poll(1).unwrap().len(), 1); - let _ = std::fs::remove_dir_all(root); - } - fn create_schema(conn: &Connection) { conn.execute_batch( " diff --git a/src/store.rs b/src/store.rs index fd19f85..18ee7bf 100644 --- a/src/store.rs +++ b/src/store.rs @@ -179,6 +179,21 @@ impl Store { self.database_path.display() ) })?; + tx.execute( + "DELETE FROM channel_pending_filename_polls + WHERE channel = ?1 + AND row_id <= COALESCE( + (SELECT cursor FROM channel_cursors WHERE channel = ?1), + 0 + )", + [channel], + ) + .with_context(|| { + format!( + "prune {channel} pending filename rows in {}", + self.database_path.display() + ) + })?; tx.commit().with_context(|| { format!( "commit {channel} cursor transaction in {}", @@ -241,6 +256,76 @@ impl Store { }) } + pub fn should_defer_pending_filename( + &mut self, + channel: &str, + row_id: i64, + poll_limit: u8, + ) -> Result { + validate_channel(channel)?; + let exhausted = i64::from(poll_limit) + 1; + let tx = self + .conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .with_context(|| { + format!( + "begin {channel} pending filename transaction in {}", + self.database_path.display() + ) + })?; + tx.execute( + "INSERT INTO channel_pending_filename_polls (channel, row_id, polls) + VALUES (?1, ?2, 1) + ON CONFLICT(channel, row_id) DO UPDATE SET + polls = MIN(channel_pending_filename_polls.polls + 1, ?3), + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')", + params![channel, row_id, exhausted], + ) + .with_context(|| { + format!( + "advance {channel} pending filename row {row_id} in {}", + self.database_path.display() + ) + })?; + let polls: i64 = tx + .query_row( + "SELECT polls FROM channel_pending_filename_polls + WHERE channel = ?1 AND row_id = ?2", + params![channel, row_id], + |row| row.get(0), + ) + .with_context(|| { + format!( + "read {channel} pending filename row {row_id} from {}", + self.database_path.display() + ) + })?; + tx.commit().with_context(|| { + format!( + "commit {channel} pending filename transaction in {}", + self.database_path.display() + ) + })?; + Ok(polls <= i64::from(poll_limit)) + } + + pub fn clear_pending_filename(&mut self, channel: &str, row_id: i64) -> Result<()> { + validate_channel(channel)?; + self.conn + .execute( + "DELETE FROM channel_pending_filename_polls + WHERE channel = ?1 AND row_id = ?2", + params![channel, row_id], + ) + .with_context(|| { + format!( + "clear {channel} pending filename row {row_id} from {}", + self.database_path.display() + ) + })?; + Ok(()) + } + /// Returns the agent session id for a thread, creating one if needed. The /// second value is true when the backend has not started that session yet. pub fn session_for( @@ -822,6 +907,43 @@ mod tests { cleanup(&database_path, &state_path); } + #[test] + fn pending_filename_grace_survives_reopen_and_clears() { + let (database_path, state_path) = temp_paths(); + let mut store = open(&database_path, &state_path); + assert!(store + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + drop(store); + + let mut reopened = open(&database_path, &state_path); + assert!(reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + drop(reopened); + + let mut reopened = open(&database_path, &state_path); + assert!(reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + assert!(!reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + assert!(!reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + + reopened.clear_pending_filename("imessage", 12).unwrap(); + assert!(reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + reopened.set_cursor("imessage", 12).unwrap(); + assert!(reopened + .should_defer_pending_filename("imessage", 12, 3) + .unwrap()); + cleanup(&database_path, &state_path); + } + #[test] fn interrupted_migration_rolls_back_and_retries_cleanly() { let (database_path, state_path) = temp_paths();