diff --git a/apps/maple-agent/Cargo.lock b/apps/maple-agent/Cargo.lock index f9690db9f..46f7bba71 100644 --- a/apps/maple-agent/Cargo.lock +++ b/apps/maple-agent/Cargo.lock @@ -10106,7 +10106,6 @@ name = "store-bench" version = "0.0.7" dependencies = [ "anyhow", - "base64 0.22.1", "maple-agent", "pi-ai", "pi-coding-agent", diff --git a/apps/maple-agent/README.md b/apps/maple-agent/README.md index 9f5a9be96..232c88b5a 100644 --- a/apps/maple-agent/README.md +++ b/apps/maple-agent/README.md @@ -271,7 +271,7 @@ the same selection state. The preview includes `j/k`, `gg/G`, counts, `Enter`, `h/l`, `y`, `/`, `ga`, `[a`/`]a`, `gi`, and `Ctrl-W h/j/k/l`. `Ctrl-W j` from the composer reaches its button row, where `h/l` (or `j/k`) and `gg/G` move between the model, thinking, -integrations, web, image, and dictation controls and `Enter` presses one; a menu +integrations, web, attachment, and dictation controls and `Enter` presses one; a menu opens on its current choice. `Space s n` starts a task and `Space ,` opens Settings. The Settings page itself is navigable, including its shortcut search and editable binding rows. Annotation motions are registered and report diff --git a/apps/maple-agent/app/src/backend.rs b/apps/maple-agent/app/src/backend.rs index 717407fa6..10c6c78d7 100644 --- a/apps/maple-agent/app/src/backend.rs +++ b/apps/maple-agent/app/src/backend.rs @@ -6,8 +6,12 @@ #[cfg_attr(not(feature = "desktop"), allow(unused_imports))] pub use local::*; +pub(crate) mod attachments; pub(crate) mod local; pub(crate) mod paging; +#[cfg(any(feature = "desktop", test))] +pub use attachments::AttachmentUploadRequest; +pub use attachments::ImportAttachmentRequest; #[cfg(feature = "desktop")] pub use paging::*; mod mcp; @@ -558,6 +562,9 @@ async fn dispatch_special( backend: &LocalBackend, request: Request, ) -> Result { + if attachments::handles(&request.method) { + return attachments::dispatch(backend, request).await; + } if paging::is_paging(&request.method) { return paging::dispatch(backend, request, maple_agent::store::PageTransport::Local).await; } @@ -670,7 +677,10 @@ pub(crate) async fn dispatch( // validation must not prevent local history, configuration, or stop calls. // The restore mutex still serializes credential validation/publication. // Cancelling an OAuth wait must remain independent. - if request.method == "cancel_oauth" { + if matches!( + request.method.as_str(), + "cancel_oauth" | "import_attachment" + ) { return dispatch_inner(backend, request).await; } if matches!( diff --git a/apps/maple-agent/app/src/backend/attachments.rs b/apps/maple-agent/app/src/backend/attachments.rs new file mode 100644 index 000000000..ce812c727 --- /dev/null +++ b/apps/maple-agent/app/src/backend/attachments.rs @@ -0,0 +1,358 @@ +//! Controller-local imports and execution-host upload receipts. Only the local +//! IPC endpoint accepts a filesystem path; remote upload metadata never does. +use super::{AgentBackend, LocalBackend}; +use crate::host::protocol::Request; +use maple_agent::store::{AccountId, Store}; +use maple_agent::store::{AttachmentUploadMetadata, AttachmentUploadStart}; +pub use maple_agent::store::{AttachmentUploadStatus, StoredAttachment}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::{ + io::{Read, Seek, SeekFrom}, + sync::Arc, +}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ImportAttachmentRequest { + pub session_id: String, + pub operation_id: String, + pub path: String, + pub name: String, + pub mime_type: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AttachmentUploadRequest { + pub session_id: String, + pub operation_id: String, +} + +#[derive(Debug)] +pub(crate) struct PreparedImport { + pub file: std::fs::File, + pub metadata: AttachmentUploadMetadata, +} + +/// Hash and transfer the same open regular file. A changed source cannot be +/// silently published: the execution store validates the advertised digest. +pub(crate) async fn prepare(request: &ImportAttachmentRequest) -> Result { + let request = request.clone(); + tokio::task::spawn_blocking(move || { + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + // A caller-controlled FIFO must not trap the import worker before + // we can verify that the opened object is a regular file. + options.custom_flags(libc::O_NONBLOCK); + } + let mut file = options + .open(&request.path) + .map_err(|_| "The staged attachment could not be opened")?; + let attributes = file + .metadata() + .map_err(|_| "Could not inspect the attachment")?; + if !attributes.is_file() + || attributes.len() > maple_agent::store::MAX_ATTACHMENT_BYTES as u64 + { + return Err("Attachments must be regular files of at most 30 MiB".into()); + } + let mut hash = Sha256::new(); + let mut count = 0u64; + let mut buffer = [0u8; 64 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|_| "Could not read the attachment")?; + if read == 0 { + break; + } + count += read as u64; + if count > maple_agent::store::MAX_ATTACHMENT_BYTES as u64 { + return Err("Attachment exceeds the 30 MiB limit".into()); + } + hash.update(&buffer[..read]); + } + if count != attributes.len() { + return Err("Attachment changed while preparing it".into()); + } + file.seek(SeekFrom::Start(0)) + .map_err(|_| "Could not read the attachment")?; + Ok(PreparedImport { + file, + metadata: AttachmentUploadMetadata { + name: request.name, + mime_type: request.mime_type, + size: count, + sha256: format!("{:x}", hash.finalize()), + }, + }) + }) + .await + .map_err(|_| "Attachment worker did not complete".to_string())? +} + +impl AgentBackend { + pub async fn import_attachment( + &self, + user: &str, + request: ImportAttachmentRequest, + ) -> Result { + self.call("import_attachment", (user, request)).await + } + pub async fn attachment_upload_status( + &self, + user: &str, + request: AttachmentUploadRequest, + ) -> Result { + self.call("attachment_upload_status", (user, request)).await + } + pub async fn cancel_attachment_upload( + &self, + user: &str, + request: AttachmentUploadRequest, + ) -> Result<(), String> { + self.call("cancel_attachment_upload", (user, request)).await + } +} + +pub(crate) fn handles(method: &str) -> bool { + matches!( + method, + "import_attachment" | "attachment_upload_status" | "cancel_attachment_upload" + ) +} + +pub(crate) async fn context( + backend: &LocalBackend, + user: &str, +) -> Result<(Arc, AccountId), String> { + backend.authorize_account(user)?; + let store = backend + .store() + .cloned() + .ok_or("Local store is unavailable")?; + let account_store = store.clone(); + let user = user.to_owned(); + let account = blocking(move || { + account_store + .account_id(&user)? + .ok_or(maple_agent::store::StoreError::NotFound) + }) + .await?; + Ok((store, account)) +} + +pub(crate) async fn blocking( + operation: impl FnOnce() -> maple_agent::store::StoreResult + Send + 'static, +) -> Result { + tokio::task::spawn_blocking(operation) + .await + .map_err(|_| "Attachment worker did not complete".to_string())? + .map_err(|error| error.to_string()) +} + +pub(crate) async fn dispatch( + backend: &LocalBackend, + request: Request, +) -> Result { + if request.method == "import_attachment" { + let (user, request): (String, ImportAttachmentRequest) = + serde_json::from_value(request.params) + .map_err(|_| "Invalid attachment import parameters")?; + backend.authorize_account(&user)?; + let prepared = prepare(&request).await?; + let (store, start) = { + let account_guard = backend.account_gate.clone().read_owned().await; + let (store, account) = context(backend, &user).await?; + let target = store.clone(); + let session = request.session_id; + let operation = request.operation_id; + let metadata = prepared.metadata; + let start = blocking(move || { + // A cancelled request must not release account ownership while + // its blocking begin/crash-reconciliation still runs. + let _account = account_guard; + target.begin_attachment_upload(account, &session, &operation, &metadata) + }) + .await?; + (store, start) + }; + let attachment = match start { + AttachmentUploadStart::Complete(attachment) => attachment, + AttachmentUploadStart::Writing(mut writer) => { + let mut file = prepared.file; + // No account or SQL lock spans the bulk copy. + let writer = tokio::task::spawn_blocking(move || { + let mut buffer = [0u8; 64 * 1024]; + loop { + let n = file + .read(&mut buffer) + .map_err(|_| "Could not read staged attachment".to_string())?; + if n == 0 { + break; + } + writer + .write_chunk(&buffer[..n]) + .map_err(|e| e.to_string())?; + } + Ok::<_, String>(writer) + }) + .await + .map_err(|_| "Attachment worker did not complete")??; + let prepare_store = store.clone(); + let prepared = blocking(move || writer.prepare(&prepare_store)).await?; + let account_guard = backend.account_gate.clone().read_owned().await; + backend.authorize_account(&user)?; + blocking(move || { + let _account = account_guard; + prepared.publish(&store) + }) + .await? + } + }; + return serde_json::to_value(attachment).map_err(|_| "Could not encode attachment".into()); + } + let (user, upload): (String, AttachmentUploadRequest) = + serde_json::from_value(request.params).map_err(|_| "Invalid attachment operation")?; + // Remote dispatch already holds a read-side account guard. The caller of + // local dispatch owns this short guard, never across an upload. + let (store, account) = context(backend, &user).await?; + let value = match request.method.as_str() { + "attachment_upload_status" => serde_json::to_value( + blocking(move || { + store.attachment_upload_status(account, &upload.session_id, &upload.operation_id) + }) + .await?, + ), + "cancel_attachment_upload" => { + blocking(move || { + store.cancel_attachment_upload(account, &upload.session_id, &upload.operation_id) + }) + .await?; + Ok(serde_json::Value::Null) + } + _ => return Err("Unknown attachment operation".into()), + }; + value.map_err(|_| "Could not encode attachment operation".into()) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use maple_agent::store::{NewTask, TaskKind}; + + pub(crate) fn fixture(root: &std::path::Path) -> LocalBackend { + let mut backend = LocalBackend::isolated_for_test(root); + backend.ui_fixture = true; + let store = backend.store().unwrap(); + let account = store.upsert_account("account", 1).unwrap(); + store + .create_task( + account, + &NewTask { + uuid: "01234567-89ab-7def-8123-456789abcdef".into(), + project_root: root.to_string_lossy().into_owned(), + kind: TaskKind::Desktop, + title: "Upload".into(), + provisional: true, + model: None, + created_ms: 1, + }, + ) + .unwrap(); + backend + } + + #[tokio::test] + async fn local_import_publishes_a_stable_copy_and_keeps_only_ids_for_send() { + let root = tempfile::tempdir().unwrap(); + let backend = fixture(root.path()); + let source = root.path().join("notes.txt"); + std::fs::write(&source, b"stable upload").unwrap(); + let import = ImportAttachmentRequest { + session_id: "01234567-89ab-7def-8123-456789abcdef".into(), + operation_id: "11111111-1111-4111-8111-111111111111".into(), + path: source.to_string_lossy().into_owned(), + name: "notes.txt".into(), + mime_type: "text/plain".into(), + }; + let request = || Request { + id: 1, + method: "import_attachment".into(), + send_id: None, + params: serde_json::json!(["account", import]), + }; + let first: StoredAttachment = + serde_json::from_value(super::super::dispatch(&backend, request()).await.unwrap()) + .unwrap(); + let second: StoredAttachment = + serde_json::from_value(super::super::dispatch(&backend, request()).await.unwrap()) + .unwrap(); + assert_eq!(first.id, second.id); + std::fs::remove_file(&source).unwrap(); + assert_eq!( + std::fs::read(&first.execution_path).unwrap(), + b"stable upload" + ); + let status = dispatch( + &backend, + Request { + id: 2, + method: "attachment_upload_status".into(), + send_id: None, + params: serde_json::json!([ + "account", + AttachmentUploadRequest { + session_id: import.session_id, + operation_id: import.operation_id + } + ]), + }, + ) + .await + .unwrap(); + assert_eq!( + serde_json::from_value::(status).unwrap(), + AttachmentUploadStatus::Complete(first) + ); + } + + #[tokio::test] + async fn local_import_refuses_directories_oversized_files_and_wrong_account_before_reading() { + let root = tempfile::tempdir().unwrap(); + let mut backend = fixture(root.path()); + let path = root.path().join("large"); + std::fs::File::create(&path) + .unwrap() + .set_len(maple_agent::store::MAX_ATTACHMENT_BYTES as u64 + 1) + .unwrap(); + let mut import = ImportAttachmentRequest { + session_id: "01234567-89ab-7def-8123-456789abcdef".into(), + operation_id: "11111111-1111-4111-8111-111111111111".into(), + path: path.to_string_lossy().into_owned(), + name: "large".into(), + mime_type: "application/octet-stream".into(), + }; + assert!(prepare(&import).await.unwrap_err().contains("30 MiB")); + import.path = root.path().to_string_lossy().into_owned(); + assert!(prepare(&import).await.is_err()); + backend.ui_fixture = false; + import.path = "/missing-do-not-open".into(); + let error = dispatch( + &backend, + Request { + id: 1, + method: "import_attachment".into(), + send_id: None, + params: serde_json::json!(["other-account", import]), + }, + ) + .await + .unwrap_err(); + assert!(error.contains("account changed")); + } +} diff --git a/apps/maple-agent/app/src/backend/local.rs b/apps/maple-agent/app/src/backend/local.rs index a1fcb2fa7..6b56ffd91 100644 --- a/apps/maple-agent/app/src/backend/local.rs +++ b/apps/maple-agent/app/src/backend/local.rs @@ -112,7 +112,7 @@ pub struct LocalBackend { row_reader: Option>, #[cfg(test)] pub(super) ui_fixture: bool, - pub(crate) account_gate: tokio::sync::RwLock<()>, + pub(crate) account_gate: Arc>, restore_lock: tokio::sync::Mutex<()>, ui_settings_gate: tokio::sync::Mutex<()>, service: MapleAgentService, @@ -1145,7 +1145,7 @@ impl LocalBackend { ui_fixture: false, restore_lock: tokio::sync::Mutex::new(()), ui_settings_gate: tokio::sync::Mutex::new(()), - account_gate: tokio::sync::RwLock::new(()), + account_gate: Arc::new(tokio::sync::RwLock::new(())), service, auth: MapleApiAuthState::new(), api_url, diff --git a/apps/maple-agent/app/src/backend/paging.rs b/apps/maple-agent/app/src/backend/paging.rs index 78537409c..51eb20c1e 100644 --- a/apps/maple-agent/app/src/backend/paging.rs +++ b/apps/maple-agent/app/src/backend/paging.rs @@ -655,6 +655,7 @@ async fn open_task( if row_limit == 0 || row_limit > 500 { return Err("Invalid open-task row limit".into()); } + handle.ensure_task_history_ready(&request.task).await?; // Start with the whole budget so immutable metadata (notably canonical // roots) need not fit an arbitrary fraction. If the merged wire response // exceeds the budget, reduce the next persisted page by the measured excess. diff --git a/apps/maple-agent/app/src/desktop.rs b/apps/maple-agent/app/src/desktop.rs index 425a86425..0bcbe3e9e 100644 --- a/apps/maple-agent/app/src/desktop.rs +++ b/apps/maple-agent/app/src/desktop.rs @@ -75,12 +75,21 @@ impl RemoteComputer { && match &self.attempt { ConnectionAttempt::Idle => true, ConnectionAttempt::Connecting(_) => false, - ConnectionAttempt::Failed { retry_at, .. } => now >= *retry_at, + ConnectionAttempt::Failed { message, retry_at } => { + !crate::remote::wire::is_protocol_incompatibility_refusal(message) + && now >= *retry_at + } } } fn error(&self) -> Option { match &self.attempt { - ConnectionAttempt::Failed { message, .. } => Some(message.clone()), + ConnectionAttempt::Failed { message, .. } => Some( + if crate::remote::wire::is_protocol_incompatibility_refusal(message) { + "Update required — select to retry after updating".into() + } else { + message.clone() + }, + ), _ => None, } } @@ -1503,6 +1512,8 @@ impl MapleApp { let Some(computer) = self.remote_computers.get_mut(&peer) else { return; }; + // Explicit navigation may retry after either computer was upgraded. + // Only background discovery uses should_connect's incompatibility gate. if matches!(computer.attempt, ConnectionAttempt::Connecting(_)) { return; } @@ -3456,6 +3467,88 @@ mod tests { assert!(other.should_connect(now)); } + #[test] + fn cold_catalog_protocol_refusal_stops_only_automatic_mismatch_retries() { + use crate::remote::wire::{Error, Refusal}; + + let now = std::time::Instant::now(); + for (message, mismatch) in [ + ( + Refusal::terminal(Error::Incompatible.to_string()).into_error(), + true, + ), + (Error::Disconnected.to_string(), false), + ( + Refusal::terminal("Remote control grant is missing or revoked").into_error(), + false, + ), + ( + Refusal::terminal(Error::InvalidFrame.to_string()).into_error(), + false, + ), + ] { + let computer = super::RemoteComputer { + name: "Office".into(), + grant: [7; 16], + chat: None, + attempt: super::ConnectionAttempt::Failed { + message: message.clone(), + retry_at: now + std::time::Duration::from_secs(30), + }, + }; + assert!(!computer.should_connect(now)); + assert_eq!( + computer.should_connect(now + std::time::Duration::from_secs(3600)), + !mismatch, + "{message}" + ); + assert_eq!( + computer.error().as_deref(), + Some(if mismatch { + "Update required — select to retry after updating" + } else { + &message + }) + ); + } + } + + #[gpui::test] + fn cold_catalog_protocol_refusal_allows_explicit_navigation_retry( + cx: &mut gpui::TestAppContext, + ) { + let app = project_shell(cx); + app.update(cx, |app, cx| { + let computer = app.remote_computers.get_mut("east").unwrap(); + computer.chat = None; + computer.attempt = ConnectionAttempt::Failed { + message: crate::remote::wire::Refusal::terminal( + crate::remote::wire::Error::Incompatible.to_string(), + ) + .into_error(), + retry_at: std::time::Instant::now(), + }; + assert!(!computer.should_connect(std::time::Instant::now())); + // Exercise dispatch, without allowing this fixture to connect to + // a developer's host: the backend rejects before any IPC or spawn. + app.backend.retire_after_update(); + app.navigate( + Navigate { + computer: Some("east".into()), + destination: Destination::Overview, + }, + cx, + ); + assert_eq!(app.next_connection_attempt, 1); + assert!(matches!( + app.remote_computers["east"].attempt, + ConnectionAttempt::Connecting(1) + )); + assert!(app.requested_navigation.is_some()); + assert_eq!(app.remote_computers["east"].grant, [1; 16]); + }); + } + #[test] fn password_and_incomplete_setup_remain_resident_when_window_closes() { use maple_agent::store::StoreState; diff --git a/apps/maple-agent/app/src/host/protocol.rs b/apps/maple-agent/app/src/host/protocol.rs index 8c55c9393..09e237a45 100644 --- a/apps/maple-agent/app/src/host/protocol.rs +++ b/apps/maple-agent/app/src/host/protocol.rs @@ -6,8 +6,9 @@ use zeroize::{Zeroize, Zeroizing}; pub const VERSION: u32 = 3; pub const BUILD: &str = concat!(env!("CARGO_PKG_VERSION"), ":", env!("MAPLE_GIT_SOURCE_SHA")); -// Ten 10 MiB image attachments fit after base64/JSON encoding. Reads check -// length before allocating. A connection admits only a bounded number of calls. +// Bounded local history/control replies retain their existing frame budget. +// Attachments are imported by staged local path, not sent as inline JSON bytes. +// Reads check length before allocating; connections bound admitted calls. pub const MAX_FRAME: usize = 160 * 1024 * 1024; #[derive(Serialize, Deserialize)] diff --git a/apps/maple-agent/app/src/remote/control.rs b/apps/maple-agent/app/src/remote/control.rs index 356307851..26d974ac1 100644 --- a/apps/maple-agent/app/src/remote/control.rs +++ b/apps/maple-agent/app/src/remote/control.rs @@ -71,7 +71,7 @@ pub(super) fn listen(node: Arc) { } wire::Mode::ConfirmGrant { window } => { let answer = { - let store = node.store.lock().unwrap(); + let store = node.store.read().unwrap(); if store.peers().get(&peer.to_string()).is_some_and(|p| p.controls_this_computer) && store.peers().get(&peer.to_string()).is_some_and(|p| p.grant == window) { Ok(window) } else { Err(wire::Refusal::terminal("No current grant matches this pairing")) } @@ -82,7 +82,7 @@ pub(super) fn listen(node: Arc) { Ok(None) } wire::Mode::Control => { - if !node.store.lock().unwrap().permits(peer) { return Err(wire::Refusal::terminal("Remote control grant is missing or revoked")); } + if !node.store.read().unwrap().permits(peer) { return Err(wire::Refusal::terminal("Remote control grant is missing or revoked")); } let slot = node.controllers.clone().try_acquire_owned().map_err(|_| "Remote connection limit reached")?; wire::write(&mut tx, &Ok::<(), wire::Refusal>(()), wire::MAX_PREFACE).await.map_err(|e| e.to_string())?; Ok(Some(slot)) @@ -125,10 +125,26 @@ async fn serve( let mut revoked = node.revoked.subscribe(); let capacity = Arc::new(streams::Capacity::new()); let subscription = Arc::new(RwLock::new(Subscription::default())); - let mut calls = tokio::task::JoinSet::new(); + let event_filter = subscription.clone(); + let mut requests = tokio::task::JoinSet::new(); + let request_node = node.clone(); + let request_connection = connection.clone(); + requests.spawn(async move { + streams::serve_requests(request_connection, move |tx, rx| { + let node = request_node.clone(); + let capacity = capacity.clone(); + let subscription = subscription.clone(); + async move { + if let Err(error) = serve_request(node, peer, tx, rx, capacity, subscription).await + { + log::debug!("Remote request stream ended: {error}"); + } + } + }) + .await; + }); // A slow event consumer cannot hold the request accept/dispatch loop. let event_node = node.clone(); - let event_filter = subscription.clone(); let _ = tx.set_priority(1); let mut event_writer = tokio::task::JoinSet::new(); event_writer.spawn(async move { @@ -143,7 +159,7 @@ async fn serve( Err(tokio::sync::broadcast::error::RecvError::Closed) => return Ok(()), }; event_node.ensure_valid().await?; - if !event_node.store.lock().unwrap().permits(peer) + if !event_node.store.read().unwrap().permits(peer) || matches!(&*event, AgentServiceEvent::AccountChanged(_)) { return Ok::<(), String>(()); @@ -167,35 +183,19 @@ async fn serve( _ = connection.closed() => break, _ = event_writer.join_next() => break, _ = tick.tick() => { - if !node.valid().await || !node.store.lock().unwrap().permits(peer) { break; } + if !node.valid().await || !node.store.read().unwrap().permits(peer) { break; } } event = revoked.recv() => { if event.is_err() || event.ok().is_some_and(|key| key == peer || key == node.endpoint.id()) { break; } } - _ = calls.join_next(), if !calls.is_empty() => {}, - incoming = connection.accept_bi() => { - let Ok((mut tx, mut rx)) = incoming else { break; }; - if calls.len() >= streams::MAX_REQUESTS { - let _ = tx.reset(1u32.into()); - let _ = rx.stop(1u32.into()); - continue; - } - let node = node.clone(); - let capacity = capacity.clone(); - let subscription = subscription.clone(); - calls.spawn(async move { - if let Err(error) = serve_request(node, peer, tx, rx, capacity, subscription).await { - log::debug!("Remote request stream ended: {error}"); - } - }); - } + _ = requests.join_next() => break, } } event_writer.abort_all(); connection.close(0u32.into(), b"remote stream closed"); // Closing streams stops incomplete transfers. An already admitted mutation // still finishes on B, as with local IPC; A never resends it. - while calls.join_next().await.is_some() {} + while requests.join_next().await.is_some() {} } async fn serve_request( @@ -206,7 +206,7 @@ async fn serve_request( capacity: Arc, subscription: Arc>, ) -> Result<(), String> { - let header = match streams::Header::read(&mut rx).await { + let header = match capacity.read_header(&mut rx).await { Ok(header) => header, Err(reason) => return Err(reason), }; @@ -215,6 +215,9 @@ async fn serve_request( let _slot = match capacity.acquire(header.method) { Ok(slot) => slot, Err(error) => { + if header.method == wire::Method::UploadAttachment { + return super::uploads::refuse(&mut tx, error).await; + } return streams::reply( &mut tx, &Message::Response { @@ -227,6 +230,16 @@ async fn serve_request( }; let id = header.id; let method = header.method; + if method == wire::Method::UploadAttachment { + if header.send_id.is_some() { + return super::uploads::refuse( + &mut tx, + "Attachment uploads cannot carry a prompt receipt".into(), + ) + .await; + } + return super::uploads::serve(node, peer, id, &mut tx, &mut rx).await; + } let request = match header.body(&mut rx).await { Ok(request) => request, Err(reason) => { @@ -242,7 +255,7 @@ async fn serve_request( }; // Recheck after the potentially slow upload, before any backend effect. node.ensure_valid().await?; - if !node.store.lock().unwrap().permits(peer) { + if !node.store.read().unwrap().permits(peer) { return Err("Remote access revoked".into()); } if request.method == wire::Method::RemoteUpdateInstall { @@ -289,7 +302,7 @@ async fn serve_request( ) .await }; - if node.valid().await && node.store.lock().unwrap().permits(peer) { + if node.valid().await && node.store.read().unwrap().permits(peer) { let result = streams::encode_result(method, result); streams::reply_with_limit( &mut tx, @@ -499,7 +512,7 @@ impl Manager { let node = self.active().await?; if !node .store - .lock() + .read() .unwrap() .peers() .get(&peer.to_string()) @@ -544,7 +557,7 @@ pub async fn bridge( let folder_grant = link .node .store - .lock() + .read() .map_err(|_| "Remote grant unavailable")? .peers() .get(&link.peer.to_string()) @@ -619,7 +632,14 @@ pub async fn bridge( local_send(&mut local_tx, &Message::Response { id: request.id, result: Ok(serde_json::Value::Null) }).await?; return Ok(true); } - let selected = if request.method == "start_runtime" { Ok(wire::Method::TaskBootstrap) } else { wire::Method::parse(&request.method) }; + let selected = match request.method.as_str() { + "start_runtime" => Ok(wire::Method::TaskBootstrap), + "import_attachment" => Ok(wire::Method::UploadAttachment), + // Binary upload is selected only by the local import + // bridge. Remote metadata contains no local source path. + "upload_attachment" => Err("Use the local attachment importer".into()), + _ => wire::Method::parse(&request.method), + }; let method = match selected { Ok(method) => method, Err(error) => { local_send(&mut local_tx, &Message::Response { id: request.id, result: Err(error) }).await?; continue; } @@ -641,13 +661,28 @@ pub async fn bridge( let send_id = request.send_id.clone(); let authorize_folders = async { node.ensure_valid().await?; - if !node.store.lock().map_err(|_| "Remote grant unavailable")?.peers() + if !node.store.read().map_err(|_| "Remote grant unavailable")?.peers() .get(&peer.to_string()) .is_some_and(|current| current.can_connect && current.grant == folder_grant) { return Err("Remote access changed. Reopen this computer.".into()); } Ok(()) }; + if method == wire::Method::UploadAttachment { + let result = async { + let (user, import): (String, crate::backend::ImportAttachmentRequest) = serde_json::from_value(request.params) + .map_err(|_| "Invalid attachment import parameters")?; + if user != node.account { return Err("Remote request account mismatch".into()); } + node.backend.authorize_account(&user)?; + let prepared = crate::backend::attachments::prepare(&import).await?; + authorize_folders.await?; + client.upload(id, user, import, prepared).await + }.await; + return match result { + Ok(frame) => frame, + Err(error) => Message::Response { id, result: Err(error) }, + }; + } match client.call(wire::Command { id, method, send_id: request.send_id, params: request.params }, authorize_folders).await { Ok(frame) => streams::decode_result(method, frame), Err(error) => Message::Response { id, result: Err(match send_id { diff --git a/apps/maple-agent/app/src/remote/manager.rs b/apps/maple-agent/app/src/remote/manager.rs index 55898b3cc..c8beaa132 100644 --- a/apps/maple-agent/app/src/remote/manager.rs +++ b/apps/maple-agent/app/src/remote/manager.rs @@ -15,7 +15,7 @@ use maple_sdk::remote::*; use serde::{Deserialize, Serialize}; use std::{ sync::{ - Arc, Mutex, + Arc, Mutex, RwLock, RwLockWriteGuard, TryLockError, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::Duration, @@ -59,7 +59,9 @@ pub struct Node { pub(super) backend: Arc, pub(super) events: broadcast::Sender>, pub(super) receipts: Arc, - pub(super) store: Mutex, + // Shared permission/publication guards keep controls responsive during disk + // publication; exclusive settings edits still serialize revoke with commit. + pub(super) store: RwLock, pub(super) session: Arc, pub(super) endpoint: Endpoint, addresses: Mutex>, @@ -293,7 +295,7 @@ impl Manager { backend: self.backend.clone(), events: self.events.clone(), receipts: self.receipts.clone(), - store: Mutex::new(store), + store: RwLock::new(store), session, endpoint, addresses: Mutex::new(std::collections::HashMap::new()), @@ -316,7 +318,7 @@ impl Manager { return Err("Remote access was stopped while connecting".into()); } self.resident - .store(node.store.lock().unwrap().enabled(), Ordering::Release); + .store(node.store.read().unwrap().enabled(), Ordering::Release); *current = Some(node.clone()); } super::control::listen(node.clone()); @@ -396,7 +398,7 @@ impl Manager { let link = self.connect(&peer).await?; self.backend.authorize_account(&account)?; if link.node.account != account { return Err("The signed-in account changed".into()); } - let paired = link.node.store.lock().map_err(|_| "Remote grant unavailable")?.peers() + let paired = link.node.store.read().map_err(|_| "Remote grant unavailable")?.peers() .get(&link.peer.to_string()).is_some_and(|value| value.can_connect && value.grant == expected_grant); if !paired { return Err("Remote update authorization changed. Reopen this computer.".into()); } let info = super::streams::peer_info(&link.connection).await @@ -410,7 +412,7 @@ impl Manager { } self.backend.authorize_account(&account)?; link.node.ensure_valid().await?; - if !link.node.store.lock().map_err(|_| "Remote grant unavailable")?.peers() + if !link.node.store.read().map_err(|_| "Remote grant unavailable")?.peers() .get(&link.peer.to_string()).is_some_and(|value| value.can_connect && value.grant == expected_grant) { return Err("Remote update authorization changed. Reopen this computer.".into()); } @@ -487,7 +489,9 @@ impl Manager { self.generation.fetch_add(1, Ordering::SeqCst); let node = current.as_ref().filter(|n| n.account == account).cloned(); if let Some(node) = &node { - node.store.lock().unwrap().set_name(&name)?; + let mut store = write_policy(&node.store).await?; + self.backend.authorize_account(&account)?; + store.set_name(&name)?; } else { self.store(&account)?.set_name(&name)?; } @@ -514,7 +518,7 @@ impl Manager { } else { None }; - // This short local lock orders settings writes against node install. + // This local lock orders settings writes against node install. // No credential validation or network operation holds it. let current = self.node.lock().await; self.generation.fetch_add(1, Ordering::SeqCst); @@ -530,12 +534,13 @@ impl Manager { let mut disk; let mut memory; let store = if let Some(node) = &node { - memory = node.store.lock().unwrap(); + memory = write_policy(&node.store).await?; &mut *memory } else { disk = self.store(&account)?; &mut disk }; + self.backend.authorize_account(&account)?; if let Some(enabled) = enable { store.set_enabled(enabled)?; self.resident.store(enabled, Ordering::Release); @@ -574,7 +579,7 @@ impl Manager { let list = node.session.list_remote_devices().await?; validate_devices(&list.devices)?; node.ensure_valid().await?; - node.remember_peer_names(&list.devices); + node.remember_peer_names(&list.devices).await; serde_json::to_value( list.devices .into_iter() @@ -610,6 +615,29 @@ impl Manager { } } +/// Publication holds a shared policy guard while it makes the attachment +/// durable. Wait for exclusive edits without blocking an async worker or +/// detaching an edit from its caller's account/window/installation guards. +/// The caller must finish its short local edit before awaiting again. +async fn write_policy(store: &RwLock) -> Result, String> { + tokio::time::timeout(super::streams::IO_TIMEOUT, async { + loop { + { + match store.try_write() { + Ok(store) => return Ok(store), + Err(TryLockError::Poisoned(_)) => { + return Err("Remote grant unavailable".into()); + } + Err(TryLockError::WouldBlock) => {} + } + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .map_err(|_| "Remote device settings are busy; try again".to_string())? +} + impl Node { pub(super) async fn valid(&self) -> bool { !self.stopped.is_cancelled() @@ -655,9 +683,9 @@ impl Node { // Local cancel/disable never wait on this network serialization gate. let _registration = self.registration.lock().await; self.ensure_valid().await?; - let key = self.store.lock().unwrap().key(); - let enabled = self.store.lock().unwrap().enabled(); - let name = self.store.lock().unwrap().name().to_owned(); + let key = self.store.read().unwrap().key(); + let enabled = self.store.read().unwrap().enabled(); + let name = self.store.read().unwrap().name().to_owned(); let window = self .window .lock() @@ -700,7 +728,7 @@ impl Node { } async fn begin_pairing(&self) -> Result { let _starting = self.registration.lock().await; - if !self.store.lock().unwrap().enabled() { + if !self.store.read().unwrap().enabled() { return Err("Enable remote access on this computer first".into()); } if let Some(window) = self.current_pairing().await { @@ -708,8 +736,8 @@ impl Node { } let generation = self.pairing_generation.load(Ordering::SeqCst); let next = pairing::Window::new(); - let key = self.store.lock().unwrap().key(); - let name = self.store.lock().unwrap().name().to_owned(); + let key = self.store.read().unwrap().key(); + let name = self.store.read().unwrap().name().to_owned(); let address = self.endpoint.addr(); let registration = register( &self.session, @@ -730,7 +758,7 @@ impl Node { self.ensure_valid().await?; let mut window = self.window.lock().await; if self.pairing_generation.load(Ordering::SeqCst) != generation - || !self.store.lock().unwrap().enabled() + || !self.store.read().unwrap().enabled() { return Err("Pairing was cancelled locally".into()); } @@ -769,7 +797,7 @@ impl Node { } }; validate_devices(&list.devices)?; - self.remember_peer_names(&list.devices); + self.remember_peer_names(&list.devices).await; let device = list .devices .into_iter() @@ -793,18 +821,16 @@ impl Node { } Ok((address, device)) } - fn remember_peer_names(&self, devices: &[RemoteDevice]) { + async fn remember_peer_names(&self, devices: &[RemoteDevice]) { let names = devices .iter() .filter(|d| d.channel == crate::profile::PROFILE) .map(|d| (d.endpoint.as_str(), d.name.as_str())); - if self - .store - .lock() - .unwrap() - .refresh_peer_names(names) - .is_err() - { + let result = match write_policy(&self.store).await { + Ok(mut store) => store.refresh_peer_names(names), + Err(error) => Err(error), + }; + if result.is_err() { log::warn!("Could not save updated remote computer names"); } } @@ -817,7 +843,7 @@ impl Node { && validate_devices(&list.devices).is_ok() && self.valid().await { - self.remember_peer_names(&list.devices); + self.remember_peer_names(&list.devices).await; } } async fn pair(&self, peer: EndpointId, code: pairing::Code) -> Result<(), String> { @@ -867,7 +893,7 @@ impl Node { // dial does not burn a window, but a partial proof can never retry. self.ensure_valid().await?; { - let mut store = self.store.lock().unwrap(); + let mut store = write_policy(&self.store).await?; if self.pairing_generation.load(Ordering::SeqCst) != generation { return Err("Pairing was cancelled locally".into()); } @@ -885,7 +911,7 @@ impl Node { { let _account = self.backend.account_gate.read().await; self.ensure_valid().await?; - let mut store = self.store.lock().unwrap(); + let mut store = write_policy(&self.store).await?; if self.pairing_generation.load(Ordering::SeqCst) != generation { return Err("Pairing was cancelled locally".into()); } @@ -904,7 +930,7 @@ impl Node { { let _account = self.backend.account_gate.read().await; self.ensure_valid().await?; - let mut store = self.store.lock().unwrap(); + let mut store = write_policy(&self.store).await?; if self.pairing_generation.load(Ordering::SeqCst) != generation { return Err("Pairing was cancelled locally".into()); } @@ -918,7 +944,7 @@ impl Node { result => { if self .store - .lock() + .read() .unwrap() .peers() .get(&peer.to_string()) @@ -934,7 +960,7 @@ impl Node { let generation = self.pairing_generation.load(Ordering::SeqCst); let window = self .store - .lock() + .read() .unwrap() .peers() .get(&peer.to_string()) @@ -969,7 +995,7 @@ impl Node { } let _account = self.backend.account_gate.read().await; self.ensure_valid().await?; - let mut store = self.store.lock().unwrap(); + let mut store = write_policy(&self.store).await?; if self.pairing_generation.load(Ordering::SeqCst) != generation { return Err("Pairing was cancelled locally".into()); } @@ -987,7 +1013,7 @@ impl Node { tx: &mut iroh::endpoint::SendStream, rx: &mut iroh::endpoint::RecvStream, ) -> Result<(), String> { - if !self.store.lock().unwrap().enabled() { + if !self.store.read().unwrap().enabled() { return Err("Remote access is disabled".into()); } if !self @@ -1048,7 +1074,13 @@ impl Node { return Err("Pairing window was closed or expired".into()); } { - let mut store = self.store.lock().unwrap(); + let mut store = write_policy(&self.store).await?; + if !pending_window + .as_ref() + .is_some_and(|w| w.id == window && !w.remaining().is_zero()) + { + return Err("Pairing window was closed or expired".into()); + } if !store.enabled() { return Err("Remote access is disabled".into()); } @@ -1247,6 +1279,70 @@ fn validate_devices(devices: &[RemoteDevice]) -> Result<(), String> { mod tests { use super::*; + fn policy_fixture() -> (tempfile::TempDir, Arc>) { + let root = tempfile::tempdir().unwrap(); + let database = Arc::new( + maple_agent::store::Store::open(maple_agent::store::StorePaths::new(root.path())) + .unwrap(), + ); + database.upsert_account("account", 1).unwrap(); + let mut policy = + Store::open(database, "local", "project", "unpackaged", "account").unwrap(); + policy.set_enabled(true).unwrap(); + (root, Arc::new(RwLock::new(policy))) + } + + fn hold_publication_policy( + policy: Arc>, + ) -> (std::sync::mpsc::Sender<()>, std::thread::JoinHandle<()>) { + let (entered, waiting) = std::sync::mpsc::channel(); + let (release, resume) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + let _publication = policy.read().unwrap(); + entered.send(()).unwrap(); + resume.recv().unwrap(); + }); + waiting.recv().unwrap(); + (release, worker) + } + + #[tokio::test(start_paused = true)] + async fn cancelling_a_waiting_policy_edit_cannot_disable_access_later() { + let (_root, policy) = policy_fixture(); + let (release, publication) = hold_publication_policy(policy.clone()); + let editing = policy.clone(); + let (started, waiting) = tokio::sync::oneshot::channel(); + let edit = tokio::spawn(async move { + started.send(()).unwrap(); + write_policy(&editing).await?.set_enabled(false) + }); + waiting.await.unwrap(); + tokio::task::yield_now().await; + assert!(!edit.is_finished()); + assert!(policy.read().unwrap().enabled()); + edit.abort(); + assert!(edit.await.unwrap_err().is_cancelled()); + release.send(()).unwrap(); + publication.join().unwrap(); + tokio::time::advance(super::super::streams::IO_TIMEOUT).await; + assert!(policy.read().unwrap().enabled()); + } + + #[tokio::test(start_paused = true)] + async fn publication_cannot_make_a_policy_edit_wait_indefinitely() { + let (_root, policy) = policy_fixture(); + let (release, publication) = hold_publication_policy(policy.clone()); + let started = tokio::time::Instant::now(); + let result = write_policy(&policy).await; + assert!( + matches!(result, Err(error) if error == "Remote device settings are busy; try again") + ); + assert_eq!(started.elapsed(), super::super::streams::IO_TIMEOUT); + release.send(()).unwrap(); + publication.join().unwrap(); + assert!(policy.read().unwrap().enabled()); + } + #[tokio::test] async fn peer_info_connection_timeout_is_unreachable_not_legacy() { let result = query_peer_information( diff --git a/apps/maple-agent/app/src/remote/mod.rs b/apps/maple-agent/app/src/remote/mod.rs index e51b7bb7a..85baf82e9 100644 --- a/apps/maple-agent/app/src/remote/mod.rs +++ b/apps/maple-agent/app/src/remote/mod.rs @@ -4,6 +4,7 @@ mod pairing; mod store; pub(crate) mod streams; pub(crate) mod update; +mod uploads; pub(crate) mod wire; pub mod control; diff --git a/apps/maple-agent/app/src/remote/streams.rs b/apps/maple-agent/app/src/remote/streams.rs index e11357f15..330388e30 100644 --- a/apps/maple-agent/app/src/remote/streams.rs +++ b/apps/maple-agent/app/src/remote/streams.rs @@ -8,7 +8,21 @@ use std::{sync::Arc, time::Duration}; use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore}; use zeroize::Zeroizing; -pub const MAX_REQUESTS: usize = 40; +pub const MAX_ORDINARY_REQUESTS: usize = 32; +pub const MAX_URGENT_REQUESTS: usize = 2; +pub const MAX_RANGE_REQUESTS: usize = 4; +pub const MAX_BULK_REQUESTS: usize = 2; +pub const MAX_REQUESTS: usize = + MAX_ORDINARY_REQUESTS + MAX_URGENT_REQUESTS + MAX_RANGE_REQUESTS + MAX_BULK_REQUESTS; +// Limit concurrent header readers independently of admitted RPCs. Requests +// queue briefly behind a burst, under one queue-and-parse deadline. A malicious +// peer withholding headers can delay classification only until that deadline. +pub const MAX_PENDING_HEADERS: usize = 8; +pub const MAX_REQUEST_TASKS: usize = MAX_REQUESTS + MAX_PENDING_HEADERS; +const MAX_BIDI_STREAMS: u32 = MAX_REQUEST_TASKS as u32 + 1; +const STREAM_RECEIVE_WINDOW: u32 = 512 * 1024; +const CONNECTION_RECEIVE_WINDOW: u32 = 32 * 1024 * 1024; +const _: () = assert!(CONNECTION_RECEIVE_WINDOW > MAX_BIDI_STREAMS * STREAM_RECEIVE_WINDOW); pub const IO_TIMEOUT: Duration = Duration::from_secs(10); const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); const FOLDERS_UNCONFIRMED: &str = @@ -24,29 +38,77 @@ pub fn transport_config() -> QuicTransportConfig { // One event/handshake stream plus bounded request streams. The aggregate // window exceeds all stream windows, so unread bulk streams cannot use // up the flow-control credit needed by another stream. - .max_concurrent_bidi_streams(41u32.into()) + .max_concurrent_bidi_streams(MAX_BIDI_STREAMS.into()) .max_concurrent_uni_streams(0u32.into()) - .stream_receive_window((512u32 * 1024).into()) - .receive_window((32u32 * 1024 * 1024).into()) + .stream_receive_window(STREAM_RECEIVE_WINDOW.into()) + .receive_window(CONNECTION_RECEIVE_WINDOW.into()) .build() } +/// Bound work before classification too: withheld headers occupy tasks until +/// their queue/read deadline. Closing the connection stops admission, while +/// already admitted effects drain exactly once. +pub(super) async fn serve_requests(connection: Connection, mut serve: F) +where + F: FnMut(SendStream, RecvStream) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, +{ + let mut calls = tokio::task::JoinSet::new(); + loop { + tokio::select! { + _ = calls.join_next(), if !calls.is_empty() => {}, + incoming = connection.accept_bi(), if calls.len() < MAX_REQUEST_TASKS => { + let Ok((tx, rx)) = incoming else { break; }; + calls.spawn(serve(tx, rx)); + } + } + } + while calls.join_next().await.is_some() {} +} + pub struct Capacity { ordinary: Arc, urgent: Arc, ranges: Arc, + bulk: Arc, + headers: Arc, } impl Capacity { pub fn new() -> Self { Self { - ordinary: Arc::new(Semaphore::new(32)), - urgent: Arc::new(Semaphore::new(2)), - ranges: Arc::new(Semaphore::new(4)), + ordinary: Arc::new(Semaphore::new(MAX_ORDINARY_REQUESTS)), + urgent: Arc::new(Semaphore::new(MAX_URGENT_REQUESTS)), + ranges: Arc::new(Semaphore::new(MAX_RANGE_REQUESTS)), + bulk: Arc::new(Semaphore::new(MAX_BULK_REQUESTS)), + headers: Arc::new(Semaphore::new(MAX_PENDING_HEADERS)), } } + pub async fn read_header(&self, rx: &mut RecvStream) -> Result { + self.read_header_with_timeout(rx, IO_TIMEOUT).await + } + async fn read_header_with_timeout( + &self, + rx: &mut RecvStream, + deadline: Duration, + ) -> Result { + tokio::time::timeout(deadline, async { + let _slot = self + .headers + .acquire() + .await + .map_err(|_| "Remote header admission closed")?; + wire::read(rx, wire::MAX_PREFACE) + .await + .map_err(|error| error.to_string()) + }) + .await + .map_err(|_| "Remote request header timed out".to_string())? + } pub fn acquire(&self, method: Method) -> Result { - let pool = if matches!(method, Method::ReadAttachment | Method::ReadRow) { + let pool = if method == Method::UploadAttachment { + &self.bulk + } else if matches!(method, Method::ReadAttachment | Method::ReadRow) { &self.ranges } else if method.urgent() { &self.urgent @@ -77,6 +139,16 @@ impl Client { } } + pub(super) async fn upload( + &self, + id: u64, + user: String, + request: crate::backend::ImportAttachmentRequest, + prepared: crate::backend::attachments::PreparedImport, + ) -> Result { + super::uploads::call(&self.connection, id, user, request, prepared).await + } + pub async fn call( &self, command: Command, @@ -120,6 +192,7 @@ pub struct Header { } impl Header { + #[cfg(test)] pub async fn read(rx: &mut RecvStream) -> Result { tokio::time::timeout(IO_TIMEOUT, wire::read(rx, wire::MAX_PREFACE)) .await @@ -145,7 +218,14 @@ impl Header { } fn body_limit(method: Method) -> usize { - if matches!(method, Method::PeerInfo | Method::RemoteUpdateInstall) { + if matches!( + method, + Method::PeerInfo + | Method::RemoteUpdateInstall + | Method::UploadAttachment + | Method::AttachmentUploadStatus + | Method::CancelAttachmentUpload + ) { wire::MAX_PREFACE } else if method.urgent() { MAX_CONTROL @@ -253,7 +333,7 @@ pub async fn call(connection: &Connection, command: Command) -> Result Result Result<(), String> { +pub(super) async fn finish_reply(rx: &mut RecvStream, deadline: Duration) -> Result<(), String> { let mut trailing = [0u8; 1]; let count = tokio::time::timeout(deadline, tokio::io::AsyncReadExt::read(rx, &mut trailing)) .await @@ -322,7 +402,7 @@ async fn finish_reply(rx: &mut RecvStream, deadline: Duration) -> Result<(), Str } #[cfg(test)] -mod tests; +pub(super) mod tests; #[cfg(test)] mod folder_tests; diff --git a/apps/maple-agent/app/src/remote/streams/tests.rs b/apps/maple-agent/app/src/remote/streams/tests.rs index 93de521f1..2947867ca 100644 --- a/apps/maple-agent/app/src/remote/streams/tests.rs +++ b/apps/maple-agent/app/src/remote/streams/tests.rs @@ -3,7 +3,7 @@ use iroh::{Endpoint, RelayMode, endpoint::presets}; use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; -pub(super) async fn endpoints() -> (Endpoint, Endpoint, Connection, Connection) { +pub(in crate::remote) async fn endpoints() -> (Endpoint, Endpoint, Connection, Connection) { async fn endpoint() -> Endpoint { Endpoint::builder(presets::Minimal) .alpns(vec![wire::ALPN.to_vec()]) @@ -26,7 +26,7 @@ pub(super) async fn endpoints() -> (Endpoint, Endpoint, Connection, Connection) (a, b, client, server) } -pub(super) fn command(id: u64, method: Method) -> Command { +pub(in crate::remote) fn command(id: u64, method: Method) -> Command { Command { id, method, @@ -46,7 +46,7 @@ async fn receive( } #[tokio::test] -async fn legacy_wire_six_rejects_optional_probe_without_losing_ordinary_requests() { +async fn unsupported_optional_probe_does_not_lose_ordinary_requests() { // This decoder is intentionally frozen to the old method set. It must not // acquire the new enum variants from the implementation under test. #[derive(Deserialize)] @@ -580,7 +580,7 @@ async fn ten_mib_image_ranges_and_oversized_page_do_not_break_other_requests() { #[test] fn ranges_have_four_slots_and_do_not_consume_urgent_capacity() { let capacity = Capacity::new(); - let _ranges: Vec<_> = (0..4) + let _ranges: Vec<_> = (0..MAX_RANGE_REQUESTS) .map(|_| capacity.acquire(Method::ReadAttachment).unwrap()) .collect(); assert!(capacity.acquire(Method::ReadAttachment).is_err()); @@ -590,3 +590,365 @@ fn ranges_have_four_slots_and_do_not_consume_urgent_capacity() { assert_eq!(Method::HistoryBefore.priority(), 5); assert_eq!(Method::CancelRun.priority(), 10); } + +#[tokio::test] +async fn header_admission_queues_a_full_legitimate_burst_without_rejecting_urgent_calls() { + let (a, b, client, server) = endpoints().await; + let capacity = Arc::new(Capacity::new()); + let mut callers = tokio::task::JoinSet::new(); + let mut receivers = tokio::task::JoinSet::new(); + // Every pool may send at once. Forty short valid headers must not race the + // eight active reader limit and receive incidental busy/refused errors. + for id in 0..MAX_REQUESTS as u64 { + let client = client.clone(); + callers.spawn(async move { + let method = if id < MAX_ORDINARY_REQUESTS as u64 { + Method::ListTasks + } else if id < (MAX_ORDINARY_REQUESTS + MAX_RANGE_REQUESTS) as u64 { + Method::ReadAttachment + } else if id < (MAX_ORDINARY_REQUESTS + MAX_RANGE_REQUESTS + MAX_BULK_REQUESTS) as u64 { + Method::UploadAttachment + } else { + Method::CancelRun + }; + call(&client, command(id, method)).await.unwrap() + }); + } + for _ in 0..MAX_REQUESTS { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let capacity = capacity.clone(); + receivers.spawn(async move { + // This is the same queued admission and parse path serve_request + // executes, with real independently accepted QUIC streams. + let header = capacity.read_header(&mut rx).await.unwrap(); + let id = header.id; + header.body(&mut rx).await.unwrap(); + reply( + &mut tx, + &Message::Response { + id, + result: Ok(serde_json::Value::Null), + }, + ) + .await + .unwrap(); + }); + } + tokio::time::timeout(Duration::from_secs(5), async { + while let Some(response) = callers.join_next().await { + assert!(matches!( + response.unwrap(), + Message::Response { result: Ok(_), .. } + )); + } + while let Some(response) = receivers.join_next().await { + response.unwrap(); + } + }) + .await + .unwrap(); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn incomplete_headers_expire_and_queued_urgent_header_recovers_without_reset() { + let (a, b, client, server) = endpoints().await; + let capacity = Arc::new(Capacity::new()); + let mut withheld = Vec::new(); + let mut readers = tokio::task::JoinSet::new(); + for _ in 0..MAX_PENDING_HEADERS { + let (mut tx, rx) = client.open_bi().await.unwrap(); + tx.write_u32(100).await.unwrap(); + tx.write_all(b"{").await.unwrap(); + withheld.push((tx, rx)); + let (response, mut body) = server.accept_bi().await.unwrap(); + let capacity = capacity.clone(); + readers.spawn(async move { + let _response = response; + capacity + .read_header_with_timeout(&mut body, Duration::from_secs(3)) + .await + }); + } + tokio::time::timeout(Duration::from_secs(2), async { + while capacity.headers.available_permits() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let control = call(&client, command(100, Method::CancelRun)); + let accept = async { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = capacity + .read_header_with_timeout(&mut rx, Duration::from_secs(5)) + .await + .unwrap(); + assert_eq!(header.method, Method::CancelRun); + let id = header.id; + header.body(&mut rx).await.unwrap(); + reply( + &mut tx, + &Message::Response { + id, + result: Ok(serde_json::Value::Null), + }, + ) + .await + .unwrap(); + }; + let (answer, ()) = tokio::time::timeout(Duration::from_secs(6), async { + tokio::join!(control, accept) + }) + .await + .unwrap(); + assert!(matches!( + answer.unwrap(), + Message::Response { result: Ok(_), .. } + )); + while let Some(result) = readers.join_next().await { + assert!(matches!(result.unwrap(),Err(reason) if reason=="Remote request header timed out")); + } + assert_eq!(capacity.headers.available_permits(), MAX_PENDING_HEADERS); + drop(withheld); + a.close().await; + b.close().await; +} + +#[test] +fn request_budgets_reserve_header_tasks_and_the_event_stream() { + let capacity = Capacity::new(); + let mut slots = Vec::new(); + for (method, limit) in [ + (Method::ListTasks, MAX_ORDINARY_REQUESTS), + (Method::CancelRun, MAX_URGENT_REQUESTS), + (Method::ReadAttachment, MAX_RANGE_REQUESTS), + (Method::UploadAttachment, MAX_BULK_REQUESTS), + ] { + slots.extend((0..limit).map(|_| capacity.acquire(method).unwrap())); + assert!(capacity.acquire(method).is_err()); + } + assert_eq!(slots.len(), MAX_REQUESTS); + assert_eq!(MAX_REQUEST_TASKS, slots.len() + MAX_PENDING_HEADERS); + assert_eq!(MAX_BIDI_STREAMS as usize, MAX_REQUEST_TASKS + 1); + assert_eq!(capacity.headers.available_permits(), MAX_PENDING_HEADERS); +} + +#[tokio::test] +async fn saturated_accept_loop_expires_incomplete_headers_and_serves_both_urgent_calls() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::{Barrier, mpsc}; + + // Exercise both occupied header readers and a completely full task gate. + // The fixture handler holds classified work; this does not exercise Node + // authorization or attachment payload decoding. + for incomplete in [ + MAX_PENDING_HEADERS, + MAX_PENDING_HEADERS + MAX_URGENT_REQUESTS, + ] { + let (a, b, client, server) = endpoints().await; + let stage = std::cell::Cell::new("handshake"); + let exercise = async { + // Control's authenticated handshake/event stream consumes the extra + // transport slot before the production request accept loop starts. + let (mut event_tx, event_rx) = client.open_bi().await.unwrap(); + event_tx.write_all(b"e").await.unwrap(); + let (event_reply, mut event_request) = server.accept_bi().await.unwrap(); + assert_eq!(event_request.read_u8().await.unwrap(), b'e'); + let _event_stream = (event_tx, event_rx, event_reply, event_request); + + let capacity = Arc::new(Capacity::new()); + let release = Arc::new(Semaphore::new(0)); + let accepted = Arc::new(AtomicUsize::new(0)); + let completed = Arc::new(AtomicUsize::new(0)); + let held = MAX_ORDINARY_REQUESTS + MAX_RANGE_REQUESTS + MAX_BULK_REQUESTS; + let header_deadline = Duration::from_secs(3); + let urgent_methods = [Method::CancelRun, Method::AnswerQuestion]; + assert_eq!(urgent_methods.len(), MAX_URGENT_REQUESTS); + let start_headers = Arc::new(Barrier::new(incomplete + 1)); + let (admitted_tx, mut admitted_rx) = mpsc::unbounded_channel(); + let (pending_tx, mut pending_rx) = mpsc::unbounded_channel(); + let (expired_tx, mut expired_rx) = mpsc::unbounded_channel(); + let (finished_tx, mut finished_rx) = mpsc::unbounded_channel(); + let mut receiver = tokio::task::JoinSet::new(); + receiver.spawn(serve_requests(server, { + let capacity = capacity.clone(); + let release = release.clone(); + let accepted = accepted.clone(); + let completed = completed.clone(); + let start_headers = start_headers.clone(); + move |mut tx, mut rx| { + let index = accepted.fetch_add(1, Ordering::SeqCst); + let withheld = (held..held + incomplete).contains(&index); + let capacity = capacity.clone(); + let release = release.clone(); + let completed = completed.clone(); + let start_headers = start_headers.clone(); + let admitted = admitted_tx.clone(); + let pending = pending_tx.clone(); + let expired = expired_tx.clone(); + let finished = finished_tx.clone(); + async move { + let header = { + if withheld { + start_headers.wait().await; + } + let mut read = Box::pin( + capacity.read_header_with_timeout(&mut rx, header_deadline), + ); + // Report readiness only after polling the real read: + // its header permit is held or queued, not merely + // waiting for this task to be scheduled. + if withheld { + let first = std::future::poll_fn(|cx| { + std::task::Poll::Ready(std::future::Future::poll( + read.as_mut(), + cx, + )) + }) + .await; + assert!(first.is_pending()); + pending.send(()).unwrap(); + } + read.await + }; + let header = match header { + Ok(header) => header, + Err(reason) => { + assert!(withheld); + assert_eq!(reason, "Remote request header timed out"); + expired.send(()).unwrap(); + completed.fetch_add(1, Ordering::SeqCst); + finished.send(()).unwrap(); + return; + } + }; + let _slot = capacity.acquire(header.method).unwrap(); + let method = header.method; + let id = header.id; + header.body(&mut rx).await.unwrap(); + admitted.send(method).unwrap(); + if !method.urgent() { + release.acquire().await.unwrap().forget(); + } + reply( + &mut tx, + &Message::Response { + id, + result: Ok(json!(null)), + }, + ) + .await + .unwrap(); + completed.fetch_add(1, Ordering::SeqCst); + finished.send(()).unwrap(); + } + } + })); + + let mut callers = tokio::task::JoinSet::new(); + let mut next_id = 0; + stage.set("class admission"); + for (method, limit) in [ + (Method::ListTasks, MAX_ORDINARY_REQUESTS), + (Method::ReadAttachment, MAX_RANGE_REQUESTS), + (Method::UploadAttachment, MAX_BULK_REQUESTS), + ] { + for _ in 0..limit { + let client = client.clone(); + let id = next_id; + next_id += 1; + callers.spawn(async move { call(&client, command(id, method)).await.unwrap() }); + assert_eq!(admitted_rx.recv().await.unwrap(), method); + } + assert!(capacity.acquire(method).is_err()); + } + let mut withheld_streams = Vec::new(); + stage.set("incomplete-header admission"); + for _ in 0..incomplete { + let (mut tx, rx) = client.open_bi().await.unwrap(); + tx.write_u32(100).await.unwrap(); + tx.write_all(b"{").await.unwrap(); + withheld_streams.push((tx, rx)); + } + start_headers.wait().await; + for _ in 0..incomplete { + pending_rx.recv().await.unwrap(); + } + // Every request uses the same queue-and-read deadline. Issue urgent + // demand on a later timer tick so its absolute deadline follows the + // already pending incomplete headers, even with coarse timers. + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(capacity.headers.available_permits(), 0); + assert_eq!(accepted.load(Ordering::SeqCst), held + incomplete); + assert_eq!(capacity.urgent.available_permits(), MAX_URGENT_REQUESTS); + + let mut urgent = tokio::task::JoinSet::new(); + stage.set("header expiry"); + for method in urgent_methods { + let client = client.clone(); + let id = next_id; + next_id += 1; + urgent.spawn(async move { call(&client, command(id, method)).await.unwrap() }); + } + for _ in 0..incomplete { + expired_rx.recv().await.unwrap(); + } + // The full-cap case also exhausts QUIC stream credit. A rejected + // client's unfinished send half must close before that credit can + // be reused: STOP_SENDING alone does not supply its final offset. + // Honor the server's expired operations while urgent calls wait, + // rather than retaining these handles until urgent completion. + drop(withheld_streams); + stage.set("urgent replies after stream retirement"); + while let Some(response) = urgent.join_next().await { + assert!(matches!( + response.unwrap(), + Message::Response { result: Ok(_), .. } + )); + } + assert!(client.close_reason().is_none()); + stage.set("classified request drain"); + release.add_permits(held); + while let Some(response) = callers.join_next().await { + assert!(matches!( + response.unwrap(), + Message::Response { result: Ok(_), .. } + )); + } + assert!(matches!( + call(&client, command(next_id, Method::CancelRun)) + .await + .unwrap(), + Message::Response { result: Ok(_), .. } + )); + let expected = held + incomplete + urgent_methods.len() + 1; + for _ in 0..expected { + finished_rx.recv().await.unwrap(); + } + client.close(0u32.into(), b"test complete"); + stage.set("accept-loop shutdown"); + receiver.join_next().await.unwrap().unwrap(); + assert!(receiver.is_empty()); + assert_eq!(accepted.load(Ordering::SeqCst), expected); + assert_eq!( + completed.load(Ordering::SeqCst), + accepted.load(Ordering::SeqCst) + ); + assert_eq!(capacity.ordinary.available_permits(), MAX_ORDINARY_REQUESTS); + assert_eq!(capacity.ranges.available_permits(), MAX_RANGE_REQUESTS); + assert_eq!(capacity.bulk.available_permits(), MAX_BULK_REQUESTS); + assert_eq!(capacity.urgent.available_permits(), MAX_URGENT_REQUESTS); + assert_eq!(capacity.headers.available_permits(), MAX_PENDING_HEADERS); + }; + let result = tokio::time::timeout(Duration::from_secs(10), exercise).await; + tokio::join!(a.close(), b.close()); + result.unwrap_or_else(|error| { + panic!( + "{incomplete} incomplete headers at {}: {error}", + stage.get() + ) + }); + } +} diff --git a/apps/maple-agent/app/src/remote/update.rs b/apps/maple-agent/app/src/remote/update.rs index 5bfe4bbe4..ea86592b2 100644 --- a/apps/maple-agent/app/src/remote/update.rs +++ b/apps/maple-agent/app/src/remote/update.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Deserializer, Serialize}; use std::sync::{Arc, Weak}; use uuid::Uuid; -pub const PROTOCOL_VERSION: &str = "6.2.0"; +pub const PROTOCOL_VERSION: &str = "7.0.0"; pub const PEER_INFO_CAPABILITY: &str = "peer_info_v1"; pub const INSTALL_CAPABILITY: &str = "remote_update_install_v1"; pub const PROJECT_FOLDERS_CAPABILITY: &str = "project_folders_v1"; @@ -357,7 +357,7 @@ impl RemoteInstallAuthority { expected_grant: [u8; 16], ) -> Result { let peer = peer.to_string(); - let store = node.store.lock().map_err(|_| "Remote grant unavailable")?; + let store = node.store.read().map_err(|_| "Remote grant unavailable")?; if !store.permits(peer.parse().map_err(|_| "Remote grant unavailable")?) || !store .peers() @@ -401,7 +401,7 @@ impl RemoteInstallAuthority { ) -> Result { let node = self.node.upgrade().ok_or("Remote access ended")?; let peer = self.peer.parse().map_err(|_| "Remote access ended")?; - let store = node.store.lock().map_err(|_| "Remote grant unavailable")?; + let store = node.store.read().map_err(|_| "Remote grant unavailable")?; if node.stopped.is_cancelled() || node.backend.saved_user_id().as_deref() != Some(&self.account) || !store.permits(peer) @@ -497,7 +497,7 @@ mod tests { #[test] fn responses_accept_future_fields_and_unknown_status_without_enabling_mutation() { let value = json!({ - "protocol":{"version":"6.3.0","capabilities":[INSTALL_CAPABILITY,"future_feature"],"future":true}, + "protocol":{"version":"7.1.0","capabilities":[INSTALL_CAPABILITY,"future_feature"],"future":true}, "running":{"version":"0.0.6","source_sha":"a".repeat(40),"platform":"linux-x86_64","future":true}, "update":{"state":"future_state","future":true}, "future":true @@ -507,10 +507,10 @@ mod tests { assert_eq!(info.update.state, UpdateState::Unknown); assert!(info.supports_remote_updates()); let mut next_major = info.clone(); - next_major.protocol.version = "7.0.0".into(); + next_major.protocol.version = "8.0.0".into(); assert!(!next_major.supports_remote_updates()); let mut previous_major = info.clone(); - previous_major.protocol.version = "5.1.0".into(); + previous_major.protocol.version = "6.2.0".into(); assert!(!previous_major.supports_remote_updates()); let mut no_capability = info; no_capability.protocol.capabilities.clear(); @@ -735,7 +735,7 @@ mod tests { } #[test] - fn released_v6_1_decoder_preserves_v6_2_updates_and_durable_receipts() { + fn released_v6_decoder_recognizes_a_new_major_without_install_authority() { let mut current = current_peer_for_compatibility(); for state in [UpdateState::Ready, UpdateState::RestartRequired] { current.update.state = state; @@ -747,8 +747,11 @@ mod tests { ); let legacy: released_v6_1::PeerInfo = serde_json::from_slice(&bytes).unwrap(); assert!(legacy.valid()); - assert_eq!(legacy.protocol.version, "6.2.0"); - assert!(legacy.supports_remote_updates()); + assert_eq!(legacy.protocol.version, "7.0.0"); + assert!( + !legacy.supports_remote_updates(), + "no cross-major management bridge" + ); assert!( legacy .protocol @@ -790,31 +793,31 @@ mod tests { } #[test] - fn project_folder_support_requires_its_capability_and_wire_six() { + fn project_folder_support_requires_its_capability_and_current_major() { let mut info = current_peer_for_compatibility(); assert!(info.supports_project_folders()); - for version in ["6.1.0", "6.2.0", "6.3.0"] { + for version in ["7.0.0", "7.1.0", "7.2.0"] { info.protocol.version = version.into(); assert!( info.supports_project_folders(), "explicit capability in {version}" ); } - for version in ["5.2.0", "7.0.0", "invalid"] { + for version in ["6.2.0", "8.0.0", "invalid"] { info.protocol.version = version.into(); assert!(!info.supports_project_folders()); } - info.protocol.version = "6.1.0".into(); + info.protocol.version = "7.0.0".into(); info.protocol.capabilities = vec!["peer_info_v1".into(), "remote_update_install_v1".into()]; assert!( info.supports_remote_updates(), - "released 6.1 update support remains" + "explicit current-major update support remains" ); assert!( !info.supports_project_folders(), - "6.1 does not advertise folder support" + "folder support still requires its capability" ); - info.protocol.version = "6.2.0".into(); + info.protocol.version = "7.1.0".into(); assert!( !info.supports_project_folders(), "a minor version alone is not a capability" diff --git a/apps/maple-agent/app/src/remote/uploads.rs b/apps/maple-agent/app/src/remote/uploads.rs new file mode 100644 index 000000000..e9582d466 --- /dev/null +++ b/apps/maple-agent/app/src/remote/uploads.rs @@ -0,0 +1,413 @@ +//! Wire 7 attachment uploads. Metadata is bounded JSON; payload bytes occupy a +//! separate low-priority QUIC stream and never an ordinary RPC/control frame. +use super::{manager::Node, streams, wire}; +use crate::{ + backend::{ + LocalBackend, + attachments::{self, ImportAttachmentRequest, PreparedImport}, + }, + host::protocol::Message, +}; +use iroh::{ + EndpointId, + endpoint::{Connection, RecvStream, SendStream}, +}; +use maple_agent::store::{ + AttachmentUploadMetadata, AttachmentUploadStart, AttachmentUploadWriter, + PreparedAttachmentUpload, StoredAttachment, +}; +use serde::{Deserialize, Serialize}; +use std::{ + sync::{Arc, RwLock}, + time::Duration, +}; +use tokio::io::AsyncReadExt; + +const CHUNK: usize = 64 * 1024; +const IDLE_TIMEOUT: Duration = Duration::from_secs(60); +// The idle bound permits useful slow links; the total bound prevents a peer +// retaining a bulk slot forever by supplying one byte before each idle timeout. +const TOTAL_TIMEOUT: Duration = Duration::from_secs(10 * 60); + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(super) struct Upload { + pub user: String, + pub session_id: String, + pub operation_id: String, + pub metadata: AttachmentUploadMetadata, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "state", content = "attachment", rename_all = "camelCase")] +pub(super) enum Admission { + Ready, + Complete(StoredAttachment), +} + +pub(super) async fn call( + connection: &Connection, + id: u64, + user: String, + request: ImportAttachmentRequest, + prepared: PreparedImport, +) -> Result { + let upload = Upload { + user, + session_id: request.session_id, + operation_id: request.operation_id, + metadata: prepared.metadata, + }; + // Query the stable upload operation after a lost acknowledgement; never + // treat a transport error as permission to replay a prompt. + tokio::time::timeout(TOTAL_TIMEOUT, async { + let (mut tx, mut rx) = tokio::time::timeout(streams::IO_TIMEOUT, connection.open_bi()) + .await + .map_err(|_| "Attachment stream capacity timed out")? + .map_err(|_| "Remote connection closed")?; + tx.set_priority(-5) + .map_err(|_| "Attachment stream closed")?; + let header = streams::Header { + id, + method: wire::Method::UploadAttachment, + send_id: None, + }; + tokio::time::timeout(streams::IO_TIMEOUT, async { + wire::write(&mut tx, &header, wire::MAX_PREFACE).await?; + wire::write(&mut tx, &upload, wire::MAX_PREFACE).await + }) + .await + .map_err(|_| "Attachment metadata timed out")? + .map_err(|e| e.to_string())?; + let admission = wire::read_progress::>( + &mut rx, + wire::MAX_MANAGEMENT_REPLY, + IDLE_TIMEOUT, + false, + ) + .await + .map_err(|e| e.to_string())??; + match admission { + Admission::Complete(attachment) => { + tx.finish().map_err(|_| "Attachment stream closed")?; + streams::finish_reply(&mut rx, streams::IO_TIMEOUT).await?; + return Ok(Message::Response { + id, + result: serde_json::to_value(attachment) + .map_err(|_| "Could not encode attachment".into()), + }); + } + Admission::Ready => {} + } + let mut source = tokio::fs::File::from_std(prepared.file); + let mut buffer = vec![0u8; CHUNK]; + let mut remaining = upload.metadata.size; + while remaining > 0 { + let count = usize::try_from(remaining.min(CHUNK as u64)).unwrap(); + let count = source + .read(&mut buffer[..count]) + .await + .map_err(|_| "Could not read staged attachment")?; + if count == 0 { + return Err("Staged attachment changed during upload".to_string()); + } + tokio::time::timeout(IDLE_TIMEOUT, tx.write_all(&buffer[..count])) + .await + .map_err(|_| "Attachment upload stalled")? + .map_err(|_| "Attachment upload disconnected")?; + remaining -= count as u64; + } + // A growing staged file is not silently accepted either. + if source + .read(&mut buffer[..1]) + .await + .map_err(|_| "Could not read staged attachment")? + != 0 + { + return Err("Staged attachment changed during upload".into()); + } + tx.finish().map_err(|_| "Attachment upload closed")?; + let response: Message = + wire::read_progress(&mut rx, wire::MAX_MANAGEMENT_REPLY, IDLE_TIMEOUT, false) + .await + .map_err(|e| e.to_string())?; + if !matches!(&response, Message::Response { id: received, .. } if *received == id) { + return Err("Invalid attachment acknowledgement".into()); + } + streams::finish_reply(&mut rx, streams::IO_TIMEOUT).await?; + Ok(response) + }) + .await + .map_err(|_| "Attachment upload timed out")? +} + +async fn read_metadata( + reader: &mut (impl tokio::io::AsyncRead + Unpin), + deadline: Duration, +) -> Result { + // Metadata is tiny. Its *total* budget must not restart with every byte. + tokio::time::timeout(deadline, wire::read(reader, wire::MAX_PREFACE)) + .await + .map_err(|_| "Attachment metadata timed out")? + .map_err(|error| error.to_string()) +} + +pub(super) async fn serve( + node: Arc, + peer: EndpointId, + id: u64, + tx: &mut SendStream, + rx: &mut RecvStream, +) -> Result<(), String> { + serve_receiver(node, peer, id, tx, rx).await +} + +// Keep account/grant admission and publication in one receiver implementation. +// Tests supply only its host context, rather than rebuilding the wire receiver. +trait UploadReceiver: Send + Sync + 'static { + fn backend(&self) -> &LocalBackend; + fn account(&self) -> &str; + fn policy(&self) -> &RwLock; + fn stopped(&self) -> bool; + fn ensure_valid(&self) -> impl std::future::Future> + Send; + #[cfg(test)] + fn before_publish(&self) {} +} + +impl UploadReceiver for Node { + fn backend(&self) -> &LocalBackend { + &self.backend + } + fn account(&self) -> &str { + &self.account + } + fn policy(&self) -> &RwLock { + &self.store + } + fn stopped(&self) -> bool { + self.stopped.is_cancelled() + } + async fn ensure_valid(&self) -> Result<(), String> { + Node::ensure_valid(self).await + } +} + +async fn serve_receiver( + node: Arc, + peer: EndpointId, + id: u64, + tx: &mut SendStream, + rx: &mut RecvStream, +) -> Result<(), String> { + let begin = async { + let upload = read_metadata(rx, streams::IO_TIMEOUT).await?; + if upload.user != node.account() { + return Err("Remote request account mismatch".to_string()); + } + let account_guard = node.backend().account_gate.clone().read_owned().await; + node.ensure_valid().await?; + let (store, account) = attachments::context(node.backend(), &upload.user).await?; + let begin_store = store.clone(); + let begin_node = node.clone(); + let start = tokio::task::spawn_blocking(move || { + // The worker retains account ownership even if its awaiting stream + // task is aborted while a durable begin/reconciliation is running. + let _account = account_guard; + let grant = begin_node + .policy() + .read() + .map_err(|_| "Remote grant unavailable".to_string())?; + begin_node + .backend() + .authorize_account(begin_node.account())?; + if begin_node.stopped() || !grant.permits(peer) { + return Err("Remote access revoked".into()); + } + let grant_id = grant + .peers() + .get(&peer.to_string()) + .ok_or("Remote access revoked")? + .grant; + let start = begin_store + .begin_attachment_upload( + account, + &upload.session_id, + &upload.operation_id, + &upload.metadata, + ) + .map_err(|e| e.to_string())?; + Ok::<_, String>((grant_id, start)) + }) + .await + .map_err(|_| "Attachment worker did not complete")??; + Ok::<_, String>((store, start)) + } + .await; + let (store, (grant_id, start)) = match begin { + Ok(value) => value, + Err(error) => { + let _ = rx.stop(1u32.into()); + return send_admission(tx, Err(error), true).await; + } + }; + check_receiver(node.as_ref(), peer, grant_id).await?; + let writer = match start { + AttachmentUploadStart::Complete(attachment) => { + return send_admission(tx, Ok(Admission::Complete(attachment)), true).await; + } + AttachmentUploadStart::Writing(writer) => *writer, + }; + send_admission(tx, Ok(Admission::Ready), false).await?; + let result = async { + let writer = tokio::time::timeout(TOTAL_TIMEOUT, receive_bytes(rx, writer)) + .await + .map_err(|_| "Attachment upload timed out")??; + // Slow file sync owns no policy guard. The sealed partial keeps its + // upload lease and still cannot be consumed as a committed attachment. + let prepare_store = store.clone(); + let prepared = attachments::blocking(move || writer.prepare(&prepare_store)).await?; + let account_guard = node.backend().account_gate.clone().read_owned().await; + node.ensure_valid().await?; + node.backend().authorize_account(node.account())?; + let node = node.clone(); + tokio::task::spawn_blocking(move || { + let _account = account_guard; + // A shared policy guard orders publication against revoke/disable + // without excluding Stop, question replies, or event permission + // checks while rename/directory sync/SQLite completion runs. + let grant = node + .policy() + .read() + .map_err(|_| "Remote grant unavailable".to_string())?; + node.backend().authorize_account(node.account())?; + #[cfg(test)] + node.before_publish(); + finish_authorized(prepared, &store, &grant, peer, grant_id, node.stopped()) + }) + .await + .map_err(|_| "Attachment worker did not complete")? + } + .await; + if result.is_err() { + let _ = rx.stop(1u32.into()); + } + // As with ordinary RPC replies, a completed effect does not authorize + // exposing its receipt to a peer whose account/grant ended while it ran. + check_receiver(node.as_ref(), peer, grant_id).await?; + let result = result.and_then(|attachment| { + serde_json::to_value(attachment).map_err(|_| "Could not encode attachment".into()) + }); + streams::reply_with_limit( + tx, + &Message::Response { id, result }, + wire::MAX_MANAGEMENT_REPLY, + ) + .await +} + +async fn check_receiver( + node: &R, + peer: EndpointId, + admitted_grant: [u8; 16], +) -> Result<(), String> { + node.ensure_valid().await?; + let policy = node + .policy() + .read() + .map_err(|_| "Remote grant unavailable")?; + if policy.permits(peer) + && policy + .peers() + .get(&peer.to_string()) + .is_some_and(|current| current.grant == admitted_grant) + { + Ok(()) + } else { + Err("Remote access revoked".into()) + } +} + +fn finish_authorized( + prepared: PreparedAttachmentUpload, + store: &maple_agent::store::Store, + policy: &super::store::Store, + peer: EndpointId, + admitted_grant: [u8; 16], + stopped: bool, +) -> Result { + if stopped + || !policy.permits(peer) + || !policy + .peers() + .get(&peer.to_string()) + .is_some_and(|current| current.grant == admitted_grant) + { + return Err("Remote access revoked".into()); + } + prepared.publish(store).map_err(|error| error.to_string()) +} + +pub(super) async fn refuse(tx: &mut SendStream, error: String) -> Result<(), String> { + send_admission(tx, Err(error), true).await +} + +async fn send_admission( + tx: &mut SendStream, + value: Result, + finish: bool, +) -> Result<(), String> { + tokio::time::timeout( + streams::IO_TIMEOUT, + wire::write(tx, &value, wire::MAX_MANAGEMENT_REPLY), + ) + .await + .map_err(|_| "Attachment admission timed out")? + .map_err(|e| e.to_string())?; + if finish { + tx.finish().map_err(|_| "Attachment admission closed")?; + tokio::time::timeout(streams::IO_TIMEOUT, tx.stopped()) + .await + .map_err(|_| "Attachment admission timed out")? + .map_err(|_| "Attachment admission closed")?; + } + Ok(()) +} + +async fn receive_bytes( + rx: &mut RecvStream, + mut writer: AttachmentUploadWriter, +) -> Result { + let mut buffer = vec![0u8; CHUNK]; + let mut cancellation = tokio::time::interval(Duration::from_millis(250)); + loop { + if writer.is_cancelled() { + return Err("Attachment upload cancelled".into()); + } + let count = { + let read = tokio::time::timeout(IDLE_TIMEOUT, AsyncReadExt::read(rx, &mut buffer)); + tokio::pin!(read); + loop { + tokio::select! { + result = &mut read => break result.map_err(|_| "Attachment upload stalled")?.map_err(|_| "Attachment upload disconnected")?, + _ = cancellation.tick() => if writer.is_cancelled() { return Err("Attachment upload cancelled".into()); } + } + } + }; + if count == 0 { + return Ok(writer); + } + // Writer enforces declared length, digest and cancellation. Only one + // bounded chunk exists in memory and no SQLite transaction spans a read. + (writer, buffer) = tokio::task::spawn_blocking(move || { + writer + .write_chunk(&buffer[..count]) + .map_err(|e| e.to_string())?; + Ok::<_, String>((writer, buffer)) + }) + .await + .map_err(|_| "Attachment worker did not complete")??; + } +} + +#[cfg(test)] +mod tests; diff --git a/apps/maple-agent/app/src/remote/uploads/tests.rs b/apps/maple-agent/app/src/remote/uploads/tests.rs new file mode 100644 index 000000000..ee1837786 --- /dev/null +++ b/apps/maple-agent/app/src/remote/uploads/tests.rs @@ -0,0 +1,956 @@ +use super::*; +use crate::remote::streams::tests::{command, endpoints}; +use maple_agent::store::{AccountId, AttachmentUploadStatus, NewTask, Store, StorePaths, TaskKind}; +use sha2::{Digest, Sha256}; + +const TASK: &str = "01234567-89ab-7def-8123-456789abcdef"; +const OP: &str = "11111111-1111-4111-8111-111111111111"; + +fn store(root: &std::path::Path) -> (Arc, AccountId) { + let store = Arc::new(Store::open(StorePaths::new(root.join("execution"))).unwrap()); + let account = store.upsert_account("account", 1).unwrap(); + store + .create_task( + account, + &NewTask { + uuid: TASK.into(), + project_root: root.to_string_lossy().into_owned(), + kind: TaskKind::Desktop, + title: "Upload".into(), + provisional: true, + model: None, + created_ms: 1, + }, + ) + .unwrap(); + (store, account) +} + +fn metadata(bytes: &[u8]) -> AttachmentUploadMetadata { + AttachmentUploadMetadata { + name: "notes.bin".into(), + mime_type: "application/octet-stream".into(), + size: bytes.len() as u64, + sha256: format!("{:x}", Sha256::digest(bytes)), + } +} + +#[test] +fn upload_metadata_cannot_name_a_controller_or_executor_source_path() { + let value = serde_json::json!({"user":"account","sessionId":TASK,"operationId":OP,"metadata":metadata(b"x"),"path":"/private/controller/secret"}); + assert!(serde_json::from_value::(value).is_err()); + assert!(wire::Method::parse("import_attachment").is_err()); +} + +#[test] +fn bulk_budgets_leave_urgent_capacity_available() { + let capacity = streams::Capacity::new(); + let _bulk: Vec<_> = (0..2) + .map(|_| capacity.acquire(wire::Method::UploadAttachment).unwrap()) + .collect(); + let _ordinary: Vec<_> = (0..32) + .map(|_| capacity.acquire(wire::Method::ListTasks).unwrap()) + .collect(); + assert!(capacity.acquire(wire::Method::UploadAttachment).is_err()); + assert!(capacity.acquire(wire::Method::CancelRun).is_ok()); + assert!( + capacity + .acquire(wire::Method::CancelAttachmentUpload) + .is_ok() + ); + assert_eq!(wire::MAX_FRAME, 16 * 1024 * 1024); +} + +#[tokio::test] +async fn binary_file_larger_than_a_control_frame_round_trips_and_reuses_operation() { + let root = tempfile::tempdir().unwrap(); + let (store, account) = store(root.path()); + let bytes = vec![0xabu8; wire::MAX_FRAME + 123]; + let path = root.path().join("staged.bin"); + std::fs::write(&path, &bytes).unwrap(); + let import = ImportAttachmentRequest { + session_id: TASK.into(), + operation_id: OP.into(), + path: path.to_string_lossy().into_owned(), + name: "notes.bin".into(), + mime_type: "application/octet-stream".into(), + }; + let (a, b, client, server) = endpoints().await; + let destination = store.clone(); + let serving = tokio::spawn(async move { + // The second import repeats metadata and returns its committed receipt + // before admitting any bytes; it never creates a second attachment. + for _ in 0..2 { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + assert_eq!(header.method, wire::Method::UploadAttachment); + assert!(header.send_id.is_none()); + let upload: Upload = wire::read(&mut rx, wire::MAX_PREFACE).await.unwrap(); + assert_eq!(upload.metadata.size, (wire::MAX_FRAME + 123) as u64); + match destination + .begin_attachment_upload(account, TASK, OP, &upload.metadata) + .unwrap() + { + AttachmentUploadStart::Complete(attachment) => { + send_admission(&mut tx, Ok(Admission::Complete(attachment)), true) + .await + .unwrap() + } + AttachmentUploadStart::Writing(writer) => { + send_admission(&mut tx, Ok(Admission::Ready), false) + .await + .unwrap(); + let writer = receive_bytes(&mut rx, *writer).await.unwrap(); + let attachment = writer.finish(&destination).unwrap(); + streams::reply( + &mut tx, + &Message::Response { + id: header.id, + result: Ok(serde_json::to_value(attachment).unwrap()), + }, + ) + .await + .unwrap(); + } + } + } + }); + let mut ids = Vec::new(); + for id in 1..=2 { + let prepared = attachments::prepare(&import).await.unwrap(); + let response = call(&client, id, "account".into(), import.clone(), prepared) + .await + .unwrap(); + let Message::Response { + result: Ok(value), .. + } = response + else { + panic!("upload failed") + }; + let result: StoredAttachment = serde_json::from_value(value).unwrap(); + assert_eq!(std::fs::read(&result.execution_path).unwrap(), bytes); + assert_ne!(std::path::Path::new(&result.execution_path), path.as_path()); + ids.push(result.id); + } + assert_eq!(ids[0], ids[1]); + serving.await.unwrap(); + assert!(matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + )); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn stalled_raw_upload_does_not_delay_stop_or_question_answer() { + let root = tempfile::tempdir().unwrap(); + let (store, account) = store(root.path()); + let (a, b, client, server) = endpoints().await; + let capacity = streams::Capacity::new(); + let (mut tx, mut rx) = client.open_bi().await.unwrap(); + wire::write( + &mut tx, + &streams::Header { + id: 1, + method: wire::Method::UploadAttachment, + send_id: None, + }, + wire::MAX_PREFACE, + ) + .await + .unwrap(); + let upload = Upload { + user: "account".into(), + session_id: TASK.into(), + operation_id: OP.into(), + metadata: metadata(&vec![7; 1024 * 1024]), + }; + wire::write(&mut tx, &upload, wire::MAX_PREFACE) + .await + .unwrap(); + let (mut response, mut body) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut body).await.unwrap(); + let slot = capacity.acquire(header.method).unwrap(); + let received: Upload = wire::read(&mut body, wire::MAX_PREFACE).await.unwrap(); + let AttachmentUploadStart::Writing(writer) = store + .begin_attachment_upload(account, TASK, OP, &received.metadata) + .unwrap() + else { + panic!() + }; + send_admission(&mut response, Ok(Admission::Ready), false) + .await + .unwrap(); + let _: Result = wire::read(&mut rx, wire::MAX_MANAGEMENT_REPLY) + .await + .unwrap(); + tx.write_all(&[7; 128]).await.unwrap(); + let upload_task = tokio::spawn(async move { + let _slot = slot; + receive_bytes(&mut body, *writer).await + }); + for (id, method) in [ + (2, wire::Method::CancelRun), + (3, wire::Method::AnswerQuestion), + ] { + let caller = streams::call(&client, command(id, method)); + let answer = async { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + let _slot = capacity.acquire(header.method).unwrap(); + header.body(&mut rx).await.unwrap(); + streams::reply( + &mut tx, + &Message::Response { + id, + result: Ok(serde_json::Value::Null), + }, + ) + .await + .unwrap(); + }; + let (reply, ()) = tokio::time::timeout(Duration::from_secs(2), async { + tokio::join!(caller, answer) + }) + .await + .unwrap(); + assert!(matches!( + reply.unwrap(), + Message::Response { result: Ok(_), .. } + )); + assert!(!upload_task.is_finished()); + } + tx.reset(1u32.into()).unwrap(); + assert!(upload_task.await.unwrap().is_err()); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn raw_truncation_digest_mismatch_and_cancel_never_publish() { + for kind in ["truncated", "digest", "cancelled", "oversized"] { + let root = tempfile::tempdir().unwrap(); + let (store, account) = store(root.path()); + let (a, b, client, server) = endpoints().await; + let declared = metadata(b"abcd"); + let AttachmentUploadStart::Writing(writer) = store + .begin_attachment_upload(account, TASK, OP, &declared) + .unwrap() + else { + panic!() + }; + let (mut tx, _rx) = client.open_bi().await.unwrap(); + tx.write_all(if kind == "truncated" { + b"ab" + } else if kind == "oversized" { + b"abcde" + } else { + b"zzzz" + }) + .await + .unwrap(); + tx.finish().unwrap(); + let (_reply, mut rx) = server.accept_bi().await.unwrap(); + if kind == "cancelled" { + store.cancel_attachment_upload(account, TASK, OP).unwrap(); + } + let result = match receive_bytes(&mut rx, *writer).await { + Ok(writer) => writer.finish(&store).map_err(|e| e.to_string()), + Err(e) => Err(e), + }; + assert!(result.is_err(), "{kind}"); + assert!( + !matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + ), + "{kind}" + ); + // A failed bulk request does not terminate the underlying connection. + assert!(client.close_reason().is_none()); + a.close().await; + b.close().await; + } +} + +#[tokio::test] +async fn lost_final_ack_is_resolved_by_the_upload_receipt_without_prompt_replay() { + let root = tempfile::tempdir().unwrap(); + let backend = Arc::new(crate::backend::attachments::tests::fixture(root.path())); + let store = backend.store().unwrap().clone(); + let account = store.account_id("account").unwrap().unwrap(); + let path = root.path().join("stable-upload.bin"); + std::fs::write(&path, b"abcd").unwrap(); + let request = ImportAttachmentRequest { + session_id: TASK.into(), + operation_id: OP.into(), + path: path.to_string_lossy().into_owned(), + name: "notes.bin".into(), + mime_type: "application/octet-stream".into(), + }; + let (a, b, client, server) = endpoints().await; + let serving = tokio::spawn(async move { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + assert_eq!(header.method, wire::Method::UploadAttachment); + let upload: Upload = read_metadata(&mut rx, streams::IO_TIMEOUT).await.unwrap(); + let AttachmentUploadStart::Writing(writer) = store + .begin_attachment_upload(account, TASK, OP, &upload.metadata) + .unwrap() + else { + panic!() + }; + send_admission(&mut tx, Ok(Admission::Ready), false) + .await + .unwrap(); + let writer = receive_bytes(&mut rx, *writer).await.unwrap(); + let attachment = writer.finish(&store).unwrap(); + // Lose the response after durable publication, not before admission. + tx.reset(1u32.into()).unwrap(); + + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + assert_eq!(header.method, wire::Method::AttachmentUploadStatus); + let command = header.body(&mut rx).await.unwrap(); + let result = crate::backend::dispatch( + &backend, + crate::host::protocol::Request { + id: command.id, + method: command.method.name(), + send_id: None, + params: command.params, + }, + ) + .await; + streams::reply( + &mut tx, + &Message::Response { + id: command.id, + result, + }, + ) + .await + .unwrap(); + + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + assert_eq!(header.method, wire::Method::UploadAttachment); + let upload: Upload = read_metadata(&mut rx, streams::IO_TIMEOUT).await.unwrap(); + let AttachmentUploadStart::Complete(repeated) = store + .begin_attachment_upload(account, TASK, OP, &upload.metadata) + .unwrap() + else { + panic!("retry must not admit another byte transfer") + }; + assert_eq!(attachment.id, repeated.id); + send_admission(&mut tx, Ok(Admission::Complete(repeated)), true) + .await + .unwrap(); + // All three requests were uploads/status: none can enter send_message. + attachment + }); + assert!( + call( + &client, + 1, + "account".into(), + request.clone(), + attachments::prepare(&request).await.unwrap() + ) + .await + .is_err() + ); + let status = streams::call( + &client, + wire::Command { + id: 2, + method: wire::Method::AttachmentUploadStatus, + send_id: None, + params: serde_json::json!([ + "account", + crate::backend::AttachmentUploadRequest { + session_id: TASK.into(), + operation_id: OP.into() + } + ]), + }, + ) + .await + .unwrap(); + let Message::Response { + result: Ok(value), .. + } = status + else { + panic!("status lookup failed") + }; + let AttachmentUploadStatus::Complete(recovered) = serde_json::from_value(value).unwrap() else { + panic!("committed receipt missing") + }; + let retry = call( + &client, + 3, + "account".into(), + request.clone(), + attachments::prepare(&request).await.unwrap(), + ) + .await + .unwrap(); + let Message::Response { + result: Ok(value), .. + } = retry + else { + panic!("idempotent retry failed") + }; + assert_eq!( + serde_json::from_value::(value) + .unwrap() + .id, + recovered.id + ); + assert_eq!(serving.await.unwrap().id, recovered.id); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn cancelling_a_stalled_upload_releases_its_stream_without_waiting_for_idle_timeout() { + let root = tempfile::tempdir().unwrap(); + let (store, account) = store(root.path()); + let (a, b, client, server) = endpoints().await; + let AttachmentUploadStart::Writing(writer) = store + .begin_attachment_upload(account, TASK, OP, &metadata(b"abcd")) + .unwrap() + else { + panic!() + }; + let (mut tx, _rx) = client.open_bi().await.unwrap(); + tx.write_all(b"a").await.unwrap(); + let (_response, mut body) = server.accept_bi().await.unwrap(); + let receiving = tokio::spawn(async move { receive_bytes(&mut body, *writer).await }); + store.cancel_attachment_upload(account, TASK, OP).unwrap(); + let result = tokio::time::timeout(Duration::from_secs(2), receiving) + .await + .unwrap() + .unwrap(); + assert!(matches!(result, Err(error) if error=="Attachment upload cancelled")); + assert_eq!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Cancelled + ); + a.close().await; + b.close().await; +} + +#[tokio::test(start_paused = true)] +async fn slowly_dripped_upload_metadata_has_an_overall_deadline() { + use tokio::io::AsyncWriteExt; + let (mut tx, mut rx) = tokio::io::duplex(64); + let send = tokio::spawn(async move { + tx.write_u32(100).await.unwrap(); + for _ in 0..100 { + tokio::time::sleep(Duration::from_secs(1)).await; + if tx.write_all(b" ").await.is_err() { + break; + } + } + }); + let started = tokio::time::Instant::now(); + let error = read_metadata(&mut rx, Duration::from_secs(3)) + .await + .unwrap_err(); + assert_eq!(error, "Attachment metadata timed out"); + assert_eq!(started.elapsed(), Duration::from_secs(3)); + send.abort(); +} + +#[test] +fn revocation_disable_and_repair_do_not_publish_an_upload_admitted_under_an_old_grant() { + for mode in ["revoke", "disable", "re-pair", "account stopped"] { + let root = tempfile::tempdir().unwrap(); + let (store, account) = store(root.path()); + let mut policy = crate::remote::store::Store::open( + store.clone(), + "local", + "project", + "unpackaged", + "account", + ) + .unwrap(); + let peer = iroh::SecretKey::generate().public(); + let grant = crate::remote::store::Peer { + name: "Controller".into(), + grant: [1; 16], + controls_this_computer: true, + can_connect: false, + pending_confirmation: None, + }; + policy.set_enabled(true).unwrap(); + policy.grant(peer, grant.clone()).unwrap(); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(account, TASK, OP, &metadata(b"abcd")) + .unwrap() + else { + panic!() + }; + writer.write_chunk(b"abcd").unwrap(); + let prepared = writer.prepare(&store).unwrap(); + match mode { + "revoke" => policy.revoke(peer).unwrap(), + "disable" => policy.set_enabled(false).unwrap(), + "re-pair" => { + policy.revoke(peer).unwrap(); + let mut next = grant; + next.grant = [2; 16]; + policy.grant(peer, next).unwrap(); + } + _ => {} + } + assert!( + finish_authorized( + prepared, + &store, + &policy, + peer, + [1; 16], + mode == "account stopped" + ) + .is_err(), + "{mode}" + ); + assert!( + !matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + ), + "{mode}" + ); + } +} + +struct PublicationPause { + entered: std::sync::Mutex>>, + release: std::sync::Mutex>, +} + +struct ReceiverFixture { + backend: Arc, + policy: RwLock, + stopped: std::sync::atomic::AtomicBool, + publication: Option, +} + +impl UploadReceiver for ReceiverFixture { + fn backend(&self) -> &LocalBackend { + &self.backend + } + fn account(&self) -> &str { + "account" + } + fn policy(&self) -> &RwLock { + &self.policy + } + fn stopped(&self) -> bool { + self.stopped.load(std::sync::atomic::Ordering::Acquire) + } + async fn ensure_valid(&self) -> Result<(), String> { + if self.stopped() { + Err("The remote account is no longer active".into()) + } else { + Ok(()) + } + } + fn before_publish(&self) { + if let Some(pause) = &self.publication { + pause + .entered + .lock() + .unwrap() + .take() + .unwrap() + .send(()) + .unwrap(); + // A dropped release sender also unblocks the worker after a failed + // assertion; no test may strand a blocking publication thread. + let _ = pause + .release + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(10)); + } + } +} + +fn receiver_fixture( + root: &std::path::Path, + peer: EndpointId, + publication: Option, +) -> Arc { + let backend = Arc::new(attachments::tests::fixture(root)); + let mut policy = crate::remote::store::Store::open( + backend.store().unwrap().clone(), + "local", + "project", + "unpackaged", + "account", + ) + .unwrap(); + policy.set_enabled(true).unwrap(); + policy + .grant( + peer, + crate::remote::store::Peer { + name: "Controller".into(), + grant: [1; 16], + controls_this_computer: true, + can_connect: false, + pending_confirmation: None, + }, + ) + .unwrap(); + Arc::new(ReceiverFixture { + backend, + policy: RwLock::new(policy), + stopped: std::sync::atomic::AtomicBool::new(false), + publication, + }) +} + +fn fixture_import(root: &std::path::Path) -> ImportAttachmentRequest { + let path = root.join("staged.bin"); + std::fs::write(&path, b"abcd").unwrap(); + ImportAttachmentRequest { + session_id: TASK.into(), + operation_id: OP.into(), + path: path.to_string_lossy().into_owned(), + name: "notes.bin".into(), + mime_type: "application/octet-stream".into(), + } +} + +#[tokio::test] +async fn production_receiver_publication_keeps_control_authorization_responsive() { + let root = tempfile::tempdir().unwrap(); + let (a, b, client, server) = endpoints().await; + let peer = a.id(); + let (entered, publishing) = tokio::sync::oneshot::channel(); + let (release, held) = std::sync::mpsc::channel(); + let receiver = receiver_fixture( + root.path(), + peer, + Some(PublicationPause { + entered: std::sync::Mutex::new(Some(entered)), + release: std::sync::Mutex::new(held), + }), + ); + let destination = receiver.clone(); + let capacity = Arc::new(streams::Capacity::new()); + let serving = tokio::spawn(streams::serve_requests(server, move |mut tx, mut rx| { + let receiver = destination.clone(); + let capacity = capacity.clone(); + async move { + let header = capacity.read_header(&mut rx).await.unwrap(); + let _slot = capacity.acquire(header.method).unwrap(); + if header.method == wire::Method::UploadAttachment { + serve_receiver(receiver, peer, header.id, &mut tx, &mut rx) + .await + .unwrap(); + } else { + let request = header.body(&mut rx).await.unwrap(); + assert!(matches!( + request.method, + wire::Method::CancelRun | wire::Method::AnswerQuestion + )); + // Exercise the same current account/grant check while the real + // upload receiver owns its publication guard. SDK session + // validation belongs to Node and the authenticated live probe. + check_receiver(receiver.as_ref(), peer, [1; 16]) + .await + .unwrap(); + let result = crate::backend::dispatch( + &receiver.backend, + crate::host::protocol::Request { + id: request.id, + method: request.method.name(), + send_id: None, + params: request.params, + }, + ) + .await; + streams::reply( + &mut tx, + &Message::Response { + id: request.id, + result, + }, + ) + .await + .unwrap(); + } + } + })); + let request = fixture_import(root.path()); + let upload_client = client.clone(); + let upload = tokio::spawn(async move { + let prepared = attachments::prepare(&request).await.unwrap(); + call(&upload_client, 1, "account".into(), request, prepared).await + }); + tokio::time::timeout(Duration::from_secs(3), publishing) + .await + .unwrap() + .unwrap(); + assert!( + receiver.policy.try_write().is_err(), + "publication must exclude revoke" + ); + for (id, method, params) in [ + ( + 2, + wire::Method::CancelRun, + serde_json::json!(["account", "run"]), + ), + ( + 3, + wire::Method::AnswerQuestion, + serde_json::json!(["account", "question", "answer"]), + ), + ] { + let response = tokio::time::timeout( + Duration::from_secs(2), + streams::call( + &client, + wire::Command { + id, + method, + send_id: None, + params, + }, + ), + ) + .await + .expect("control authorization waited for disk publication") + .unwrap(); + // The fixture has no running agent/question; a backend response still + // proves the control reached dispatch while publication was held. + assert!(matches!(response, Message::Response { id: received, .. } if received == id)); + assert!(!upload.is_finished()); + } + release.send(()).unwrap(); + assert!(matches!( + upload.await.unwrap().unwrap(), + Message::Response { result: Ok(_), .. } + )); + client.close(0u32.into(), b"done"); + serving.await.unwrap(); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn production_receiver_rechecks_revoke_disable_and_repair_after_payload() { + for mode in ["revoke", "disable", "re-pair", "stopped"] { + let root = tempfile::tempdir().unwrap(); + let (a, b, client, server) = endpoints().await; + let peer = a.id(); + let receiver = receiver_fixture(root.path(), peer, None); + let destination = receiver.clone(); + let serving = tokio::spawn(async move { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + serve_receiver(destination, peer, header.id, &mut tx, &mut rx).await + }); + let (mut tx, mut rx) = client.open_bi().await.unwrap(); + wire::write( + &mut tx, + &streams::Header { + id: 1, + method: wire::Method::UploadAttachment, + send_id: None, + }, + wire::MAX_PREFACE, + ) + .await + .unwrap(); + wire::write( + &mut tx, + &Upload { + user: "account".into(), + session_id: TASK.into(), + operation_id: OP.into(), + metadata: metadata(b"abcd"), + }, + wire::MAX_PREFACE, + ) + .await + .unwrap(); + let admission: Result = wire::read(&mut rx, wire::MAX_MANAGEMENT_REPLY) + .await + .unwrap(); + assert!(matches!(admission, Ok(Admission::Ready))); + tx.write_all(b"abcd").await.unwrap(); + { + let mut policy = receiver.policy.write().unwrap(); + match mode { + "revoke" => policy.revoke(peer).unwrap(), + "disable" => policy.set_enabled(false).unwrap(), + "re-pair" => { + let mut grant = policy.peers().get(&peer.to_string()).unwrap().clone(); + grant.grant = [2; 16]; + policy.revoke(peer).unwrap(); + policy.grant(peer, grant).unwrap(); + } + _ => receiver + .stopped + .store(true, std::sync::atomic::Ordering::Release), + } + } + tx.finish().unwrap(); + assert!( + tokio::time::timeout(Duration::from_secs(3), serving) + .await + .unwrap() + .unwrap() + .is_err(), + "{mode}" + ); + let store = receiver.backend.store().unwrap(); + let account = store.account_id("account").unwrap().unwrap(); + assert!( + !matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + ), + "{mode}" + ); + a.close().await; + b.close().await; + } +} + +#[tokio::test] +async fn aborted_receiver_keeps_account_guard_until_publication_worker_exits() { + let root = tempfile::tempdir().unwrap(); + let (a, b, client, server) = endpoints().await; + let peer = a.id(); + let (entered, publishing) = tokio::sync::oneshot::channel(); + let (release, held) = std::sync::mpsc::channel(); + let receiver = receiver_fixture( + root.path(), + peer, + Some(PublicationPause { + entered: std::sync::Mutex::new(Some(entered)), + release: std::sync::Mutex::new(held), + }), + ); + let destination = receiver.clone(); + let serving = tokio::spawn(async move { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + serve_receiver(destination, peer, header.id, &mut tx, &mut rx).await + }); + let request = fixture_import(root.path()); + let upload_client = client.clone(); + let upload = tokio::spawn(async move { + let prepared = attachments::prepare(&request).await.unwrap(); + call(&upload_client, 1, "account".into(), request, prepared).await + }); + tokio::time::timeout(Duration::from_secs(3), publishing) + .await + .unwrap() + .unwrap(); + serving.abort(); + assert!(serving.await.unwrap_err().is_cancelled()); + assert!( + receiver.backend.account_gate.try_write().is_err(), + "aborting the stream released account ownership before its disk worker" + ); + receiver + .stopped + .store(true, std::sync::atomic::Ordering::Release); + release.send(()).unwrap(); + let transition = tokio::time::timeout( + Duration::from_secs(3), + receiver.backend.account_gate.write(), + ) + .await + .unwrap(); + let store = receiver.backend.store().unwrap(); + let account = store.account_id("account").unwrap().unwrap(); + assert!(!matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + )); + drop(transition); + client.close(0u32.into(), b"done"); + assert!(upload.await.unwrap().is_err()); + a.close().await; + b.close().await; +} + +#[tokio::test] +async fn production_receiver_requires_authority_for_begin_crash_reconciliation() { + fn partial_file(directory: &std::path::Path) -> Option { + for entry in std::fs::read_dir(directory).unwrap() { + let entry = entry.unwrap(); + if entry.file_type().unwrap().is_dir() { + if let Some(path) = partial_file(&entry.path()) { + return Some(path); + } + } else if entry.file_name() == ".part" { + return Some(entry.path()); + } + } + None + } + for permitted in [true, false] { + let root = tempfile::tempdir().unwrap(); + let (a, b, client, server) = endpoints().await; + let peer = a.id(); + let receiver = receiver_fixture(root.path(), peer, None); + let store = receiver.backend.store().unwrap().clone(); + let account = store.account_id("account").unwrap().unwrap(); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(account, TASK, OP, &metadata(b"abcd")) + .unwrap() + else { + panic!() + }; + writer.write_chunk(b"abcd").unwrap(); + drop(writer); + // Simulate a crash after rename but before the ready-row commit, using + // only this disposable fixture's files. Begin must authorize recovery. + let partial = partial_file(root.path()).unwrap(); + std::fs::rename(&partial, partial.with_file_name("notes.bin")).unwrap(); + if !permitted { + receiver.policy.write().unwrap().revoke(peer).unwrap(); + } + let destination = receiver.clone(); + let serving = tokio::spawn(async move { + let (mut tx, mut rx) = server.accept_bi().await.unwrap(); + let header = streams::Header::read(&mut rx).await.unwrap(); + serve_receiver(destination, peer, header.id, &mut tx, &mut rx).await + }); + let request = fixture_import(root.path()); + let prepared = attachments::prepare(&request).await.unwrap(); + let response = tokio::time::timeout( + Duration::from_secs(3), + call(&client, 1, "account".into(), request, prepared), + ) + .await + .unwrap(); + assert_eq!(response.is_ok(), permitted); + assert_eq!( + matches!( + store.attachment_upload_status(account, TASK, OP).unwrap(), + AttachmentUploadStatus::Complete(_) + ), + permitted + ); + let _ = tokio::time::timeout(Duration::from_secs(3), serving) + .await + .unwrap() + .unwrap(); + a.close().await; + b.close().await; + } +} diff --git a/apps/maple-agent/app/src/remote/wire.rs b/apps/maple-agent/app/src/remote/wire.rs index e07479c19..8ed2c6f16 100644 --- a/apps/maple-agent/app/src/remote/wire.rs +++ b/apps/maple-agent/app/src/remote/wire.rs @@ -6,7 +6,7 @@ use zeroize::Zeroizing; // Keep the ALPN family stable so older peers receive the explicit version refusal. pub const ALPN: &[u8] = b"maple-agent/remote/1"; -pub const VERSION: u32 = 6; +pub const VERSION: u32 = 7; pub const MAX_FRAME: usize = 16 * 1024 * 1024; pub const MAX_PREFACE: usize = 4096; pub const MAX_MANAGEMENT_REPLY: usize = 16 * 1024; @@ -124,6 +124,9 @@ pub enum Method { TaskFacts, ReadRow, ReadAttachment, + UploadAttachment, + AttachmentUploadStatus, + CancelAttachmentUpload, QueuedMessage, Subscribe, ListMcpDefinitions, @@ -179,6 +182,8 @@ impl Method { pub fn priority(self) -> i32 { if self.urgent() { 10 + } else if self == Self::UploadAttachment { + -5 } else if matches!(self, Self::OpenTask | Self::HistoryBefore) { 5 } else { @@ -187,7 +192,14 @@ impl Method { } pub fn reply_limit(self) -> usize { - if matches!(self, Self::PeerInfo | Self::RemoteUpdateInstall) { + if matches!( + self, + Self::PeerInfo + | Self::RemoteUpdateInstall + | Self::UploadAttachment + | Self::AttachmentUploadStatus + | Self::CancelAttachmentUpload + ) { MAX_MANAGEMENT_REPLY } else if matches!(self, Self::BrowseProjectFolders | Self::CreateProjectFolder) { 256 * 1024 @@ -216,7 +228,10 @@ impl Method { pub fn urgent(self) -> bool { matches!( self, - Self::CancelRun | Self::AnswerQuestion | Self::CancelExternalAgent + Self::CancelRun + | Self::AnswerQuestion + | Self::CancelExternalAgent + | Self::CancelAttachmentUpload ) } pub fn parse(value: &str) -> Result { @@ -249,13 +264,25 @@ pub enum Error { InvalidFrame, Incompatible, } + +const INCOMPATIBLE_PROTOCOL_MESSAGE: &str = + "These Maple remote protocol versions are incompatible. Update both computers."; + +/// The GUI receives terminal refusals through the host's string error boundary. +/// Only the explicit protocol refusal establishes that an update is required; +/// malformed frames, network failures and revoked authority do not. +#[cfg(any(feature = "desktop", test))] +pub(crate) fn is_protocol_incompatibility_refusal(message: &str) -> bool { + message.strip_prefix("Remote access ended: ") == Some(INCOMPATIBLE_PROTOCOL_MESSAGE) +} + impl std::fmt::Display for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(match self { Self::TimedOut => "Remote transfer stopped making progress", Self::Disconnected => "Remote computer disconnected", Self::InvalidFrame => "The remote computer sent an invalid message. Update both computers before reconnecting.", - Self::Incompatible => "These Maple remote protocol versions are incompatible. Update both computers.", + Self::Incompatible => INCOMPATIBLE_PROTOCOL_MESSAGE, }) } } @@ -492,6 +519,23 @@ mod tests { )); } + #[test] + fn only_explicit_terminal_protocol_refusal_requires_an_update() { + let incompatible = Error::Incompatible.to_string(); + assert!(is_protocol_incompatibility_refusal( + &Refusal::terminal(incompatible.clone()).into_error() + )); + for error in [ + Refusal::temporary(incompatible).into_error(), + Refusal::terminal(Error::InvalidFrame.to_string()).into_error(), + Refusal::terminal("This computer is no longer paired").into_error(), + Error::Disconnected.to_string(), + Error::TimedOut.to_string(), + ] { + assert!(!is_protocol_incompatibility_refusal(&error), "{error}"); + } + } + #[test] fn compatible_versions_do_not_depend_on_build_identity() { let mut hello = Hello { @@ -519,7 +563,7 @@ mod tests { #[test] fn optional_management_keeps_the_shipped_control_ack_and_alpn() { assert_eq!(ALPN, b"maple-agent/remote/1"); - assert_eq!(VERSION, 6); + assert_eq!(VERSION, 7); assert_eq!( serde_json::to_string(&Ok::<(), Refusal>(())).unwrap(), r#"{"Ok":null}"# @@ -548,18 +592,66 @@ mod tests { } #[tokio::test] - async fn optional_folder_commands_preserve_the_shipped_wire_six_handshake() { - // Folder commands require optional capability discovery. They do not - // change the legacy handshake or prevent ordinary version-6 access. + async fn file_upload_wire_upgrade_refuses_six_before_decoding_requests() { + // Frozen released fixture: version 6 carries inline image data URLs. + // A clean baseline refuses it before those requests reach the decoder. let previous: serde_json::Value = serde_json::from_str(include_str!("wire_v6.json")).unwrap(); for kind in ["hello", "pair"] { let (mut tx, mut rx) = tokio::io::duplex(4096); write(&mut tx, &previous[kind], MAX_PREFACE).await.unwrap(); - let hello = read_hello(&mut rx).await.unwrap(); - assert_eq!(hello.version, 6); - assert_eq!(hello.minimum_version, 6); - assert_eq!(serde_json::to_value(hello).unwrap(), previous[kind]); + let refusal = match read_hello(&mut rx).await { + Err(refusal) => refusal, + Ok(_) => panic!("wire 6 must update before file attachments are sent"), + }; + assert!(is_protocol_incompatibility_refusal(&refusal.into_error())); + } + } + + #[tokio::test] + async fn released_wire_six_refuses_seven_before_reading_any_request() { + // Freeze the released reader's version policy and refusal shape. Do not + // reuse the current VERSION/check_version when testing the other side. + #[derive(Deserialize)] + struct ReleasedVersion { + version: u32, + minimum_version: u32, + } + #[derive(Serialize)] + struct ReleasedRefusal { + terminal: bool, + message: &'static str, + } + let current: serde_json::Value = + serde_json::from_str(include_str!("wire_v7.json")).unwrap(); + for kind in ["hello", "pair"] { + let (mut client, mut server) = tokio::io::duplex(4096); + write(&mut client, ¤t[kind], MAX_PREFACE) + .await + .unwrap(); + let untouched = b"request-body-must-not-be-read"; + client.write_all(untouched).await.unwrap(); + let value: serde_json::Value = read(&mut server, 4096).await.unwrap(); + let version: ReleasedVersion = serde_json::from_value(value).unwrap(); + let refused = version.minimum_version == 0 + || version.minimum_version > version.version + || version.minimum_version > 6 + || version.version < 6; + assert!(refused, "released wire 6 must refuse the new baseline"); + write(&mut server, &Err::<(), _>(ReleasedRefusal { + terminal: true, + message: "These Maple remote protocol versions are incompatible. Update both computers.", + }), 4096).await.unwrap(); + let reply: Result<(), Refusal> = read(&mut client, MAX_PREFACE).await.unwrap(); + assert!(is_protocol_incompatibility_refusal( + &reply.unwrap_err().into_error() + )); + let mut remaining = vec![0; untouched.len()]; + server.read_exact(&mut remaining).await.unwrap(); + assert_eq!( + remaining, untouched, + "version refusal does not decode or admit queued sends" + ); } } @@ -615,9 +707,9 @@ mod tests { ); } #[test] - fn shipped_wire_examples_are_stable() { + fn wire_seven_examples_are_stable() { let examples: serde_json::Value = - serde_json::from_str(include_str!("wire_v6.json")).unwrap(); + serde_json::from_str(include_str!("wire_v7.json")).unwrap(); fn check(value: &serde_json::Value) { let decoded: T = serde_json::from_value(value.clone()).unwrap(); assert_eq!(serde_json::to_value(decoded).unwrap(), *value); @@ -631,6 +723,33 @@ mod tests { check::(&examples[key]); } check::(&examples["send"]); + check::(&examples["upload_header"]); + check::(&examples["upload"]); + for key in [ + "upload_admission_ready", + "upload_admission_complete", + "upload_admission_error", + ] { + check::>(&examples[key]); + } + check::(&examples["upload_ack"]); + check::( + &examples["upload_ack"]["Response"]["result"]["Ok"], + ); + for key in ["upload_status", "cancel_upload"] { + check::(&examples[key]); + check::<(String, crate::backend::AttachmentUploadRequest)>(&examples[key]["params"]); + } + for value in examples["upload_status_replies"].as_array().unwrap() { + check::(value); + check::(&value["Response"]["result"]["Ok"]); + } + check::(&examples["cancel_upload_reply"]); + check::<()>(&examples["cancel_upload_reply"]["Response"]["result"]["Ok"]); + assert!( + Method::parse("import_attachment").is_err(), + "local source paths never reach a remote filesystem importer" + ); check::(&examples["rename_project"]); check::( &examples["project"]["Response"]["result"]["Ok"], diff --git a/apps/maple-agent/app/src/remote/wire_v7.json b/apps/maple-agent/app/src/remote/wire_v7.json new file mode 100644 index 000000000..aeb4c4713 --- /dev/null +++ b/apps/maple-agent/app/src/remote/wire_v7.json @@ -0,0 +1,519 @@ +{ + "hello": { + "version": 7, + "minimum_version": 7, + "account": "account", + "project": "project", + "deployment": "local", + "channel": "unpackaged", + "mode": { + "kind": "control" + } + }, + "pair": { + "version": 7, + "minimum_version": 7, + "account": "account", + "project": "project", + "deployment": "local", + "channel": "unpackaged", + "mode": { + "kind": "pair", + "window": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1 + ], + "admission": "opaque-admission" + } + }, + "header": { + "id": 7, + "method": "send_message", + "send_id": "original-send" + }, + "command": { + "id": 8, + "method": "send_outcome", + "send_id": null, + "params": "original-send" + }, + "refusal": { + "Err": { + "terminal": false, + "message": "Pairing window is closed" + } + }, + "image": { + "Response": { + "id": 9, + "result": { + "Ok": { + "mime": "image/png", + "size": 2, + "data": "AP8=" + } + } + } + }, + "failure": { + "Response": { + "id": 10, + "result": { + "Err": "Remote reply is too large to transfer" + } + } + }, + "events": [ + { + "Event": { + "SessionUpdated": { + "session_id": "01234567-89ab-7def-8123-456789abcdef", + "run_id": "run", + "session": { + "id": "01234567-89ab-7def-8123-456789abcdef", + "title": "Fixture task", + "projectRoot": "/project", + "createdMs": 1, + "updatedMs": 2, + "messageCount": 1, + "model": "fixture-model", + "webEnabled": false, + "state": "active", + "acp": false + } + } + } + }, + { + "Event": { + "Run": { + "session_id": "01234567-89ab-7def-8123-456789abcdef", + "run_id": "run", + "event": { + "TimelineItem": { + "revision": 3, + "id": "reply", + "itemType": "message", + "role": "assistant", + "text": "hello", + "createdMs": 2, + "merge": "append", + "persisted": { + "taskId": "01234567-89ab-7def-8123-456789abcdef", + "rowsEpoch": 2, + "rowSeq": 7, + "updateSeq": 9 + }, + "live": { + "runId": "run", + "occurrence": 1 + } + } + } + } + } + }, + { + "Event": { + "Run": { + "session_id": "01234567-89ab-7def-8123-456789abcdef", + "run_id": "run", + "event": { + "Finished": "Completed" + } + } + } + }, + { + "Event": { + "QuestionResolved": { + "request_id": "question" + } + } + }, + { + "Event": { + "HistoryReplaced": { + "user_id": "account", + "session_id": "01234567-89ab-7def-8123-456789abcdef", + "rows_epoch": 2 + } + } + }, + { + "Event": { + "ProjectsChanged": { + "user_id": "account", + "roots": [ + { + "path": "/project", + "name": "Shared project", + "lastUsedMs": 2, + "isDefaultWorkspace": false, + "nameRevision": 1 + } + ] + } + } + }, + { + "Event": { + "ProjectRemoved": { + "user_id": "account", + "project_root": "/project" + } + } + } + ], + "send": { + "id": 11, + "method": "send_message", + "send_id": "d0c0caaa-90f4-46eb-b78a-747130b56187", + "params": [ + "account", + { + "sessionId": "01234567-89ab-7def-8123-456789abcdef", + "text": "Inspect this", + "model": "fixture-model", + "contextLimit": 32000, + "visionCapable": true, + "steer": true, + "queueId": null, + "attachments": [ + "0123456789abcdef0123456789abcdef" + ] + } + ] + }, + "receipt": { + "Response": { + "id": 13, + "result": { + "Ok": "Interrupted" + } + } + }, + "tasks": { + "Response": { + "id": 12, + "result": { + "Ok": { + "tasks": [ + { + "uuid": "01234567-89ab-7def-8123-456789abcdef", + "title": "Fixture task", + "project_root": "/project", + "state": "active", + "kind": "desktop", + "pinned_ms": null, + "web_enabled": false, + "model": "fixture-model", + "message_count": 1, + "created_ms": 1, + "updated_ms": 2, + "truncated": false + } + ] + } + } + } + }, + "resync": "Resync", + "paged_commands": [ + { + "id": 20, + "method": "list_tasks", + "send_id": null, + "params": [ + "account", + { + "section": "active", + "project": null, + "pins": "exclude", + "cursor": null, + "limit": 50 + } + ] + }, + { + "id": 21, + "method": "open_task", + "send_id": null, + "params": [ + "account", + { + "task": "01234567-89ab-7def-8123-456789abcdef", + "limitRows": 50, + "limitBytes": 524288 + } + ] + }, + { + "id": 22, + "method": "read_row", + "send_id": null, + "params": [ + "account", + { + "task": "01234567-89ab-7def-8123-456789abcdef", + "rowsEpoch": 2, + "rowSeq": 7, + "updateSeq": 9, + "offset": 0, + "len": 1024 + } + ] + }, + { + "id": 23, + "method": "read_attachment", + "send_id": null, + "params": [ + "account", + { + "task": "01234567-89ab-7def-8123-456789abcdef", + "attachment": "0123456789abcdef0123456789abcdef", + "offset": 0, + "len": 1048576 + } + ] + }, + { + "id": 24, + "method": "subscribe", + "send_id": null, + "params": [ + "account", + { + "intent": 1, + "tasks": [ + "01234567-89ab-7def-8123-456789abcdef" + ] + } + ] + } + ], + "rename_project": { + "id": 14, + "method": "rename_project_root", + "send_id": null, + "params": [ + "account", + "/project", + "Shared project" + ] + }, + "project": { + "Response": { + "id": 14, + "result": { + "Ok": { + "path": "/project", + "name": "Shared project", + "lastUsedMs": 2, + "isDefaultWorkspace": false, + "nameRevision": 1 + } + } + } + }, + "project_page": { + "projects": [ + { + "root": "/project", + "display_name": "Shared project", + "position": 0, + "last_used_ms": 2, + "truncated": false, + "name_revision": 1, + "is_default_workspace": false + } + ] + }, + "bootstrap": { + "runtimeStatus": { + "running": false, + "projectRoot": "/project", + "model": null, + "activeRuns": {} + }, + "projectRoot": "/project", + "recentProjects": [ + { + "root": "/project", + "display_name": "Shared project", + "position": 0, + "last_used_ms": 2, + "truncated": false, + "name_revision": 1, + "is_default_workspace": false + } + ], + "projectsNext": null, + "pendingQuestions": [], + "persistenceStatus": { + "store_id": null, + "revision": 0, + "failure": null + }, + "defaultProject": null + }, + "selected_project": { + "projectRoot": "/project", + "project": { + "path": "/project", + "name": "Shared project", + "lastUsedMs": 2, + "isDefaultWorkspace": false, + "nameRevision": 1 + } + }, + "upload_header": { + "id": 27, + "method": "upload_attachment", + "send_id": null + }, + "upload": { + "user": "account", + "sessionId": "01234567-89ab-7def-8123-456789abcdef", + "operationId": "11111111-1111-4111-8111-111111111111", + "metadata": { + "name": "notes.txt", + "mimeType": "text/plain", + "size": 2, + "sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + }, + "upload_admission_ready": { + "Ok": { + "state": "ready" + } + }, + "upload_admission_complete": { + "Ok": { + "state": "complete", + "attachment": { + "id": "0123456789abcdef0123456789abcdef", + "name": "notes.txt", + "mimeType": "text/plain", + "source": "maple-attachment://0123456789abcdef0123456789abcdef", + "size": 2, + "executionPath": "/store/attachments/0123456789abcdef0123456789abcdef/notes.txt" + } + } + }, + "upload_admission_error": { + "Err": "Remote access revoked" + }, + "upload_ack": { + "Response": { + "id": 27, + "result": { + "Ok": { + "id": "0123456789abcdef0123456789abcdef", + "name": "notes.txt", + "mimeType": "text/plain", + "source": "maple-attachment://0123456789abcdef0123456789abcdef", + "size": 2, + "executionPath": "/store/attachments/0123456789abcdef0123456789abcdef/notes.txt" + } + } + } + }, + "upload_status": { + "id": 28, + "method": "attachment_upload_status", + "send_id": null, + "params": [ + "account", + { + "sessionId": "01234567-89ab-7def-8123-456789abcdef", + "operationId": "11111111-1111-4111-8111-111111111111" + } + ] + }, + "upload_status_replies": [ + { + "Response": { + "id": 28, + "result": { + "Ok": { + "state": "missing" + } + } + } + }, + { + "Response": { + "id": 28, + "result": { + "Ok": { + "state": "uploading" + } + } + } + }, + { + "Response": { + "id": 28, + "result": { + "Ok": { + "state": "complete", + "attachment": { + "id": "0123456789abcdef0123456789abcdef", + "name": "notes.txt", + "mimeType": "text/plain", + "source": "maple-attachment://0123456789abcdef0123456789abcdef", + "size": 2, + "executionPath": "/store/attachments/0123456789abcdef0123456789abcdef/notes.txt" + } + } + } + } + }, + { + "Response": { + "id": 28, + "result": { + "Ok": { + "state": "cancelled" + } + } + } + } + ], + "cancel_upload": { + "id": 29, + "method": "cancel_attachment_upload", + "send_id": null, + "params": [ + "account", + { + "sessionId": "01234567-89ab-7def-8123-456789abcdef", + "operationId": "11111111-1111-4111-8111-111111111111" + } + ] + }, + "cancel_upload_reply": { + "Response": { + "id": 29, + "result": { + "Ok": null + } + } + } +} diff --git a/apps/maple-agent/app/src/ui/chat/attachments.rs b/apps/maple-agent/app/src/ui/chat/attachments.rs new file mode 100644 index 000000000..5ea53e0d7 --- /dev/null +++ b/apps/maple-agent/app/src/ui/chat/attachments.rs @@ -0,0 +1,486 @@ +//! Generic composer attachments, staged once on the controller before import. +//! Copies and image decoding stay off the UI thread; only metadata reaches it. + +use super::{ChatScreen, staging}; +use gpui::{Context, SharedString}; +use std::{path::Path, sync::Arc}; + +pub(super) const MAX_DRAFT_ATTACHMENTS: usize = 10; +const MAX_DRAFT_IMAGE_BYTES: u64 = 10 * 1024 * 1024; + +#[derive(Clone)] +pub(super) struct DraftAttachment { + pub id: u64, + pub name: String, + pub mime_type: String, + pub detail: SharedString, + pub size: u64, + pub staged: Option>, + pub operation_id: String, + pub upload_session: Option, + /// Reused only for the same execution task. A failed send never uploads a + /// second copy merely because its admission response was lost. + pub committed: Option<(String, String)>, + pub thumbnail: Option>, +} + +impl DraftAttachment { + pub(super) fn preparing(id: u64, name: String) -> Self { + Self { + id, + name, + mime_type: String::new(), + detail: "Preparing…".into(), + size: 0, + staged: None, + operation_id: uuid::Uuid::new_v4().to_string(), + upload_session: None, + committed: None, + thumbnail: None, + } + } + pub(super) fn reset_upload(&mut self) { + self.operation_id = uuid::Uuid::new_v4().to_string(); + self.committed = None; + } + + pub(super) fn prepare_upload_for(&mut self, session: &str) { + if self + .upload_session + .as_deref() + .is_some_and(|previous| previous != session) + { + self.reset_upload(); + } + self.upload_session = Some(session.to_owned()); + } + + pub(super) fn ready(&self) -> bool { + self.staged.is_some() + } + + #[cfg(test)] + pub(super) fn for_test(id: u64, name: &str, ready: bool) -> Self { + let mut draft = Self::preparing(id, name.to_string()); + if ready { + draft.staged = Some(staging::copy_bytes(&staging::root(), b"test attachment").unwrap()); + draft.mime_type = "application/octet-stream".into(); + draft.size = 15; + draft.detail = attachment_detail(&draft.mime_type, draft.size).into(); + } + draft + } +} + +impl ChatScreen { + fn remaining_attachment_slots(&mut self, cx: &mut Context) -> Option { + let remaining = MAX_DRAFT_ATTACHMENTS.saturating_sub(self.draft_attachments.len()); + if remaining == 0 { + self.notice = + Some(format!("Attach at most {MAX_DRAFT_ATTACHMENTS} files at a time").into()); + cx.notify(); + return None; + } + Some(remaining) + } + + fn image_plan_available(&self) -> bool { + self.sidebar_plan.as_ref().is_none_or(|plan| { + let label = plan.plan_label.to_lowercase(); + label.contains("pro") || label.contains("max") || label.contains("team") + }) + } + + pub(super) fn paste_image(&mut self, image: gpui::Image, cx: &mut Context) { + if self.remaining_attachment_slots(cx).is_none() { + return; + } + if !self.image_plan_available() { + self.notice = Some("Image attachments need a Pro, Max, or Team plan".into()); + cx.notify(); + return; + } + let extension = match image.format { + gpui::ImageFormat::Jpeg => "jpg", + gpui::ImageFormat::Webp => "webp", + _ => "png", + }; + let name = format!("pasted-{}.{extension}", self.draft_attachments.len() + 1); + if let Err(message) = draft_image_mime(&name, &image.bytes) { + self.notice = Some(message.into()); + cx.notify(); + return; + } + let id = self.next_draft_id(); + self.draft_attachments + .push(DraftAttachment::preparing(id, name.clone())); + self.notice = None; + self.call( + async move { + let permit = preparation_slots() + .clone() + .acquire_owned() + .await + .map_err(|_| "Attachment preparation stopped")?; + tokio::task::spawn_blocking(move || { + let _permit = permit; + let mime = draft_image_mime(&name, &image.bytes)?; + let staged = staging::copy_bytes(&staging::root(), &image.bytes)?; + let thumbnail = square_thumbnail(&image.bytes).ok().map(Arc::new); + Ok(( + staged, + mime.to_string(), + image.bytes.len() as u64, + thumbnail, + )) + }) + .await + .map_err(|e| format!("Attachment preparation failed: {e}"))? + }, + cx, + move |this, result, cx| this.finish_attachment_preparation(id, result, cx), + ); + cx.notify(); + } + + fn next_draft_id(&mut self) -> u64 { + self.draft_counter += 1; + self.draft_counter + } + + pub(super) fn finish_attachment_preparation( + &mut self, + id: u64, + result: Result, + cx: &mut Context, + ) { + let images_allowed = self.image_plan_available(); + let Some(draft) = self + .draft_attachments + .iter_mut() + .find(|draft| draft.id == id) + else { + return; + }; + match result { + Ok((staged, mime_type, size, thumbnail)) + if !mime_type.starts_with("image/") || images_allowed => + { + draft.detail = attachment_detail(&mime_type, size).into(); + draft.staged = Some(staged); + draft.mime_type = mime_type; + draft.size = size; + draft.thumbnail = thumbnail; + } + result => { + let message = result + .err() + .unwrap_or_else(|| "Image attachments need a Pro, Max, or Team plan".into()); + self.draft_attachments.retain(|draft| draft.id != id); + self.notice = Some(message.into()); + } + } + cx.notify(); + } + + pub(super) fn remove_attachment(&mut self, id: u64, cx: &mut Context) { + let Some(index) = self + .draft_attachments + .iter() + .position(|draft| draft.id == id) + else { + return; + }; + let draft = self.draft_attachments.remove(index); + if let Some(session_id) = draft.upload_session { + let backend = self.backend.clone(); + let user = self.user_id.clone(); + let operation_id = draft.operation_id; + // A restored draft can still have a pending/ready upload. Removing + // its chip explicitly cancels that operation; accepted-message + // references remain protected by the execution store. + self.call( + async move { + let _ = tokio::time::timeout( + std::time::Duration::from_secs(3), + backend.cancel_attachment_upload( + &user, + crate::backend::AttachmentUploadRequest { + session_id, + operation_id, + }, + ), + ) + .await; + Ok(()) + }, + cx, + |_, _, _| {}, + ); + } + cx.notify(); + } + + pub(super) fn pick_attachments(&mut self, cx: &mut Context) { + if self.attachment_picking || self.remaining_attachment_slots(cx).is_none() { + return; + } + self.attachment_picking = true; + self.notice = None; + cx.notify(); + let receiver = cx.prompt_for_paths(gpui::PathPromptOptions { + files: true, + directories: false, + multiple: true, + prompt: None, + }); + let bridge = cx.spawn(async move |this, cx| { + let picked = receiver.await; + this.update(cx, |this, cx| { + this.attachment_picking = false; + match picked { + Ok(Ok(Some(paths))) => this.add_attachment_paths(paths, cx), + Ok(Ok(None)) | Err(_) => {} + Ok(Err(error)) => { + this.notice = + Some(format!("Could not open the file dialog: {error}").into()) + } + } + cx.notify(); + }) + .ok(); + }); + crate::ui::task::retain(&self.bridged_tasks, bridge); + } + + pub(super) fn add_attachment_paths( + &mut self, + paths: Vec, + cx: &mut Context, + ) { + let Some(remaining) = self.remaining_attachment_slots(cx) else { + return; + }; + self.notice = if paths.len() > remaining { + Some(format!("Attach at most {MAX_DRAFT_ATTACHMENTS} files at a time").into()) + } else { + None + }; + // Reserve every slot immediately, before asynchronous work can race a + // second drop. Dropping a preparing chip also discards its late result. + let work: Vec<_> = paths + .into_iter() + .take(remaining) + .map(|path| { + let id = self.next_draft_id(); + let name = path + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| "attachment".into()); + self.draft_attachments + .push(DraftAttachment::preparing(id, name)); + (id, path) + }) + .collect(); + for (id, path) in work { + self.call( + async move { + let permit = preparation_slots() + .clone() + .acquire_owned() + .await + .map_err(|_| "Attachment preparation stopped")?; + tokio::task::spawn_blocking(move || { + let _permit = permit; + prepare_file(&path) + }) + .await + .map_err(|e| format!("Attachment preparation failed: {e}"))? + }, + cx, + move |this, result, cx| this.finish_attachment_preparation(id, result, cx), + ); + } + cx.notify(); + } +} + +fn preparation_slots() -> &'static Arc { + static SLOTS: std::sync::OnceLock> = std::sync::OnceLock::new(); + SLOTS.get_or_init(|| Arc::new(tokio::sync::Semaphore::new(2))) +} + +type PreparedAttachment = ( + Arc, + String, + u64, + Option>, +); + +fn prepare_file(path: &Path) -> Result { + use std::io::Read; + let (staged, size) = staging::copy_file(&staging::root(), path)?; + let mut signature = [0; 16]; + let mut file = std::fs::File::open(&staged.path).map_err(|e| e.to_string())?; + let read = file.read(&mut signature).map_err(|e| e.to_string())?; + let name = path.file_name().unwrap_or_default().to_string_lossy(); + if image_format_from_bytes(&signature[..read]).is_some() { + if size > MAX_DRAFT_IMAGE_BYTES { + return Err("Images must be 10 MiB or smaller".into()); + } + let bytes = std::fs::read(&staged.path).map_err(|e| e.to_string())?; + let mime = draft_image_mime(&name, &bytes)?.to_string(); + Ok(( + staged, + mime, + size, + square_thumbnail(&bytes).ok().map(Arc::new), + )) + } else { + let mime = match path + .extension() + .and_then(|ext| ext.to_str()) + .unwrap_or("") + .to_ascii_lowercase() + .as_str() + { + "zip" => "application/zip", + "pdf" => "application/pdf", + "txt" | "md" | "rs" | "py" | "js" | "ts" | "log" => "text/plain", + "json" => "application/json", + "csv" => "text/csv", + _ => "application/octet-stream", + }; + Ok((staged, mime.to_string(), size, None)) + } +} + +pub(super) fn attachment_detail(mime: &str, size: u64) -> String { + let kind = match mime { + "application/zip" => "ZIP", + "application/pdf" => "PDF", + "application/octet-stream" => "File", + mime if mime.starts_with("image/") => "Image", + mime if mime.starts_with("text/") => "Text", + _ => "File", + }; + let size = if size >= 1024 * 1024 { + format!("{:.1} MiB", size as f64 / (1024. * 1024.)) + } else if size >= 1024 { + format!("{:.0} KiB", size as f64 / 1024.) + } else { + format!("{size} B") + }; + format!("{kind} · {size}") +} + +fn draft_image_mime(name: &str, bytes: &[u8]) -> Result<&'static str, String> { + if bytes.len() as u64 > MAX_DRAFT_IMAGE_BYTES { + return Err(format!("{name} is larger than 10 MiB")); + } + match image_format_from_bytes(bytes) { + Some(gpui::ImageFormat::Png) => Ok("image/png"), + Some(gpui::ImageFormat::Jpeg) => Ok("image/jpeg"), + Some(gpui::ImageFormat::Webp) => Ok("image/webp"), + _ => Err(format!("{name} is not a PNG, JPEG, or WebP image")), + } +} + +/// Thumbnail edge in physical pixels: 2x the 64pt box so it stays sharp +/// on HiDPI screens. +const DRAFT_THUMBNAIL_PX: u32 = 128; + +/// Center-crop to a square and scale down, so the composer can show a +/// rounded square that is the picture itself. gpui clips with rectangular +/// masks only, so cropping the pixels is the one way to get round corners +/// on a cover-fit thumbnail. +fn square_thumbnail(bytes: &[u8]) -> Result { + let mut reader = image::ImageReader::new(std::io::Cursor::new(bytes)) + .with_guessed_format() + .map_err(|error| error.to_string())?; + let mut limits = image::Limits::default(); + limits.max_alloc = Some(64 * 1024 * 1024); + limits.max_image_width = Some(16_384); + limits.max_image_height = Some(16_384); + reader.limits(limits); + let decoded = reader.decode().map_err(|error| error.to_string())?; + let (width, height) = (decoded.width(), decoded.height()); + let edge = width.min(height); + if edge == 0 { + return Err("empty image".to_string()); + } + let cropped = decoded.crop_imm((width - edge) / 2, (height - edge) / 2, edge, edge); + let scaled = if edge > DRAFT_THUMBNAIL_PX { + cropped.resize_exact( + DRAFT_THUMBNAIL_PX, + DRAFT_THUMBNAIL_PX, + image::imageops::FilterType::Triangle, + ) + } else { + cropped + }; + let mut png = std::io::Cursor::new(Vec::new()); + scaled + .write_to(&mut png, image::ImageFormat::Png) + .map_err(|error| error.to_string())?; + Ok(gpui::Image::from_bytes( + gpui::ImageFormat::Png, + png.into_inner(), + )) +} + +/// Image format from the file signature; `None` for unsupported data. +pub(super) fn image_format_from_bytes(bytes: &[u8]) -> Option { + if bytes.starts_with(&[0x89, b'P', b'N', b'G']) { + Some(gpui::ImageFormat::Png) + } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) { + Some(gpui::ImageFormat::Jpeg) + } else if bytes.len() > 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" { + Some(gpui::ImageFormat::Webp) + } else { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn changing_the_draft_task_uses_a_new_upload_operation() { + let mut draft = DraftAttachment::preparing(1, "report.zip".into()); + draft.prepare_upload_for("original-task"); + let original = draft.operation_id.clone(); + draft.prepare_upload_for("original-task"); + assert_eq!(draft.operation_id, original); + draft.committed = Some(("original-task".into(), "old-file".into())); + draft.prepare_upload_for("another-task"); + assert_ne!(draft.operation_id, original); + assert!(draft.committed.is_none()); + } + + #[test] + fn zip_attachment_is_a_stable_generic_file() { + let directory = tempfile::tempdir().unwrap(); + let original = directory.path().join("report.zip"); + std::fs::write(&original, b"PK\x03\x04an archive").unwrap(); + let (staged, mime, size, thumbnail) = prepare_file(&original).unwrap(); + std::fs::remove_file(original).unwrap(); + assert_eq!(mime, "application/zip"); + assert_eq!(size, 14); + assert!(thumbnail.is_none()); + assert_eq!( + std::fs::read(&staged.path).unwrap(), + b"PK\x03\x04an archive" + ); + } + + #[test] + fn generic_file_limit_does_not_raise_the_image_limit() { + let directory = tempfile::tempdir().unwrap(); + let original = directory.path().join("large.png"); + let mut bytes = vec![0; MAX_DRAFT_IMAGE_BYTES as usize + 1]; + bytes[..8].copy_from_slice(&[0x89, b'P', b'N', b'G', 13, 10, 26, 10]); + std::fs::write(&original, bytes).unwrap(); + assert!(prepare_file(&original).err().unwrap().contains("10 MiB")); + } +} diff --git a/apps/maple-agent/app/src/ui/chat/cache.rs b/apps/maple-agent/app/src/ui/chat/cache.rs index 7d949f203..21467f469 100644 --- a/apps/maple-agent/app/src/ui/chat/cache.rs +++ b/apps/maple-agent/app/src/ui/chat/cache.rs @@ -249,10 +249,19 @@ impl MarkdownCache { /// Strings a tool or text card shows, derived once per item revision /// instead of on every frame. +#[derive(Default)] +pub(super) struct AttachmentDisplay { + pub id: String, + pub name: SharedString, + pub detail: SharedString, + pub unavailable: bool, +} + #[derive(Default)] pub(super) struct ItemDerived { /// Display text of thinking rows and user bubbles. pub(super) text: SharedString, + pub(super) attachments: Vec, /// Readable tool output for the expanded card. pub(super) output_text: Option, /// `input: {json}` for the expanded card. @@ -271,6 +280,15 @@ impl ItemDerived { let output_text = tool_output_markdown(item).map(SharedString::from); Self { text: SharedString::from(item.text.clone().unwrap_or_default()), + attachments: super::transcript::attachment_metadata(item) + .map(|attachment| AttachmentDisplay { + id: attachment.id.to_string(), + name: attachment.name.to_string().into(), + detail: super::attachments::attachment_detail(attachment.mime, attachment.size) + .into(), + unavailable: attachment.unavailable, + }) + .collect(), output_text, input_line: tool_input_line(item).map(SharedString::from), diff_lines: Arc::new(diff_lines_for(item)), diff --git a/apps/maple-agent/app/src/ui/chat/composer.rs b/apps/maple-agent/app/src/ui/chat/composer.rs index a1bc3ac12..cfb173d36 100644 --- a/apps/maple-agent/app/src/ui/chat/composer.rs +++ b/apps/maple-agent/app/src/ui/chat/composer.rs @@ -10,7 +10,7 @@ use super::cache::MarkdownKind; use super::commands::ChatCommand; use super::transcript::{render_plan_row, render_subagent_row}; use super::{ - COMPOSER_PLACEHOLDER, ChatPopup, ChatScreen, DraftImage, OpenSettingsSection, + COMPOSER_PLACEHOLDER, ChatPopup, ChatScreen, DraftAttachment, OpenSettingsSection, ROOT_MENU_RECENTS, SIDE_THREAD_PLACEHOLDER, SIDEBAR_COLLAPSED_INSET, Section, }; use crate::ui::icons::{icon, spinner}; @@ -69,7 +69,7 @@ pub(super) enum ComposerControl { Thinking, Integrations, Web, - Images, + Attachments, Dictate, } @@ -79,7 +79,7 @@ impl ComposerControl { Self::Thinking, Self::Integrations, Self::Web, - Self::Images, + Self::Attachments, Self::Dictate, ]; } @@ -254,7 +254,7 @@ impl ChatScreen { let next = !self.web_enabled; self.set_web_enabled(next, cx); } - ComposerControl::Images => self.pick_images(cx), + ComposerControl::Attachments => self.pick_attachments(cx), ComposerControl::Dictate => self.toggle_recording(cx), } } @@ -343,9 +343,9 @@ impl ChatScreen { this.activate_composer_control(ComposerControl::Web, cx); })), ), - ComposerControl::Images => div().child( + ComposerControl::Attachments => div().child( div() - .id("add-images") + .id("add-attachments") .size_8() .flex() .items_center() @@ -357,12 +357,12 @@ impl ChatScreen { .cursor_pointer() }) .active(|style| style.bg(gpui::rgb(theme::bg_sidebar_pill()))) - .when(self.image_picking, |el| el.opacity(0.5)) - .tooltip(widgets::tooltip("Attach images", None)) + .when(self.attachment_picking, |el| el.opacity(0.5)) + .tooltip(widgets::tooltip("Attach files or images", None)) .on_click(cx.listener(|this, _event, _window, cx| { - this.activate_composer_control(ComposerControl::Images, cx); + this.activate_composer_control(ComposerControl::Attachments, cx); })) - .child(icon("image", px(16.), theme::text_secondary())), + .child(icon("paperclip", px(16.), theme::text_secondary())), ), ComposerControl::Dictate => { let recording = self.recording; @@ -982,14 +982,14 @@ impl ChatScreen { "Sending…" }; let has_text = self.composer_has_text; - let has_images = !self.draft_images.is_empty(); - let images_ready = self.draft_images.iter().all(DraftImage::ready); - let can_send = !disabled && images_ready && (has_text || has_images); + let has_attachments = !self.draft_attachments.is_empty(); + let attachments_ready = self.draft_attachments.iter().all(DraftAttachment::ready); + let can_send = !disabled && attachments_ready && (has_text || has_attachments); // One control. Stop only while a run is active and nothing is - // waiting to send. Text, or an image that is still encoding, + // waiting to send. Text, or an file that is still being prepared, // turns it back into Send, which queues behind the run. - let drafting = has_text || has_images; - let show_stop = running && !disabled && !drafting; + let drafting = has_text || has_attachments; + let show_stop = running && !self.booting && !pending_create && !drafting; let queue_chips = self.render_queue(cx); let expanded = self.composer_expanded; let composer = self.composer.clone(); @@ -1042,7 +1042,10 @@ impl ChatScreen { self.render_composer_control(control, selected_control == Some(control), window, cx) }) .collect::>(); - let drafts = &self.draft_images; + let drafts = &self.draft_attachments; + let upload = self + .selected_upload() + .map(|(id, upload)| (id.to_string(), upload.label.clone())); div() .w_full() .flex() @@ -1057,7 +1060,7 @@ impl ChatScreen { .when(disabled, |container| { container.border_color(gpui::rgb(theme::border())) }) - // Files dragged from the desktop land as image attachments. + // Native files are copied locally before any execution-side upload. .can_drop(|dragged, _window, _cx| { dragged.downcast_ref::().is_some() }) @@ -1069,38 +1072,92 @@ impl ChatScreen { }) .on_drop( cx.listener(|this, paths: &gpui::ExternalPaths, _window, cx| { - this.add_image_paths(paths.paths().to_vec(), cx); + this.add_attachment_paths(paths.paths().to_vec(), cx); }), ) .children(queue_chips) + .when_some(upload, |container, (id, label)| { + container.child( + div() + .flex() + .items_center() + .gap_2() + .px_4() + .pt_3() + .child( + div() + .flex_1() + .text_xs() + .text_color(gpui::rgb(theme::text_muted())) + .child(label), + ) + .child( + widgets::secondary_button("cancel-attachment-upload") + .text_xs() + .child("Cancel") + .on_click(cx.listener(move |this, _, _, cx| { + this.cancel_attachment_send(&id, cx) + })), + ), + ) + }) .when(!drafts.is_empty(), |container| { container.child(div().flex().flex_wrap().gap_2().px_4().pt_4().children( - drafts.iter().enumerate().map(|(index, image)| { + drafts.iter().map(|attachment| { + let id = attachment.id; div() .relative() - .size_16() + .max_w(px(260.)) + .h_16() .rounded(theme::RADIUS_LG) .border_1() .border_color(gpui::rgb(theme::border())) .bg(gpui::rgb(theme::bg_elevated())) .flex() .items_center() - .justify_center() - .map(|frame| match &image.thumbnail { - // The thumbnail is already a square crop, so - // it fills the frame and the corners round. - Some(thumbnail) => frame.child( + .gap_2() + .pr_4() + .child(match &attachment.thumbnail { + Some(thumbnail) => div().size_16().flex_none().child( gpui::img(gpui::ImageSource::Image(Arc::clone(thumbnail))) .size_full() .rounded(theme::RADIUS_LG), ), - None => { - frame.child(icon("image", px(20.), theme::text_secondary())) - } + None => div() + .size_12() + .flex_none() + .flex() + .items_center() + .justify_center() + .child(icon( + if attachment.mime_type.starts_with("image/") { + "image" + } else { + "paperclip" + }, + px(20.), + theme::text_secondary(), + )), }) .child( div() - .id(gpui::SharedString::from(format!("draft-remove-{index}"))) + .min_w_0() + .flex() + .flex_col() + .gap_1() + .child( + div().truncate().text_sm().child(attachment.name.clone()), + ) + .child( + div() + .text_xs() + .text_color(gpui::rgb(theme::text_muted())) + .child(attachment.detail.clone()), + ), + ) + .child( + div() + .id(gpui::SharedString::from(format!("draft-remove-{id}"))) .absolute() .top(px(-4.)) .right(px(-4.)) @@ -1114,10 +1171,7 @@ impl ChatScreen { .border_color(gpui::rgb(theme::border())) .hover(|style| style.cursor_pointer()) .on_click(cx.listener(move |this, _event, _window, cx| { - if index < this.draft_images.len() { - this.draft_images.remove(index); - } - cx.notify(); + this.remove_attachment(id, cx); })) .child(icon("x", px(10.), theme::text_primary())), ) diff --git a/apps/maple-agent/app/src/ui/chat/computers.rs b/apps/maple-agent/app/src/ui/chat/computers.rs index e38364184..e1e3df5cd 100644 --- a/apps/maple-agent/app/src/ui/chat/computers.rs +++ b/apps/maple-agent/app/src/ui/chat/computers.rs @@ -230,7 +230,15 @@ impl ChatScreen { loading: !self.navigation_ready && self.host_recovery.events.is_some(), error: self.host_recovery.unavailable.then(|| { if self.host_recovery.ended { - "Offline — select to reconnect" + if self + .runtime_error + .as_deref() + .is_some_and(crate::remote::wire::is_protocol_incompatibility_refusal) + { + "Update required — select to retry after updating" + } else { + "Offline — select to reconnect" + } } else { "Connecting…" } @@ -414,3 +422,55 @@ impl ChatScreen { ) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::wire::{Error, Refusal}; + use gpui::TestAppContext; + + #[gpui::test] + fn catalog_marks_only_actual_protocol_refusal_as_update_required(cx: &mut TestAppContext) { + cx.executor().allow_parking(); + let cases = [ + ( + Refusal::terminal(Error::Incompatible.to_string()).into_error(), + "Update required — select to retry after updating", + ), + ( + Refusal::terminal("This computer is no longer paired").into_error(), + "Offline — select to reconnect", + ), + ( + Refusal::terminal(Error::InvalidFrame.to_string()).into_error(), + "Offline — select to reconnect", + ), + ( + Error::Disconnected.to_string(), + "Offline — select to reconnect", + ), + ]; + for (reason, expected) in cases { + let chat = { + let _guard = crate::settings::SETTINGS_IO_LOCK.lock(); + let backend = Arc::new( + AgentBackend::for_test("http://127.0.0.1:9".into(), String::new()).unwrap(), + ); + cx.new(|cx| ChatScreen::new_inner(backend, "account".into(), cx)) + }; + chat.update(cx, |chat, cx| { + chat.handle_service_event(AgentServiceEvent::RemoteAccessEnded(reason), cx); + let snapshot = chat.computer_snapshot("Office".into(), cx); + assert_eq!(snapshot.error.as_deref(), Some(expected)); + assert!(!snapshot.loading); + // The generic disconnect following a terminal refusal must + // not hide the reason the authenticated connection ended. + chat.handle_service_event(AgentServiceEvent::HostDisconnected, cx); + assert_eq!( + chat.computer_snapshot("Office".into(), cx).error.as_deref(), + Some(expected) + ); + }); + } + } +} diff --git a/apps/maple-agent/app/src/ui/chat/dialogs.rs b/apps/maple-agent/app/src/ui/chat/dialogs.rs index d6739cc94..d03e07af1 100644 --- a/apps/maple-agent/app/src/ui/chat/dialogs.rs +++ b/apps/maple-agent/app/src/ui/chat/dialogs.rs @@ -460,6 +460,21 @@ impl ChatScreen { /// Forget a task the backend no longer has. A deleted selection /// moves to the newest live task in the same root. pub(super) fn remove_session(&mut self, session_id: &str, cx: &mut Context) { + let sends: Vec<_> = self + .sending_previews + .iter() + .filter(|(_, preview)| preview.session.as_deref() == Some(session_id)) + .map(|(id, _)| id.clone()) + .collect(); + self.discarded_sends.extend(sends.iter().cloned()); + for id in sends { + self.cancel_attachment_send(&id, cx); + self.attachment_uploads.remove(&id); + self.sending_previews.remove(&id); + self.discarded_sends.insert(id); + } + self.unresolved_sends + .retain(|_, send| send.session != session_id); self.sidebar.update(cx, |sidebar, cx| { sidebar.forget_tasks(&[session_id.to_owned()], cx) }); diff --git a/apps/maple-agent/app/src/ui/chat/images.rs b/apps/maple-agent/app/src/ui/chat/images.rs deleted file mode 100644 index 99d7b6c8b..000000000 --- a/apps/maple-agent/app/src/ui/chat/images.rs +++ /dev/null @@ -1,312 +0,0 @@ -//! Image attachments for the composer: staging drafts from the file -//! picker, the clipboard, and drag-and-drop, plus the decoding helpers -//! that run off the UI thread. - -use std::sync::Arc; - -use gpui::Context; - -use super::{ChatScreen, DraftImage}; - -impl ChatScreen { - /// How many more images the draft can take, or `None` with a notice - /// set when the plan or the limit blocks attachments. - fn remaining_image_slots(&mut self, cx: &mut Context) -> Option { - if let Some(plan) = self.sidebar_plan.as_ref() { - let label = plan.plan_label.to_lowercase(); - if !(label.contains("pro") || label.contains("max") || label.contains("team")) { - self.notice = Some("Image attachments need a Pro, Max, or Team plan".into()); - cx.notify(); - return None; - } - } - let remaining = MAX_DRAFT_IMAGES.saturating_sub(self.draft_images.len()); - if remaining == 0 { - self.notice = - Some(format!("Attach at most {MAX_DRAFT_IMAGES} images at a time").into()); - cx.notify(); - return None; - } - Some(remaining) - } - - /// Stage an image pasted into the composer from the clipboard. - pub(super) fn paste_image(&mut self, image: gpui::Image, cx: &mut Context) { - if self.remaining_image_slots(cx).is_none() { - return; - } - let extension = match image.format { - gpui::ImageFormat::Jpeg => "jpg", - gpui::ImageFormat::Webp => "webp", - _ => "png", - }; - let name = format!("pasted-{}.{extension}", self.draft_images.len() + 1); - match draft_image_mime(&name, &image.bytes) { - Ok(mime) => { - let id = self.next_draft_id(); - self.notice = None; - self.draft_images.push(DraftImage { - id, - name, - data_url: None, - thumbnail: None, - }); - self.prepare_draft_image(id, mime, image.bytes, cx); - } - Err(message) => self.notice = Some(message.into()), - } - cx.notify(); - } - - fn next_draft_id(&mut self) -> u64 { - self.draft_counter += 1; - self.draft_counter - } - - /// Encode the data URL and center-crop the thumbnail off the UI - /// thread, then attach both to the draft with `id` if it is still - /// staged. A pasted image can be 10 MB; encoding it inline stalled - /// the frame. - fn prepare_draft_image( - &mut self, - id: u64, - mime: &'static str, - bytes: Vec, - cx: &mut Context, - ) { - self.call( - async move { - tokio::task::spawn_blocking(move || { - let data_url = encode_data_url(mime, &bytes); - let thumbnail = square_thumbnail(&bytes); - Ok::<_, String>((data_url, thumbnail)) - }) - .await - .map_err(|error| format!("Image task failed: {error}"))? - }, - cx, - move |this, result, cx| { - let Some(draft) = this.draft_images.iter_mut().find(|draft| draft.id == id) else { - return; - }; - match result { - Ok((data_url, thumbnail)) => { - draft.data_url = Some(data_url); - match thumbnail { - Ok(image) => draft.thumbnail = Some(Arc::new(image)), - Err(message) => { - log::debug!("thumbnail for {}: {message}", draft.name) - } - } - } - Err(message) => { - // Without a data URL the draft can never be sent. - let name = std::mem::take(&mut draft.name); - this.draft_images.retain(|draft| draft.id != id); - this.notice = Some(format!("Could not prepare {name}: {message}").into()); - } - } - cx.notify(); - }, - ); - } - - pub(super) fn pick_images(&mut self, cx: &mut Context) { - if self.image_picking { - return; - } - if self.remaining_image_slots(cx).is_none() { - return; - } - self.image_picking = true; - self.notice = None; - cx.notify(); - // The platform file picker through gpui, native on every OS. It - // has no file-type filter; `add_image_paths` re-checks the limit - // and skips non-image files with a notice. - let receiver = cx.prompt_for_paths(gpui::PathPromptOptions { - files: true, - directories: false, - multiple: true, - prompt: None, - }); - let bridge = cx.spawn(async move |this, cx| { - let picked = receiver.await; - this.update(cx, |this, cx| { - this.image_picking = false; - match picked { - Ok(Ok(Some(paths))) => this.add_image_paths(paths, cx), - // Cancelled, or the picker dropped its channel. - Ok(Ok(None)) | Err(_) => {} - Ok(Err(error)) => { - this.notice = - Some(format!("Could not open the file dialog: {error}").into()); - } - } - cx.notify(); - }) - .ok(); - }); - crate::ui::task::retain(&self.bridged_tasks, bridge); - } - - /// Stage images dropped onto the composer. Files that are not PNG, - /// JPEG, or WebP are skipped with a notice. - pub(super) fn add_image_paths( - &mut self, - paths: Vec, - cx: &mut Context, - ) { - let (images, other): (Vec<_>, Vec<_>) = paths.into_iter().partition(|path| { - path.extension() - .and_then(|extension| extension.to_str()) - .is_some_and(|extension| { - matches!( - extension.to_ascii_lowercase().as_str(), - "png" | "jpg" | "jpeg" | "webp" - ) - }) - }); - if !other.is_empty() { - self.notice = Some( - format!( - "Only PNG, JPEG, and WebP images can be attached ({} file{} skipped)", - other.len(), - if other.len() == 1 { "" } else { "s" } - ) - .into(), - ); - cx.notify(); - } - if images.is_empty() { - return; - } - let Some(remaining) = self.remaining_image_slots(cx) else { - return; - }; - self.call( - async move { - tokio::task::spawn_blocking(move || { - images - .iter() - .take(remaining) - .map(|path| load_draft_image(path)) - .collect::, String>>() - }) - .await - .map_err(|error| format!("Image load failed: {error}"))? - }, - cx, - |this, result, cx| { - match result { - Ok(images) => { - for mut image in images { - image.id = this.next_draft_id(); - this.draft_images.push(image); - } - } - Err(message) => this.notice = Some(message.into()), - } - cx.notify(); - }, - ); - } -} - -pub(super) const MAX_DRAFT_IMAGES: usize = 10; -const MAX_DRAFT_IMAGE_BYTES: usize = 10 * 1024 * 1024; - -/// Read an image file into a draft with its thumbnail, checking size and -/// format the same way the runtime does so errors surface before the send. -/// Runs on a blocking thread; the caller assigns the id. -fn load_draft_image(path: &std::path::Path) -> Result { - let name = path - .file_name() - .map(|n| n.to_string_lossy().to_string()) - .unwrap_or_else(|| "image".to_string()); - let bytes = std::fs::read(path).map_err(|error| format!("Could not read {name}: {error}"))?; - let mime = draft_image_mime(&name, &bytes)?; - Ok(DraftImage { - id: 0, - name, - data_url: Some(encode_data_url(mime, &bytes)), - thumbnail: square_thumbnail(&bytes).ok().map(Arc::new), - }) -} - -/// Check size and format the same way the runtime does, so errors -/// surface before the send. Cheap: reads the signature only. -fn draft_image_mime(name: &str, bytes: &[u8]) -> Result<&'static str, String> { - if bytes.len() > MAX_DRAFT_IMAGE_BYTES { - return Err(format!("{name} is larger than 10 MB")); - } - if bytes.starts_with(&[0x89, b'P', b'N', b'G']) { - Ok("image/png") - } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) { - Ok("image/jpeg") - } else if bytes.len() > 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" { - Ok("image/webp") - } else { - Err(format!("{name} is not a PNG, JPEG, or WebP image")) - } -} - -/// The `data:` URL the runtime stores with the message. Runs on a -/// blocking thread: the payload can be 10 MB. -pub(super) fn encode_data_url(mime: &str, bytes: &[u8]) -> Arc { - use base64::Engine as _; - let mut url = String::with_capacity(mime.len() + 16 + bytes.len().div_ceil(3) * 4); - url.push_str("data:"); - url.push_str(mime); - url.push_str(";base64,"); - base64::engine::general_purpose::STANDARD.encode_string(bytes, &mut url); - Arc::from(url) -} - -/// Thumbnail edge in physical pixels: 2x the 64pt box so it stays sharp -/// on HiDPI screens. -const DRAFT_THUMBNAIL_PX: u32 = 128; - -/// Center-crop to a square and scale down, so the composer can show a -/// rounded square that is the picture itself. gpui clips with rectangular -/// masks only, so cropping the pixels is the one way to get round corners -/// on a cover-fit thumbnail. -fn square_thumbnail(bytes: &[u8]) -> Result { - let decoded = image::load_from_memory(bytes).map_err(|error| error.to_string())?; - let (width, height) = (decoded.width(), decoded.height()); - let edge = width.min(height); - if edge == 0 { - return Err("empty image".to_string()); - } - let cropped = decoded.crop_imm((width - edge) / 2, (height - edge) / 2, edge, edge); - let scaled = if edge > DRAFT_THUMBNAIL_PX { - cropped.resize_exact( - DRAFT_THUMBNAIL_PX, - DRAFT_THUMBNAIL_PX, - image::imageops::FilterType::Triangle, - ) - } else { - cropped - }; - let mut png = std::io::Cursor::new(Vec::new()); - scaled - .write_to(&mut png, image::ImageFormat::Png) - .map_err(|error| error.to_string())?; - Ok(gpui::Image::from_bytes( - gpui::ImageFormat::Png, - png.into_inner(), - )) -} - -/// Image format from the file signature; `None` for unsupported data. -pub(super) fn image_format_from_bytes(bytes: &[u8]) -> Option { - if bytes.starts_with(&[0x89, b'P', b'N', b'G']) { - Some(gpui::ImageFormat::Png) - } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) { - Some(gpui::ImageFormat::Jpeg) - } else if bytes.len() > 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" { - Some(gpui::ImageFormat::Webp) - } else { - None - } -} diff --git a/apps/maple-agent/app/src/ui/chat/mod.rs b/apps/maple-agent/app/src/ui/chat/mod.rs index a5331a4e9..00f54cbcd 100644 --- a/apps/maple-agent/app/src/ui/chat/mod.rs +++ b/apps/maple-agent/app/src/ui/chat/mod.rs @@ -10,7 +10,7 @@ use gpui::{ Window, div, prelude::*, px, }; use maple_agent::agent::{ - AgentCreateSessionRequest, AgentImageUpload, AgentIntegration, AgentIntegrationAvailability, + AgentCreateSessionRequest, AgentIntegration, AgentIntegrationAvailability, AgentIntegrationBackend, AgentMcpServer, AgentMcpTransport, AgentProjectTrustStatus, AgentQueuedMessage, AgentSendMessageRequest, AgentServiceEvent, AgentSessionIntegrationKind, AgentSessionMcpServer, AgentSessionSummary, AgentSlashCommand, AgentSubagent, AgentTaskState, @@ -34,8 +34,8 @@ mod commands; mod composer; pub(crate) mod computers; use computers::{Computer, Destination, Navigate, NavigationChanged}; +mod attachments; mod dialogs; -mod images; mod navigation; mod paging; mod project_folder; @@ -43,11 +43,15 @@ mod queue; mod sending; mod sidebar; mod speech; +mod staging; mod summaries; #[cfg(test)] mod tests; mod thinking; mod transcript; +#[cfg(test)] +mod upload_tests; +mod uploads; use self::cache::{DerivedCache, INLINE_PARSE_LIMIT, MarkdownCache, MarkdownKind}; use self::commands::ChatCommand; @@ -168,32 +172,14 @@ const PROMPT_HISTORY_LIMIT: usize = 50; /// elapsed time. It runs only while a subagent works. const SUBAGENT_TICK: std::time::Duration = std::time::Duration::from_secs(1); -/// An image staged in the composer: the data URL the runtime stores with -/// the message and a square thumbnail, both built off the UI thread. -#[derive(Clone)] -struct DraftImage { - /// Unique per staged draft, so a late thumbnail finds its owner even - /// after other drafts were removed. - id: u64, - name: String, - /// `None` while the base64 encode is still running; the send waits - /// for it. Shared so the send does not copy up to 13 MB per image. - data_url: Option>, - /// `None` until the crop finishes (or if the image did not decode). - thumbnail: Option>, -} - -impl DraftImage { - fn ready(&self) -> bool { - self.data_url.is_some() - } -} +use self::attachments::DraftAttachment; /// Shared read-only state the transcript rows render from. struct TranscriptCtx<'a> { markdown_cache: &'a MarkdownCache, derived: &'a DerivedCache, attachment_images: &'a HashMap>, + attachment_failures: &'a HashSet, chat: &'a gpui::WeakEntity, tool_summaries: &'a HashMap, render: &'a RenderCtx, @@ -237,7 +223,7 @@ pub(super) enum ChatPopup { /// showing the text meanwhile; it clears when the send goes out. #[derive(Debug, Clone, PartialEq, Eq)] enum FirstSend { - /// Plain text (with the staged images) for the new task. + /// Plain text (with the staged attachments) for the new task. Message { text: String, steer: bool }, /// A `/command` that needs a task, as typed. Command(String), @@ -387,11 +373,19 @@ fn draft_mcp_rows( rows } +struct SendReceiptQuery { + id: String, + outcome: Option, + backend: Arc, +} + +#[derive(Clone)] struct UnresolvedSend { + backend: Arc, outcome: Option, session: String, text: String, - images: Vec, + attachments: Vec, from_composer: bool, } @@ -576,14 +570,17 @@ pub struct ChatScreen { watched_git_dir: Option, /// Sidebar hidden; a toggle in the main pane brings it back. sidebar_collapsed: bool, - /// Images staged for the next message. - draft_images: Vec, + /// Files staged for the next message. + draft_attachments: Vec, unresolved_sends: HashMap, + attachment_uploads: HashMap, + accept_send_results: bool, + discarded_sends: HashSet, checking_sends: bool, - /// Source of `DraftImage::id`. + /// Source of `DraftAttachment::id`. draft_counter: u64, /// True while the native file dialog is open. - image_picking: bool, + attachment_picking: bool, /// The catalog's vision and context window per model id, looked up /// when a model is chosen. model_catalog: HashMap, @@ -652,6 +649,7 @@ pub struct ChatScreen { /// Attachment ids with a read in flight or already failed; never asked /// for twice. attachment_requests: HashSet, + attachment_failures: HashSet, attachment_queue: std::collections::VecDeque, pending_attachments: usize, attachment_viewport: std::ops::Range, @@ -1283,11 +1281,14 @@ impl ChatScreen { dialog_focus_pending: false, trust_prompts: true, sidebar_collapsed: false, - draft_images: Vec::new(), + draft_attachments: Vec::new(), unresolved_sends: HashMap::new(), + attachment_uploads: HashMap::new(), + accept_send_results: true, + discarded_sends: HashSet::new(), checking_sends: false, draft_counter: 0, - image_picking: false, + attachment_picking: false, model_catalog: HashMap::new(), model_labels: HashMap::new(), session_mcp: Vec::new(), @@ -1313,6 +1314,7 @@ impl ChatScreen { reload_generation: 0, attachment_images: HashMap::new(), attachment_requests: HashSet::new(), + attachment_failures: HashSet::new(), attachment_queue: std::collections::VecDeque::new(), pending_attachments: 0, attachment_viewport: 0..0, @@ -1511,6 +1513,7 @@ impl ChatScreen { } } this.booting = false; + this.dispatch_prepared_uploads(cx); cx.notify(); this.refresh_models(cx); this.refresh_roots(cx); @@ -2200,7 +2203,7 @@ impl ChatScreen { return; }; self.first_send_preview = matches!(action, FirstSend::Message { .. }).then(|| { - sending::SendPreview::new(None, &action.text(), &self.draft_images).status_only() + sending::SendPreview::new(None, &action.text(), &self.draft_attachments).status_only() }); self.pending_first_send = Some(action); cx.notify(); @@ -3143,6 +3146,7 @@ impl ChatScreen { self.tool_summaries.clear(); self.attachment_images.clear(); self.attachment_requests.clear(); + self.attachment_failures.clear(); self.attachment_queue.clear(); self.attachment_viewport = 0..0; self.toggled_tools.clear(); @@ -4378,12 +4382,12 @@ impl ChatScreen { cx.notify(); return; } - if text.trim().is_empty() && self.draft_images.is_empty() { + if text.trim().is_empty() && self.draft_attachments.is_empty() { return; } - if !self.draft_images.iter().all(DraftImage::ready) { - // The encode is still on its way; keep the text for the retry. - self.notice = Some("Images are still being prepared; try again in a moment".into()); + if !self.draft_attachments.iter().all(DraftAttachment::ready) { + // The stable copy is still being prepared; keep the text for the retry. + self.notice = Some("Files are still being prepared; try again in a moment".into()); if let Some(composer) = self.composer.clone() { composer.update(cx, |input, cx| input.set_text(&text, cx)); } @@ -4672,6 +4676,17 @@ impl ChatScreen { queue_id: Option, cx: &mut Context, ) -> bool { + if queue_id.is_none() + && !self.draft_attachments.is_empty() + && self + .attachment_uploads + .values() + .any(|upload| upload.session == session_id && !upload.fence.handed_off()) + { + self.notice = Some("Wait for the current file upload or cancel it before sending more files. Text can still be queued or sent to the running task.".into()); + cx.notify(); + return false; + } if self .sending_previews .values() @@ -4694,7 +4709,6 @@ impl ChatScreen { return false; } let send_id = uuid::Uuid::new_v4().to_string(); - let outgoing_id = send_id.clone(); let session_id = session_id.to_string(); let backend = self.backend.clone(); let user_id = self.user_id.clone(); @@ -4702,17 +4716,20 @@ impl ChatScreen { let vision_capable = self.selected_model_supports_vision(); let context_limit = self.selected_model_context_window(); let run_active = self.active_runs.contains_key(&session_id); - // A queued chip keeps its own attachments; new images stay staged. - let drafts = if queue_id.is_some() { + // A queued chip keeps its own attachments; new files stay staged. + let mut drafts = if queue_id.is_some() { Vec::new() } else { - std::mem::take(&mut self.draft_images) + std::mem::take(&mut self.draft_attachments) }; + for draft in &mut drafts { + draft.prepare_upload_for(&session_id); + } let editing = self .queue_edit .as_ref() .is_some_and(|edit| Some(&edit.queue_id) == queue_id.as_ref()); - let request = AgentSendMessageRequest { + let mut request = AgentSendMessageRequest { session_id: session_id.clone(), // A steer identifies the queued item atomically. A fetched preview // must never overwrite another client's intervening queue edit. @@ -4726,16 +4743,7 @@ impl ChatScreen { vision_capable, steer: steer && (run_active || queue_id.is_some()), queue_id, - attachments: drafts - .iter() - .filter_map(|image| { - let data_url = image.data_url.as_deref()?; - Some(AgentImageUpload { - name: image.name.clone(), - data_url: data_url.to_string(), - }) - }) - .collect(), + attachments: Vec::new(), }; self.notice = None; @@ -4760,6 +4768,11 @@ impl ChatScreen { } let preview = sending::SendPreview::new(Some(session_id.clone()), &text, &drafts) .capture_timeline(&self.timeline); + let preview = if drafts.is_empty() { + preview + } else { + preview.uploading() + }; self.sending_previews.insert( send_id.clone(), if run_active || request.queue_id.is_some() { @@ -4769,53 +4782,121 @@ impl ChatScreen { }, ); cx.notify(); + let retained_send = UnresolvedSend { + backend: backend.clone(), + outcome: None, + session: session_id.clone(), + text: text.clone(), + attachments: drafts.clone(), + from_composer, + }; + let (upload_fence, progress) = self.begin_uploads(&send_id, &session_id, &drafts, cx); + if drafts.is_empty() { + self.admit_prepared_send( + send_id, + uploads::PreparedSend { + request, + send: retained_send, + fence: upload_fence, + }, + cx, + ); + return true; + } self.call( async move { - backend - .send_message_identified(&user_id, request, outgoing_id) - .await - }, - cx, - move |this, result, cx| { - this.finish_send( - send_id, + let result = match uploads::commit_attachments( + &backend, + &user_id, + &session_id, + &mut drafts, + &upload_fence, + progress, + ) + .await + { + Ok(ids) => { + request.attachments = ids; + Ok(()) + } + Err(error) => Err(error), + }; + Ok(( UnresolvedSend { + backend: backend.clone(), outcome: None, session: session_id, text, - images: drafts, + attachments: drafts, from_composer, }, result, - cx, - ); + request, + upload_fence, + )) + }, + cx, + move |this, result, cx| { + match result { + Ok((send, Ok(()), request, fence)) + if this.accept_send_results && !this.discarded_sends.remove(&send_id) => + { + if let Some(preview) = this.sending_previews.get_mut(&send_id) { + preview.prepare(uploads::PreparedSend { + request, + send, + fence, + }); + this.dispatch_prepared_uploads(cx); + } + } + Ok((send, Err(error), _, _)) => this.finish_attachment_send( + send_id, + retained_send, + Ok((send, Err(error))), + cx, + ), + Err(error) => { + this.finish_attachment_send(send_id, retained_send, Err(error), cx) + } + Ok(_) => {} + } + cx.notify(); }, ); true } + fn pending_receipt_queries(&self) -> Vec { + self.unresolved_sends + .iter() + .map(|(id, send)| SendReceiptQuery { + id: id.clone(), + outcome: send + .outcome + .clone() + .filter(|o| !matches!(o, crate::host::receipts::Outcome::Pending)), + backend: send.backend.clone(), + }) + .collect() + } + fn check_unresolved_sends(&mut self, cx: &mut Context) { if self.checking_sends || self.unresolved_sends.is_empty() || self.host_recovery.unavailable { return; } self.checking_sends = true; - let ids: Vec<_> = self - .unresolved_sends - .iter() - .map(|(id, send)| { - let settled = send - .outcome - .clone() - .filter(|o| !matches!(o, crate::host::receipts::Outcome::Pending)); - (id.clone(), settled) - }) - .collect(); - let backend = self.backend.clone(); + let queries = self.pending_receipt_queries(); self.call( async move { let mut results = Vec::new(); - for (id, settled) in ids { + for SendReceiptQuery { + id, + outcome: settled, + backend, + } in queries + { if settled.is_some() { results.push((id, settled)); continue; @@ -4848,6 +4929,16 @@ impl ChatScreen { id: &str, outcome: Option, cx: &mut Context, + ) { + self.apply_send_receipt_result(id, outcome, cx); + self.dispatch_prepared_uploads(cx); + } + + fn apply_send_receipt_result( + &mut self, + id: &str, + outcome: Option, + cx: &mut Context, ) { use crate::host::receipts::Outcome; let Some(send) = self.unresolved_sends.get_mut(id) else { @@ -4875,13 +4966,13 @@ impl ChatScreen { .composer .as_ref() .is_some_and(|c| c.read(cx).text_ref().trim().is_empty()) - && self.draft_images.is_empty(); + && self.draft_attachments.is_empty(); if self.selected_session.as_deref() == Some(&send.session) && empty { let send = self.unresolved_sends.remove(id).unwrap(); if let Some(composer) = self.composer.clone() { composer.update(cx, |input, cx| input.set_text(&send.text, cx)); } - self.draft_images = send.images; + self.draft_attachments = send.attachments; self.notice = Some(if matches!(outcome, Some(Outcome::Interrupted)) { "The host restarted before recording this send’s result. It may have taken effect. Your draft is restored; check the task before sending again." } else { @@ -4896,7 +4987,7 @@ impl ChatScreen { } } _ => { - self.notice = Some("Send outcome is still unknown. Your message and images are retained; Maple will check again after reconnecting.".into()); + self.notice = Some("Send outcome is still unknown. Your message and files are retained; Maple will check again after reconnecting.".into()); } } } @@ -5215,6 +5306,17 @@ impl ChatScreen { } pub(crate) fn sign_out(&mut self, cx: &mut Context) { + // A logout during upload must win before prompt admission, even while + // the host is still shutting down an earlier run. + self.accept_send_results = false; + let uploads: Vec<_> = self.attachment_uploads.keys().cloned().collect(); + for id in uploads { + self.cancel_attachment_send(&id, cx); + } + self.attachment_uploads.clear(); + self.draft_attachments.clear(); + self.unresolved_sends.clear(); + self.sending_previews.clear(); self.stop_speech(cx); if self.recording { self.recording = false; @@ -5730,6 +5832,7 @@ impl ChatScreen { this.runtime_error = None; this.booting = !running && !this.navigation_only; this.check_unresolved_sends(cx); + this.dispatch_prepared_uploads(cx); cx.emit(NavigationChanged); if !this.navigation_only { this.refresh_models(cx); @@ -5739,7 +5842,8 @@ impl ChatScreen { if this.booting { // A host another client restarted may not have a // runtime yet. Restore catalog/tool availability; - // this does not resume or submit any prompt. + // this never replays an admitted prompt. Prepared + // uploads still wait for this readiness result. this.start_runtime(cx); } this.finish_navigation_bootstrap(cx); diff --git a/apps/maple-agent/app/src/ui/chat/paging.rs b/apps/maple-agent/app/src/ui/chat/paging.rs index 7758d04d0..c0a5f91b6 100644 --- a/apps/maple-agent/app/src/ui/chat/paging.rs +++ b/apps/maple-agent/app/src/ui/chat/paging.rs @@ -227,6 +227,7 @@ impl ChatScreen { if changed { self.attachment_images.clear(); self.attachment_requests.clear(); + self.attachment_failures.clear(); self.attachment_queue.clear(); } } diff --git a/apps/maple-agent/app/src/ui/chat/sending.rs b/apps/maple-agent/app/src/ui/chat/sending.rs index 102e87745..9768eebcd 100644 --- a/apps/maple-agent/app/src/ui/chat/sending.rs +++ b/apps/maple-agent/app/src/ui/chat/sending.rs @@ -6,28 +6,46 @@ use super::*; pub(super) struct SendPreview { pub session: Option, text: SharedString, - images: Vec<(SharedString, Option>)>, + attachments: Vec<(SharedString, SharedString, Option>)>, show_body: bool, acknowledged_run: Option, projected_runs: HashSet, baseline_user_rows: HashSet, snapshot_projected: bool, + phase: SendPhase, +} + +enum SendPhase { + Uploading, + Prepared(Box), + Admitting, } impl SendPreview { - pub(super) fn new(session: Option, text: &str, images: &[DraftImage]) -> Self { + pub(super) fn new( + session: Option, + text: &str, + attachments: &[DraftAttachment], + ) -> Self { Self { session, text: text.to_string().into(), - images: images + attachments: attachments .iter() - .map(|image| (image.name.clone().into(), image.thumbnail.clone())) + .map(|attachment| { + ( + attachment.name.clone().into(), + attachment.detail.clone(), + attachment.thumbnail.clone(), + ) + }) .collect(), show_body: true, acknowledged_run: None, projected_runs: HashSet::new(), baseline_user_rows: HashSet::new(), snapshot_projected: false, + phase: SendPhase::Admitting, } } @@ -46,7 +64,33 @@ impl SendPreview { } pub(super) fn admitting(&self) -> bool { - self.acknowledged_run.is_none() + matches!(self.phase, SendPhase::Admitting) && self.acknowledged_run.is_none() + } + + pub(super) fn uploading(mut self) -> Self { + self.phase = SendPhase::Uploading; + self + } + + pub(super) fn prepare(&mut self, send: uploads::PreparedSend) { + self.phase = SendPhase::Prepared(Box::new(send)); + } + + pub(super) fn prepared(&self) -> Option<&uploads::PreparedSend> { + match &self.phase { + SendPhase::Prepared(send) => Some(send), + _ => None, + } + } + + pub(super) fn take_prepared(&mut self) -> Option { + if !matches!(self.phase, SendPhase::Prepared(_)) { + return None; + } + match std::mem::replace(&mut self.phase, SendPhase::Admitting) { + SendPhase::Prepared(send) => Some(*send), + _ => unreachable!(), + } } pub(super) fn body_visible(&self) -> bool { @@ -192,7 +236,7 @@ impl ChatScreen { self.awaiting_first_token = false; } if uncertain { - self.notice = Some("Delivery is uncertain. Your message and images are retained while Maple checks the host; nothing will be resent.".into()); + self.notice = Some("Delivery is uncertain. Your message and files are retained while Maple checks the host; nothing will be resent.".into()); self.check_unresolved_sends(cx); } else { if self.is_selected(&session) { @@ -206,6 +250,7 @@ impl ChatScreen { } } } + self.dispatch_prepared_uploads(cx); cx.notify(); } @@ -273,16 +318,26 @@ impl ChatScreen { .border_1() .border_color(gpui::rgb(theme::user_bubble_border())) .text_color(gpui::rgb(theme::text_primary())) - .when(!preview.images.is_empty(), |body| { + .when(!preview.attachments.is_empty(), |body| { body.child(div().flex().flex_wrap().gap_2().mb_1().children( - preview.images.iter().map(|(name, image)| { + preview.attachments.iter().map(|(name, detail, image)| { match image { Some(image) => div().child( gpui::img(gpui::ImageSource::Image(image.clone())) .size_16() .rounded(theme::RADIUS_SM), ), - None => div().text_xs().child(name.clone()), + None => div() + .flex() + .flex_col() + .text_xs() + .gap_1() + .child(name.clone()) + .child( + div() + .text_color(gpui::rgb(theme::text_muted())) + .child(detail.clone()), + ), } }), )) @@ -302,11 +357,17 @@ impl ChatScreen { px(12.), theme::text_muted(), )) - .child(if self.send_pending() { - "Sending…" - } else { - "Waiting for response…" - }), + .child( + self.selected_upload() + .map(|(_, upload)| upload.label.clone()) + .unwrap_or_else(|| { + if self.send_pending() { + "Sending…".into() + } else { + "Waiting for response…".into() + } + }), + ), ), ) } diff --git a/apps/maple-agent/app/src/ui/chat/staging.rs b/apps/maple-agent/app/src/ui/chat/staging.rs new file mode 100644 index 000000000..fa82959fb --- /dev/null +++ b/apps/maple-agent/app/src/ui/chat/staging.rs @@ -0,0 +1,419 @@ +//! Controller-local stable attachment copies. No staged path is interpreted by +//! a remote execution host. Disk work is called only on a blocking worker. + +use std::{ + fs::{self, File, OpenOptions}, + io::{Read, Write}, + path::{Path, PathBuf}, + sync::{ + Arc, OnceLock, + atomic::{AtomicU64, AtomicUsize, Ordering}, + mpsc, + }, + time::{Duration, SystemTime}, +}; + +pub(super) const MAX_FILE_BYTES: u64 = 30 * 1024 * 1024; +/// Crashed processes cannot replay drafts. Reclaim their abandoned copies after +/// seven days; an OS lock excludes every still-live window/process. +const STALE_AFTER: Duration = Duration::from_secs(7 * 24 * 60 * 60); +const MAX_SWEEP_ENTRIES: usize = 512; + +const MAX_LIVE_FILES: usize = 64; +const MAX_LIVE_BYTES: u64 = 300 * 1024 * 1024; + +#[derive(Default)] +struct StagingBudget { + files: AtomicUsize, + bytes: AtomicU64, +} + +struct Reservation { + budget: Arc, + bytes: u64, +} + +impl StagingBudget { + fn reserve(self: &Arc, bytes: u64) -> Result { + self.files + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |files| { + (files < MAX_LIVE_FILES).then_some(files + 1) + }) + .map_err( + |_| "Too many retained attachments. Remove some draft files before attaching more.", + )?; + if self + .bytes + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |used| { + used.checked_add(bytes) + .filter(|next| *next <= MAX_LIVE_BYTES) + }) + .is_err() + { + self.files.fetch_sub(1, Ordering::AcqRel); + return Err("Retained attachments exceed 300 MiB. Remove some draft files before attaching more.".into()); + } + Ok(Reservation { + budget: self.clone(), + bytes, + }) + } +} + +impl Drop for Reservation { + fn drop(&mut self) { + self.budget.bytes.fetch_sub(self.bytes, Ordering::AcqRel); + self.budget.files.fetch_sub(1, Ordering::AcqRel); + } +} + +fn budget() -> &'static Arc { + static BUDGET: OnceLock> = OnceLock::new(); + BUDGET.get_or_init(|| Arc::new(StagingBudget::default())) +} + +pub(super) struct StagedFile { + pub path: PathBuf, + owner: Option, + _reservation: Reservation, +} + +impl Drop for StagedFile { + fn drop(&mut self) { + // Last-drop can happen during a GPUI update. Never remove files there. + self.owner.take(); + if let Some(directory) = self.path.parent() + && let Some(sender) = cleanup_sender() + { + let _ = sender.send(directory.to_path_buf()); + } + } +} + +fn cleanup_sender() -> &'static Option> { + static SENDER: OnceLock>> = OnceLock::new(); + SENDER.get_or_init(|| { + let (sender, receiver) = mpsc::channel::(); + std::thread::Builder::new() + .name("attachment-cleanup".into()) + .spawn(move || { + for directory in receiver { + remove_staged_directory(&directory); + } + }) + .ok()?; + Some(sender) + }) +} + +fn remove_staged_directory(directory: &Path) { + // Remove only our known files, never recurse through unknown entries. + let _ = fs::remove_file(directory.join("payload")); + let _ = fs::remove_file(directory.join("owner")); + let _ = fs::remove_dir(directory); +} + +pub(super) fn root() -> PathBuf { + #[cfg(not(test))] + return crate::backend::profile_data_root().join("draft-attachments"); + #[cfg(test)] + { + static ROOT: OnceLock = OnceLock::new(); + ROOT.get_or_init(|| tempfile::tempdir().unwrap()) + .path() + .join("draft-attachments") + } +} + +fn private_directory(path: &Path) -> Result<(), String> { + let builder = fs::DirBuilder::new(); + #[cfg(unix)] + let builder = { + use std::os::unix::fs::DirBuilderExt; + let mut builder = builder; + builder.mode(0o700); + builder + }; + match builder.create(path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + let metadata = fs::symlink_metadata(path).map_err(|e| e.to_string())?; + if !metadata.is_dir() || metadata.file_type().is_symlink() { + return Err("Attachment staging directory is unsafe".into()); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + if metadata.permissions().mode() & 0o077 != 0 { + return Err("Attachment staging directory is not private".into()); + } + } + Ok(()) + } + Err(error) => Err(format!("Could not prepare attachment staging: {error}")), + } +} + +fn file_options(create: bool) -> OpenOptions { + let mut options = OpenOptions::new(); + options.read(true).write(true).create_new(create); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options + .mode(0o600) + .custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC); + } + options +} + +fn sweep(root: &Path, now: SystemTime) { + let Ok(entries) = fs::read_dir(root) else { + return; + }; + for entry in entries.take(MAX_SWEEP_ENTRIES).flatten() { + if uuid::Uuid::parse_str(&entry.file_name().to_string_lossy()).is_err() { + continue; + } + let path = entry.path(); + let Ok(metadata) = fs::symlink_metadata(&path) else { + continue; + }; + if !metadata.is_dir() || metadata.file_type().is_symlink() { + continue; + } + let stale = metadata + .modified() + .ok() + .and_then(|time| now.duration_since(time).ok()) + .is_some_and(|age| age >= STALE_AFTER); + if !stale { + continue; + } + let Ok(owner) = file_options(false).open(path.join("owner")) else { + continue; + }; + if owner.try_lock().is_err() { + continue; + } + drop(owner); + remove_staged_directory(&path); + } +} + +fn begin(root: &Path, reservation: Reservation) -> Result<(Arc, File), String> { + // The profile root is already owned by Maple; tests use an isolated parent. + // Create the staging child with its own private permissions in both cases. + if let Some(parent) = root.parent() { + fs::create_dir_all(parent).map_err(|e| e.to_string())?; + } + private_directory(root)?; + sweep(root, SystemTime::now()); + let directory = root.join(uuid::Uuid::new_v4().to_string()); + private_directory(&directory)?; + let mut staged = StagedFile { + path: directory.join("payload"), + owner: None, + _reservation: reservation, + }; + let owner = file_options(true) + .open(directory.join("owner")) + .map_err(|e| e.to_string())?; + owner.try_lock().map_err(|e| e.to_string())?; + staged.owner = Some(owner); + let output = file_options(true) + .open(&staged.path) + .map_err(|e| e.to_string())?; + let staged = Arc::new(staged); + Ok((staged, output)) +} + +pub(super) fn copy_file(root: &Path, source: &Path) -> Result<(Arc, u64), String> { + let mut input_options = OpenOptions::new(); + input_options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + input_options.custom_flags(libc::O_NONBLOCK | libc::O_CLOEXEC); + } + let input = input_options + .open(source) + .map_err(|e| format!("Could not open the attachment: {e}"))?; + let metadata = input.metadata().map_err(|e| e.to_string())?; + if !metadata.is_file() { + return Err("Choose a file, not a folder".into()); + } + if metadata.len() > MAX_FILE_BYTES { + return Err("Files must be 30 MiB or smaller".into()); + } + let reserved_size = metadata.len(); + let reservation = budget().reserve(reserved_size)?; + let (staged, mut output) = begin(root, reservation)?; + let copied = std::io::copy(&mut input.take(reserved_size + 1), &mut output) + .map_err(|e| e.to_string())?; + if copied > reserved_size { + return Err("The file grew while it was being attached. Please attach it again.".into()); + } + output.sync_all().map_err(|e| e.to_string())?; + Ok((staged, copied)) +} + +pub(super) fn copy_bytes(root: &Path, bytes: &[u8]) -> Result, String> { + if bytes.len() as u64 > MAX_FILE_BYTES { + return Err("Files must be 30 MiB or smaller".into()); + } + let reservation = budget().reserve(bytes.len() as u64)?; + let (staged, mut output) = begin(root, reservation)?; + output.write_all(bytes).map_err(|e| e.to_string())?; + output.sync_all().map_err(|e| e.to_string())?; + Ok(staged) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn live_staging_budget_counts_files_and_bytes_and_releases_on_drop() { + let budget = Arc::new(StagingBudget::default()); + let reservation = budget.reserve(MAX_LIVE_BYTES).unwrap(); + assert!(budget.reserve(1).is_err()); + assert_eq!(budget.files.load(Ordering::Acquire), 1); + drop(reservation); + assert_eq!(budget.bytes.load(Ordering::Acquire), 0); + let reservations: Vec<_> = (0..MAX_LIVE_FILES) + .map(|_| budget.reserve(0).unwrap()) + .collect(); + assert!(budget.reserve(0).is_err()); + drop(reservations); + assert_eq!(budget.files.load(Ordering::Acquire), 0); + assert!(budget.reserve(MAX_LIVE_BYTES).is_ok()); + } + + #[test] + fn staging_path_failure_releases_its_reservation() { + let directory = tempfile::tempdir().unwrap(); + let not_a_directory = directory.path().join("occupied"); + fs::write(¬_a_directory, b"unrelated").unwrap(); + let budget = Arc::new(StagingBudget::default()); + let reserved = budget.reserve(123).unwrap(); + assert!(begin(¬_a_directory, reserved).is_err()); + assert_eq!(budget.files.load(Ordering::Acquire), 0); + assert_eq!(budget.bytes.load(Ordering::Acquire), 0); + assert_eq!(fs::read(not_a_directory).unwrap(), b"unrelated"); + } + + #[test] + fn zip_copy_survives_source_changes_and_removal() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("report.zip"); + fs::write(&source, b"PK\x03\x04original archive").unwrap(); + let (staged, size) = copy_file(&directory.path().join("staging"), &source).unwrap(); + fs::write(&source, b"changed").unwrap(); + fs::remove_file(source).unwrap(); + assert_eq!( + fs::read(&staged.path).unwrap(), + b"PK\x03\x04original archive" + ); + assert_eq!(size, 20); + } + + #[cfg(unix)] + #[test] + fn special_files_are_rejected_without_waiting_for_a_writer() { + use std::os::unix::ffi::OsStrExt; + let directory = tempfile::tempdir().unwrap(); + let fifo = directory.path().join("upload.pipe"); + let path = std::ffi::CString::new(fifo.as_os_str().as_bytes()).unwrap(); + assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0); + let result = copy_file(&directory.path().join("staging"), &fifo); + assert!(result.err().unwrap().contains("not a folder")); + assert!(!directory.path().join("staging").exists()); + } + + #[test] + fn final_draft_drop_cleans_only_its_staged_copy() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("staging"); + let staged = copy_bytes(&root, b"kept").unwrap(); + let path = staged.path.clone(); + let second_owner = staged.clone(); + drop(staged); + assert!(path.exists()); + drop(second_owner); + let deadline = std::time::Instant::now() + Duration::from_secs(3); + while path.exists() && std::time::Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(10)); + } + assert!(!path.exists()); + assert!(root.exists()); + } + + #[test] + fn restart_sweep_keeps_live_files_even_after_seven_days() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("staging"); + let staged = copy_bytes(&root, b"live draft").unwrap(); + sweep( + &root, + SystemTime::now() + STALE_AFTER + Duration::from_secs(1), + ); + assert_eq!(fs::read(&staged.path).unwrap(), b"live draft"); + // Simulate a crash: release the OS lease without running our drop. + let path = staged.path.clone(); + let mut staged = Arc::try_unwrap(staged).ok().unwrap(); + staged.owner.take(); + sweep( + &root, + SystemTime::now() + STALE_AFTER + Duration::from_secs(1), + ); + assert!(!path.exists()); + drop(staged); + } + + #[cfg(unix)] + #[test] + fn permissive_staging_root_is_rejected_without_changing_permissions() { + use std::os::unix::fs::PermissionsExt; + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("staging"); + fs::create_dir(&root).unwrap(); + fs::set_permissions(&root, fs::Permissions::from_mode(0o755)).unwrap(); + let error = copy_bytes(&root, b"private attachment").err().unwrap(); + assert!(error.contains("not private")); + assert_eq!( + fs::metadata(&root).unwrap().permissions().mode() & 0o777, + 0o755 + ); + assert_eq!(fs::read_dir(&root).unwrap().count(), 0); + } + + #[test] + fn concurrent_first_copies_share_a_privately_created_staging_root() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("staging"); + let barrier = Arc::new(std::sync::Barrier::new(8)); + let workers: Vec<_> = (0..8) + .map(|_| { + let root = root.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + copy_bytes(&root, b"private attachment").unwrap() + }) + }) + .collect(); + let staged: Vec<_> = workers + .into_iter() + .map(|worker| worker.join().unwrap()) + .collect(); + for file in &staged { + assert_eq!(fs::read(&file.path).unwrap(), b"private attachment"); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!(fs::metadata(&root).unwrap().permissions().mode() & 0o077, 0); + } + } +} diff --git a/apps/maple-agent/app/src/ui/chat/summaries.rs b/apps/maple-agent/app/src/ui/chat/summaries.rs index 70e28f6c7..79d4a1dae 100644 --- a/apps/maple-agent/app/src/ui/chat/summaries.rs +++ b/apps/maple-agent/app/src/ui/chat/summaries.rs @@ -9,15 +9,16 @@ use gpui::{Context, SharedString}; use maple_agent::agent::AgentTimelineItem; use super::ChatScreen; -use super::transcript::{attachment_refs, has_tool_input, tool_output_markdown}; +use super::transcript::{has_tool_input, image_attachment_refs, tool_output_markdown}; use crate::backend::{AgentBackend, ReadAttachmentRequest, StoreToolSummaryRequest}; -use crate::ui::chat::images::image_format_from_bytes; +use crate::ui::chat::attachments::image_format_from_bytes; /// A thinking block shorter than this reads faster than its summary; /// skip the model call. const THINKING_SUMMARY_MIN_CHARS: usize = 200; const ATTACHMENT_CHUNK_BYTES: usize = 1024 * 1024; -const MAX_ATTACHMENT_BYTES: usize = 10 * 1024 * 1024; +// Generated/tool images can be 20 MiB; user image admission remains 10 MiB. +const MAX_ATTACHMENT_BYTES: usize = 20 * 1024 * 1024; const MAX_CACHED_ATTACHMENT_BYTES: usize = 32 * 1024 * 1024; const MAX_CACHED_ATTACHMENTS: usize = 64; const MAX_ATTACHMENT_READS: usize = 4; @@ -258,7 +259,7 @@ impl ChatScreen { self.attachment_viewport = start..end; let visible: HashSet = self.timeline[start..end] .iter() - .flat_map(attachment_refs) + .flat_map(image_attachment_refs) .map(|(id, _)| id.to_string()) .collect(); // A fast scroll must not leave a backlog downloading distant pages. @@ -298,7 +299,7 @@ impl ChatScreen { return; } } - let wanted: Vec = attachment_refs(item) + let wanted: Vec = image_attachment_refs(item) .map(|(id, _)| id.to_string()) .filter(|id| { !self.attachment_images.contains_key(id) && !self.attachment_requests.contains(id) @@ -344,7 +345,14 @@ impl ChatScreen { if this.selection_generation == generation && this.is_selected(&target) { match result { Ok(image) => this.insert_attachment_image(&id, image, cx), - Err(message) => log::debug!("attachment {id} not loaded: {message}"), + Err(_) => { + this.attachment_failures.insert(id.clone()); + let rows: Vec<_> = this.timeline.iter().enumerate().filter(|(_, item)| + image_attachment_refs(item).any(|(candidate, _)| candidate == id)) + .map(|(index, _)| index).collect(); + for index in rows { this.remeasure_item(index); } + cx.notify(); + }, } } // Stale reads still occupied a slot; their completion can @@ -369,7 +377,7 @@ impl ChatScreen { .min(self.timeline.len()); let visible: HashSet = self.timeline[start..end] .iter() - .flat_map(attachment_refs) + .flat_map(image_attachment_refs) .map(|(id, _)| id.to_string()) .collect(); self.attachment_images.remove(id); @@ -408,7 +416,7 @@ impl ChatScreen { // requested and measured when that row becomes visible again. let rows: Vec = (start..end) .filter(|&index| { - attachment_refs(&self.timeline[index]) + image_attachment_refs(&self.timeline[index]) .any(|(candidate, _)| changed.contains(candidate)) }) .collect(); diff --git a/apps/maple-agent/app/src/ui/chat/tests.rs b/apps/maple-agent/app/src/ui/chat/tests.rs index c44623224..1f00bfb9c 100644 --- a/apps/maple-agent/app/src/ui/chat/tests.rs +++ b/apps/maple-agent/app/src/ui/chat/tests.rs @@ -3,10 +3,10 @@ mod state_tests { use std::sync::Arc; + use crate::ui::chat::attachments::MAX_DRAFT_ATTACHMENTS; use crate::ui::chat::cache::{INLINE_PARSE_LIMIT, MAX_DIFF_LINES, ORDINAL_SPACING}; use crate::ui::chat::composer::SideThreadTurn; use crate::ui::chat::computers::ProjectTarget; - use crate::ui::chat::images::{MAX_DRAFT_IMAGES, encode_data_url}; use crate::ui::chat::sidebar::{RenameTarget, SidebarEvent, TaskMove}; use crate::ui::chat::transcript::{diff_lines_for, tool_label_title}; use crate::ui::chat::*; @@ -1020,7 +1020,7 @@ mod state_tests { .expect("backend"), ); // This gpui's test scheduler flags activity on other threads unless - // parking is allowed; the backend runtime and image encoder run on tokio. + // parking is allowed; the backend runtime and attachment worker run on tokio. cx.executor().allow_parking(); let screen = cx.new(|cx| { let mut screen = ChatScreen::new_inner(backend, "user".to_string(), cx); @@ -1377,7 +1377,7 @@ mod state_tests { .unwrap(); crate::settings::install_test_store(store); // This gpui's test scheduler flags activity on other threads unless - // parking is allowed; the backend runtime and image encoder run on tokio. + // parking is allowed; the backend runtime and attachment worker run on tokio. cx.executor().allow_parking(); let screen = cx.new(|cx| { ChatScreen::new_inner( @@ -3334,12 +3334,7 @@ mod state_tests { assert_eq!(shown(cx), "stop", "whitespace is not a draft"); chat.update(cx, |this, cx| { - this.draft_images = vec![DraftImage { - id: 1, - name: "shot.png".into(), - data_url: None, - thumbnail: None, - }]; + this.draft_attachments = vec![DraftAttachment::for_test(1, "shot.png", false)]; cx.notify(); }); assert_eq!( @@ -3349,7 +3344,7 @@ mod state_tests { ); chat.update(cx, |this, cx| { - this.draft_images.clear(); + this.draft_attachments.clear(); this.booting = true; cx.notify(); }); @@ -4974,13 +4969,53 @@ mod state_tests { screen.update(cx, |this, cx| { let mut sent = user_item("u1", "see this"); sent.input = Some(serde_json::json!({ - "imageAttachments": [{"id": "att-1", "name": "a.png"}] + "attachments": [{"id": "att-1", "name": "a.png", "mimeType": "image/png", "size": 24, "executionPath": "/uploads/a.png"}] })); this.apply_incoming_item("s1", sent, cx); assert!(this.attachment_requests.contains("att-1")); }); } + #[gpui::test] + fn generic_and_retired_attachments_never_request_image_bytes(cx: &mut TestAppContext) { + let screen = screen(cx); + screen.update(cx, |this, cx| { + let mut sent = user_item("files", ""); + sent.input = Some(serde_json::json!({"attachments": [ + {"id":"archive", "name":"report.zip", "mimeType":"application/zip", "size":123, "executionPath":"/uploads/report.zip"}, + {"id":"retired", "name":"old.png", "mimeType":"image/unknown", "size":0, "executionPath":""} + ]})); + this.apply_incoming_item("s1", sent.clone(), cx); + for _ in 0..3 { this.load_attachment_images_for(&sent, cx); } + assert!(this.attachment_requests.is_empty()); + assert!(this.attachment_queue.is_empty()); + let derived = this.derived.get(&sent, 0); + assert_eq!(derived.attachments.len(), 2); + assert_eq!(derived.attachments[0].detail.as_ref(), "ZIP · 123 B"); + assert!(derived.attachments[1].unavailable); + }); + } + + #[gpui::test] + fn dropping_a_preparing_attachment_discards_its_late_copy(cx: &mut TestAppContext) { + let screen = screen(cx); + screen.update(cx, |this, cx| { + let draft = DraftAttachment::for_test(999, "report.zip", true); + let path = draft.staged.as_ref().unwrap().path.clone(); + this.finish_attachment_preparation( + 999, + Ok((draft.staged.unwrap(), "application/zip".into(), 15, None)), + cx, + ); + assert!(this.draft_attachments.is_empty()); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(3); + while path.exists() && std::time::Instant::now() < deadline { + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert!(!path.exists()); + }); + } + #[test] fn test_diff_lines_stop_at_the_cap() { let mut tool = item("edit", "tool", None); @@ -6985,7 +7020,7 @@ mod state_tests { screen.update(cx, |this, cx| { let mut promoted = user_item("u1", "queued text"); promoted.input = Some(serde_json::json!({ - "imageAttachments": [{ "id": "att-9", "name": "shot.png" }] + "attachments": [{ "id": "att-9", "name": "shot.png", "mimeType": "image/png", "size": 24, "executionPath": "/uploads/shot.png" }] })); this.handle_run_event( "s1", @@ -7487,21 +7522,15 @@ mod state_tests { screen.update(cx, |this, cx| { let image = gpui::Image::from_bytes(gpui::ImageFormat::Png, png_bytes()); this.paste_image(image, cx); - assert_eq!(this.draft_images.len(), 1); - assert_eq!(this.draft_images[0].name, "pasted-1.png"); - // The encode runs off the UI thread; the draft is staged at + assert_eq!(this.draft_attachments.len(), 1); + assert_eq!(this.draft_attachments[0].name, "pasted-1.png"); + // The stable copy is made off the UI thread; the draft is staged at // once and cannot be sent until it lands. - assert!(!this.draft_images[0].ready()); + assert!(!this.draft_attachments[0].ready()); assert!(this.notice.is_none()); }); } - #[test] - fn test_data_url_encodes_the_signature() { - let url = encode_data_url("image/png", &png_bytes()); - assert!(url.starts_with("data:image/png;base64,iVBORw0KGgo")); - } - #[gpui::test] fn test_paste_image_rejects_unsupported_format(cx: &mut TestAppContext) { cx.executor().allow_parking(); @@ -7509,7 +7538,7 @@ mod state_tests { screen.update(cx, |this, cx| { let image = gpui::Image::from_bytes(gpui::ImageFormat::Gif, b"GIF89a".to_vec()); this.paste_image(image, cx); - assert!(this.draft_images.is_empty()); + assert!(this.draft_attachments.is_empty()); assert!( this.notice .as_deref() @@ -7523,11 +7552,11 @@ mod state_tests { cx.executor().allow_parking(); let screen = screen(cx); screen.update(cx, |this, cx| { - for _ in 0..=MAX_DRAFT_IMAGES { + for _ in 0..=MAX_DRAFT_ATTACHMENTS { let image = gpui::Image::from_bytes(gpui::ImageFormat::Png, png_bytes()); this.paste_image(image, cx); } - assert_eq!(this.draft_images.len(), MAX_DRAFT_IMAGES); + assert_eq!(this.draft_attachments.len(), MAX_DRAFT_ATTACHMENTS); assert!( this.notice .as_deref() @@ -7557,7 +7586,7 @@ mod state_tests { } // This gpui's test scheduler flags activity on other threads unless - // parking is allowed; the backend runtime and image encoder run on tokio. + // parking is allowed; the backend runtime and attachment worker run on tokio. cx.executor().allow_parking(); let chat = cx.new(|cx| { let _guard = SETTINGS_LOCK.lock(); @@ -8194,7 +8223,7 @@ mod state_tests { } // This gpui's test scheduler flags activity on other threads unless - // parking is allowed; the backend runtime and image encoder run on tokio. + // parking is allowed; the backend runtime and attachment worker run on tokio. cx.executor().allow_parking(); let chat = cx.new(|cx| { let _guard = SETTINGS_LOCK.lock(); @@ -8550,7 +8579,7 @@ mod state_tests { cx.simulate_keystrokes("G ctrl-w k"); assert_eq!(row(cx).2, Some(composer_focus.clone())); cx.simulate_keystrokes("ctrl-w j"); - assert_eq!(row(cx).1, Some(ComposerControl::Images)); + assert_eq!(row(cx).1, Some(ComposerControl::Attachments)); cx.simulate_keystrokes("ctrl-w h"); assert_eq!(row(cx).0, ChatRegion::Sidebar); cx.simulate_keystrokes("ctrl-w l"); @@ -8558,7 +8587,7 @@ mod state_tests { row(cx), ( ChatRegion::Controls, - Some(ComposerControl::Images), + Some(ComposerControl::Attachments), Some(application_focus) ) ); @@ -8988,35 +9017,106 @@ mod state_tests { composer.update(cx, |input, cx| input.set_text(draft, cx)); let id = uuid::Uuid::new_v4().to_string(); chat.finish_send(id.clone(), UnresolvedSend { + backend: chat.backend.clone(), outcome: None, session: "s1".into(), text: "already admitted original message".into(), from_composer: true, - images: vec![DraftImage { - id: 50, - name: "image.png".into(), - data_url: Some(Arc::from("data:image/png;base64,cGljdHVyZQ==")), - thumbnail: None, - }], + attachments: vec![DraftAttachment::for_test(50, "image.png", true)], }, Err(format!("Send outcome unknown ({id}). Reload this task before sending again; it may already be running.")), cx); assert_eq!(composer.read(cx).text_ref(), draft); - assert!(chat.draft_images.is_empty()); - assert_eq!(chat.unresolved_sends[&id].images.len(), 1); + assert!(chat.draft_attachments.is_empty()); + assert_eq!(chat.unresolved_sends[&id].attachments.len(), 1); + assert!(chat.unresolved_sends[&id].attachments[0].staged.as_ref().unwrap().path.exists()); assert!(chat.unresolved_sends[&id].awaiting_receipt()); chat.apply_send_receipt(&id, Some(Outcome::Pending), cx); assert_eq!(composer.read(cx).text_ref(), draft); - assert!(chat.draft_images.is_empty()); + assert!(chat.draft_attachments.is_empty()); assert!(chat.unresolved_sends[&id].awaiting_receipt()); chat.apply_send_receipt(&id, Some(Outcome::Finished(Ok(serde_json::json!("accepted-run")))), cx); assert_eq!(composer.read(cx).text_ref(), draft); - assert!(chat.draft_images.is_empty()); + assert!(chat.draft_attachments.is_empty()); assert!(!chat.unresolved_sends.contains_key(&id)); } }); } #[gpui::test] - fn uncertain_send_keeps_images_and_does_not_overwrite_another_draft(cx: &mut TestAppContext) { + fn deleting_a_task_cancels_its_pending_upload_and_discards_send_recovery( + cx: &mut TestAppContext, + ) { + let chat = screen(cx); + chat.update(cx, |this, cx| { + let drafts = vec![DraftAttachment::for_test(1, "report.zip", true)]; + let (fence, progress) = this.begin_uploads("pending", "s1", &drafts, cx); + drop(progress); + this.sending_previews.insert( + "pending".into(), + sending::SendPreview::new(Some("s1".into()), "read this", &drafts), + ); + this.unresolved_sends.insert( + "uncertain".into(), + UnresolvedSend { + backend: this.backend.clone(), + outcome: None, + session: "s1".into(), + text: "older".into(), + attachments: drafts, + from_composer: true, + }, + ); + this.remove_session("s1", cx); + assert!(fence.cancelled()); + assert!(this.attachment_uploads.is_empty()); + assert!(this.unresolved_sends.is_empty()); + assert!(this.discarded_sends.contains("pending")); + }); + } + + #[gpui::test] + fn view_teardown_cancels_upload_permission_before_a_detached_import_can_send( + cx: &mut TestAppContext, + ) { + let chat = screen(cx); + chat.update(cx, |this, cx| { + let drafts = vec![DraftAttachment::for_test(1, "report.zip", true)]; + let (fence, progress) = this.begin_uploads("pending", "s1", &drafts, cx); + drop(progress); + this.attachment_uploads.clear(); + assert!(fence.cancelled()); + assert!(!fence.hand_off()); + }); + } + + #[gpui::test] + fn pending_receipts_keep_the_computer_that_received_the_send(cx: &mut TestAppContext) { + let chat = screen(cx); + chat.update(cx, |this, _| { + let origin = Arc::new(AgentBackend::for_remote_test("execution-a").unwrap()); + this.unresolved_sends.insert( + "original".into(), + UnresolvedSend { + backend: origin.clone(), + outcome: None, + session: "task-a".into(), + text: "keep".into(), + attachments: vec![], + from_composer: true, + }, + ); + this.backend = Arc::new(AgentBackend::for_remote_test("execution-b").unwrap()); + let queries = this.pending_receipt_queries(); + assert_eq!(queries.len(), 1); + assert_eq!(queries[0].id, "original"); + assert!(Arc::ptr_eq(&queries[0].backend, &origin)); + assert_eq!(queries[0].backend.remote_host(), Some("execution-a")); + }); + } + + #[gpui::test] + fn uncertain_send_keeps_attachments_and_does_not_overwrite_another_draft( + cx: &mut TestAppContext, + ) { use crate::host::receipts::Outcome; let chat = screen(cx); chat.update(cx, |chat, cx| { @@ -9026,20 +9126,16 @@ mod state_tests { chat.unresolved_sends.insert( "send-original".into(), UnresolvedSend { + backend: chat.backend.clone(), outcome: None, session: "s1".into(), text: "original message".into(), from_composer: true, - images: vec![DraftImage { - id: 50, - name: "image.png".into(), - data_url: Some(Arc::from("data:image/png;base64,cGljdHVyZQ==")), - thumbnail: None, - }], + attachments: vec![DraftAttachment::for_test(50, "image.png", true)], }, ); chat.apply_send_receipt("send-original", None, cx); - assert_eq!(chat.unresolved_sends["send-original"].images.len(), 1); + assert_eq!(chat.unresolved_sends["send-original"].attachments.len(), 1); chat.apply_send_receipt("send-original", Some(Outcome::Missing), cx); assert_eq!(composer.read(cx).text_ref(), "new draft"); assert!(chat.unresolved_sends.contains_key("send-original")); @@ -9050,7 +9146,7 @@ mod state_tests { chat.selected_session = Some("s1".into()); chat.apply_send_receipt("send-original", Some(Outcome::Missing), cx); assert_eq!(composer.read(cx).text_ref(), "original message"); - assert_eq!(chat.draft_images[0].name, "image.png"); + assert_eq!(chat.draft_attachments[0].name, "image.png"); assert!(chat.unresolved_sends.is_empty()); }); } @@ -9072,10 +9168,11 @@ mod state_tests { chat.unresolved_sends.insert( "original".into(), UnresolvedSend { + backend: chat.backend.clone(), outcome: None, session: "s1".into(), text: "original draft".into(), - images: vec![], + attachments: vec![], from_composer: true, }, ); @@ -9577,24 +9674,15 @@ mod state_tests { .clone() .unwrap() .update(cx, |input, cx| input.set_text("newer draft", cx)); - this.draft_images = vec![DraftImage { - id: 2, - name: "new.png".into(), - data_url: None, - thumbnail: None, - }]; + this.draft_attachments = vec![DraftAttachment::for_test(2, "new.png", false)]; this.finish_send( "send".into(), UnresolvedSend { + backend: this.backend.clone(), outcome: None, session: "s1".into(), text: "original".into(), - images: vec![DraftImage { - id: 1, - name: "original.png".into(), - data_url: None, - thumbnail: None, - }], + attachments: vec![DraftAttachment::for_test(1, "original.png", false)], from_composer: true, }, Err("failed".into()), @@ -9607,16 +9695,19 @@ mod state_tests { "another task's waiting state stays intact" ); assert_eq!(this.composer_text(cx).as_deref(), Some("newer draft")); - assert_eq!(this.draft_images[0].name, "new.png"); + assert_eq!(this.draft_attachments[0].name, "new.png"); assert_eq!(this.unresolved_sends["send"].text, "original"); - assert_eq!(this.unresolved_sends["send"].images[0].name, "original.png"); + assert_eq!( + this.unresolved_sends["send"].attachments[0].name, + "original.png" + ); // Returning to an empty original task recovers the retained draft. this.selected_session = Some("s1".into()); this.composer .clone() .unwrap() .update(cx, |input, cx| input.clear(cx)); - this.draft_images.clear(); + this.draft_attachments.clear(); this.apply_send_receipt( "send", Some(crate::host::receipts::Outcome::Finished(Err( @@ -9625,7 +9716,7 @@ mod state_tests { cx, ); assert_eq!(this.composer_text(cx).as_deref(), Some("original")); - assert_eq!(this.draft_images[0].name, "original.png"); + assert_eq!(this.draft_attachments[0].name, "original.png"); }); } diff --git a/apps/maple-agent/app/src/ui/chat/transcript.rs b/apps/maple-agent/app/src/ui/chat/transcript.rs index e98ad95eb..29ab64f6a 100644 --- a/apps/maple-agent/app/src/ui/chat/transcript.rs +++ b/apps/maple-agent/app/src/ui/chat/transcript.rs @@ -68,6 +68,7 @@ impl ChatScreen { markdown_cache: &chat.markdown_cache, derived: &chat.derived, attachment_images: &chat.attachment_images, + attachment_failures: &chat.attachment_failures, chat: &entity, tool_summaries: &chat.tool_summaries, render: &render_ctx, @@ -389,21 +390,59 @@ pub(super) fn render_timeline_item( } } -/// Attachment `(id, name)` pairs stored on a user message. -pub(super) fn attachment_refs(item: &AgentTimelineItem) -> impl Iterator { +pub(super) struct AttachmentRef<'a> { + pub id: &'a str, + pub name: &'a str, + pub mime: &'a str, + pub size: u64, + pub unavailable: bool, +} + +pub(super) fn attachment_metadata( + item: &AgentTimelineItem, +) -> impl Iterator> { item.input .as_ref() - .and_then(|input| input.get("imageAttachments")) + .and_then(|input| input.get("attachments")) .and_then(|items| items.as_array()) .into_iter() .flatten() .filter_map(|entry| { - let id = entry.get("id").and_then(|id| id.as_str())?; - let name = entry.get("name").and_then(|name| name.as_str())?; - Some((id, name)) + let id = entry.get("id")?.as_str()?; + let name = entry.get("name")?.as_str()?; + let mime = entry + .get("mimeType") + .and_then(|value| value.as_str()) + .unwrap_or("application/octet-stream"); + let size = entry + .get("size") + .and_then(|value| value.as_u64()) + .unwrap_or(0); + Some(AttachmentRef { + id, + name, + mime, + size, + unavailable: mime == "image/unknown", + }) }) } +/// All file references count as message content, including retired images. +pub(super) fn attachment_refs(item: &AgentTimelineItem) -> impl Iterator { + attachment_metadata(item).map(|attachment| (attachment.id, attachment.name)) +} + +/// Only renderable image references need bytes. Generic files and known retired +/// images never occupy the thumbnail request queue. +pub(super) fn image_attachment_refs( + item: &AgentTimelineItem, +) -> impl Iterator { + attachment_metadata(item) + .filter(|attachment| attachment.mime.starts_with("image/") && !attachment.unavailable) + .map(|attachment| (attachment.id, attachment.name)) +} + /// Hover-revealed button that copies one message's text. fn copy_message_button( item_id: &str, @@ -444,9 +483,9 @@ fn render_message(item: &AgentTimelineItem, revision: u64, transcript: &Transcri let chat = transcript.chat; let is_user = item.role.as_deref() == Some("user"); let text = item.text.as_deref().unwrap_or(""); - let mut attachments = attachment_refs(item).peekable(); - let has_images = attachments.peek().is_some(); - if text.trim().is_empty() && !(is_user && has_images) { + let derived = transcript.derived.get(item, revision); + let has_attachments = !derived.attachments.is_empty(); + if text.trim().is_empty() && !(is_user && has_attachments) { return div(); } let group = SharedString::from(format!("message-{}", item.id)); @@ -480,9 +519,12 @@ fn render_message(item: &AgentTimelineItem, revision: u64, transcript: &Transcri .border_1() .border_color(gpui::rgb(theme::user_bubble_border())) .text_color(gpui::rgb(theme::text_primary())) - .when(has_images, |bubble| { + .when(has_attachments, |bubble| { bubble.child(div().flex().flex_wrap().gap_2().mb_1().children( - attachments.map(|(id, name)| { + derived.attachments.iter().map(|attachment| { + let id = attachment.id.as_str(); + let unavailable = attachment.unavailable + || transcript.attachment_failures.contains(id); match attachment_images.get(id) { // The picture itself, scaled to fit; click // opens it full size. gpui keeps the aspect @@ -534,7 +576,22 @@ fn render_message(item: &AgentTimelineItem, revision: u64, transcript: &Transcri .text_xs() .text_color(gpui::rgb(theme::text_secondary())) .child(icon("paperclip", px(12.), theme::text_secondary())) - .child(name.to_string()), + .child( + div() + .flex() + .flex_col() + .gap_0p5() + .child(attachment.name.clone()) + .child( + div() + .text_color(gpui::rgb(theme::text_muted())) + .child(if unavailable { + SharedString::from("Image unavailable") + } else { + attachment.detail.clone() + }), + ), + ), } }), )) diff --git a/apps/maple-agent/app/src/ui/chat/upload_tests.rs b/apps/maple-agent/app/src/ui/chat/upload_tests.rs new file mode 100644 index 000000000..64599cb0c --- /dev/null +++ b/apps/maple-agent/app/src/ui/chat/upload_tests.rs @@ -0,0 +1,473 @@ +//! Composer upload/admission interleavings through the real backend facade and +//! framed IPC client. The host holds individual replies; no model or account is used. + +use super::*; +use crate::host::{RemoteClient, ipc, protocol, receipts::Outcome}; +use gpui::TestAppContext; +use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; +use tokio::sync::mpsc; + +#[derive(Clone)] +struct ObservedRequest { + id: u64, + method: String, + send_id: Option, + params: Value, +} + +struct UploadHost { + backend: Arc, + requests: Arc>>, + replies: mpsc::UnboundedSender<(u64, Result)>, + runtime: Option, + _root: tempfile::TempDir, +} + +impl UploadHost { + fn new() -> Self { + let root = tempfile::tempdir().unwrap(); + let endpoint = ipc::Endpoint::for_test(root.path().to_path_buf()); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let recorded = requests.clone(); + let (replies, mut responses) = mpsc::unbounded_channel::<(u64, Result)>(); + let backend = runtime.block_on(async { + let (client, server) = tokio::io::duplex(64 * 1024); + tokio::spawn(async move { + let (mut input, mut output) = tokio::io::split(server); + // Keep reads outside select! so a released reply cannot cancel + // a partially consumed request frame. + let (incoming, mut received) = mpsc::channel(32); + let reader = tokio::spawn(async move { + while let Ok(request) = protocol::read::(&mut input).await { + if incoming.send(request).await.is_err() { + break; + } + } + }); + loop { + let response = tokio::select! { + reply = responses.recv() => { + let Some(reply) = reply else { break }; + reply + } + request = received.recv() => { + let Some(request) = request else { break }; + recorded.lock().unwrap().push(ObservedRequest { + id: request.id, + method: request.method.clone(), + send_id: request.send_id, + params: request.params, + }); + if matches!(request.method.as_str(), "import_attachment" | "attachment_upload_status" | "send_message" | "send_outcome") { + continue; + } + let result = match request.method.as_str() { + "cancel_run" | "cancel_attachment_upload" | "set_ui_settings" | "set_remote_sidebar" | "flush_ui_settings" | "detach" => Ok(Value::Null), + "get_ui_settings" => Ok(serde_json::to_value(crate::settings::AppSettings::default()).unwrap()), + "get_remote_sidebars" => Ok(json!({})), + "list_slash_commands" => Ok(json!([])), + _ => Err("Optional upload fixture operation unavailable".into()), + }; + (request.id, result) + } + }; + if protocol::write(&mut output, &protocol::Message::Response { + id: response.0, + result: response.1, + }).await.is_err() { + break; + } + } + reader.abort(); + }); + Arc::new(AgentBackend::from_host_client_for_test(RemoteClient::attach_for_test( + endpoint, + Box::new(client), + protocol::Welcome { + version: protocol::VERSION, + pid: std::process::id(), + build: protocol::BUILD.into(), + user_id: Some("user".into()), + store: maple_agent::store::StoreState::Open, + }, + ))) + }); + Self { + backend, + requests, + replies, + runtime: Some(runtime), + _root: root, + } + } + + fn requests(&self, method: &str) -> Vec { + self.requests + .lock() + .unwrap() + .iter() + .filter(|request| request.method == method) + .cloned() + .collect() + } + + fn reply(&self, request: &ObservedRequest, result: Result) { + self.replies.send((request.id, result)).unwrap(); + } + + fn finish_upload(&self) { + self.reply( + &self.requests("import_attachment")[0], + Ok(serde_json::to_value(maple_agent::store::StoredAttachment { + id: "0123456789abcdef0123456789abcdef".into(), + name: "file.txt".into(), + mime_type: "application/octet-stream".into(), + source: "maple-attachment://0123456789abcdef0123456789abcdef".into(), + size: 15, + execution_path: "/synthetic-execution/file.txt".into(), + }) + .unwrap()), + ); + } +} + +impl Drop for UploadHost { + fn drop(&mut self) { + if let Some(runtime) = self.runtime.take() { + runtime.shutdown_background(); + } + } +} + +fn chat_window<'a>( + host: &UploadHost, + cx: &'a mut TestAppContext, +) -> (Entity, &'a mut gpui::VisualTestContext) { + struct WindowHost(Entity); + impl Render for WindowHost { + fn render(&mut self, _: &mut Window, _: &mut Context) -> impl IntoElement { + div() + .w(px(1200.)) + .h(px(800.)) + .flex() + .flex_col() + .child(self.0.clone()) + } + } + cx.executor().allow_parking(); + let chat = cx.new(|cx| { + let _guard = crate::settings::SETTINGS_IO_LOCK.lock(); + crate::desktop::register_key_bindings(cx); + let mut chat = ChatScreen::new_without_start(host.backend.clone(), "user".into(), cx); + chat.selected_session = Some("s1".into()); + chat.booting = false; + chat.trust_prompts = false; + chat.active_runs.insert("s1".into(), "run-1".into()); + chat.sidebar + .update(cx, |sidebar, _| sidebar.reset_persisted_for_test()); + chat + }); + let (_, cx) = cx.add_window_view(|_, _| WindowHost(chat.clone())); + cx.simulate_resize(gpui::size(px(1200.), px(800.))); + cx.update(|window, _| window.activate_window()); + cx.run_until_parked(); + (chat, cx) +} + +fn wait_for( + cx: &mut gpui::VisualTestContext, + label: &str, + mut ready: impl FnMut(&mut gpui::VisualTestContext) -> bool, +) { + for _ in 0..500 { + cx.run_until_parked(); + if ready(cx) { + return; + } + std::thread::sleep(std::time::Duration::from_millis(10)); + } + panic!("upload fixture timed out waiting for {label}"); +} + +fn start_upload(chat: &Entity, cx: &mut gpui::VisualTestContext, host: &UploadHost) { + chat.update(cx, |chat, cx| { + chat.draft_attachments = vec![DraftAttachment::for_test(1, "file.txt", true)]; + chat.composer + .clone() + .unwrap() + .update(cx, |input, cx| input.set_text("with file", cx)); + chat.send_inner(cx); + }); + wait_for(cx, "held import", |_| { + host.requests("import_attachment").len() == 1 + }); + assert!( + host.requests("send_message").is_empty(), + "bytes precede the identified prompt" + ); +} + +fn type_followup( + chat: &Entity, + cx: &mut gpui::VisualTestContext, + text: &str, + steer: bool, +) { + chat.update(cx, |chat, cx| { + chat.composer + .clone() + .unwrap() + .update(cx, |input, cx| input.set_text(text, cx)); + if steer { + chat.steer_text(text.into(), cx); + } else { + chat.send_inner(cx); + } + }); +} + +fn wait_prepared(chat: &Entity, cx: &mut gpui::VisualTestContext) { + wait_for(cx, "uploaded files waiting for admission", |cx| { + chat.update(cx, |chat, _| { + chat.sending_previews + .values() + .any(|preview| preview.prepared().is_some()) + }) + }); +} + +#[gpui::test] +fn slow_upload_keeps_stop_queue_and_steer_operational(cx: &mut TestAppContext) { + let host = UploadHost::new(); + let (chat, cx) = chat_window(&host, cx); + start_upload(&chat, cx, &host); + cx.run_until_parked(); + cx.update(|window, _| window.refresh()); + cx.run_until_parked(); + let stop = cx + .debug_bounds("stop-run") + .expect("Stop remains visible during upload"); + cx.simulate_click(stop.center(), gpui::Modifiers::default()); + wait_for(cx, "Stop RPC while import remains held", |_| { + host.requests("cancel_run").len() == 1 + }); + assert_eq!( + host.requests("cancel_run")[0].params, + json!(["user", "run-1"]) + ); + + type_followup(&chat, cx, "queue this", false); + wait_for(cx, "queued text admission", |_| { + host.requests("send_message").len() == 1 + }); + let queued = host.requests("send_message")[0].clone(); + assert_eq!(queued.params[1]["text"], "queue this"); + assert_eq!(queued.params[1]["steer"], false); + assert_eq!(queued.params[1]["attachments"], json!([])); + host.reply(&queued, Ok(json!("run-1"))); + wait_for(cx, "queue acknowledgement", |cx| { + chat.update(cx, |chat, _| !chat.send_pending()) + }); + + type_followup(&chat, cx, "steer this", true); + wait_for(cx, "steering admission", |_| { + host.requests("send_message").len() == 2 + }); + let steering = host.requests("send_message")[1].clone(); + assert_eq!(steering.params[1]["text"], "steer this"); + assert_eq!(steering.params[1]["steer"], true); + host.finish_upload(); + wait_prepared(&chat, cx); + assert_eq!( + host.requests("send_message").len(), + 2, + "upload waits for the newer text admission" + ); + host.reply(&steering, Ok(json!("run-1"))); + wait_for(cx, "file prompt after steering settles", |_| { + host.requests("send_message").len() == 3 + }); + let file_send = host.requests("send_message")[2].clone(); + assert_eq!(file_send.params[1]["text"], "with file"); + assert_eq!( + file_send.params[1]["attachments"], + json!(["0123456789abcdef0123456789abcdef"]) + ); + assert_ne!(file_send.send_id, queued.send_id); + assert_ne!(file_send.send_id, steering.send_id); + chat.update(cx, |chat, cx| { + let (id, upload) = chat + .attachment_uploads + .iter() + .next() + .expect("send acknowledgement is held"); + assert!(upload.fence.handed_off()); + let id = id.clone(); + chat.cancel_attachment_send(&id, cx); + }); + assert!( + host.requests("cancel_attachment_upload").is_empty(), + "cancellation cannot win after prompt handoff" + ); + host.reply(&file_send, Ok(json!("run-1"))); + wait_for(cx, "file acknowledgement", |cx| { + chat.update(cx, |chat, _| chat.attachment_uploads.is_empty()) + }); +} + +#[gpui::test] +fn uploaded_prompt_waits_for_unknown_text_receipt_without_replaying_text(cx: &mut TestAppContext) { + let host = UploadHost::new(); + let (chat, cx) = chat_window(&host, cx); + start_upload(&chat, cx, &host); + type_followup(&chat, cx, "possibly received", false); + wait_for(cx, "text admission", |_| { + host.requests("send_message").len() == 1 + }); + let text = host.requests("send_message")[0].clone(); + host.finish_upload(); + wait_prepared(&chat, cx); + host.reply( + &text, + Err("Send outcome unknown. Synthetic lost acknowledgement.".into()), + ); + wait_for(cx, "receipt query", |_| { + host.requests("send_outcome").len() == 1 + }); + let receipt = host.requests("send_outcome")[0].clone(); + assert_eq!(receipt.params, json!(text.send_id.as_ref().unwrap())); + assert_eq!(host.requests("send_message").len(), 1); + host.reply( + &receipt, + Ok(serde_json::to_value(Outcome::Pending).unwrap()), + ); + wait_for(cx, "pending receipt retained", |cx| { + chat.update(cx, |chat, _| { + !chat.checking_sends + && chat + .unresolved_sends + .values() + .any(UnresolvedSend::awaiting_receipt) + }) + }); + assert_eq!( + host.requests("send_message").len(), + 1, + "Pending authorizes neither replay nor the prepared prompt" + ); + chat.update(cx, |chat, cx| chat.check_unresolved_sends(cx)); + wait_for(cx, "same receipt queried again", |_| { + host.requests("send_outcome").len() == 2 + }); + let resolved = host.requests("send_outcome")[1].clone(); + assert_eq!(resolved.params, receipt.params); + host.reply( + &resolved, + Ok(serde_json::to_value(Outcome::Finished(Ok(json!("run-1")))).unwrap()), + ); + wait_for(cx, "prepared prompt admitted after receipt", |_| { + host.requests("send_message").len() == 2 + }); + let file_send = host.requests("send_message")[1].clone(); + assert_eq!(file_send.params[1]["text"], "with file"); + assert_ne!(file_send.send_id, text.send_id); + assert_eq!( + host.requests("send_message") + .iter() + .filter(|request| request.send_id == text.send_id) + .count(), + 1 + ); + host.reply(&file_send, Ok(json!("run-1"))); + wait_for(cx, "uploaded send settled", |cx| { + chat.update(cx, |chat, _| chat.attachment_uploads.is_empty()) + }); +} + +#[gpui::test] +fn cancelling_prepared_upload_before_handoff_never_dispatches_its_prompt(cx: &mut TestAppContext) { + let host = UploadHost::new(); + let (chat, cx) = chat_window(&host, cx); + start_upload(&chat, cx, &host); + type_followup(&chat, cx, "keep this text", false); + wait_for(cx, "text admission", |_| { + host.requests("send_message").len() == 1 + }); + let text = host.requests("send_message")[0].clone(); + host.finish_upload(); + wait_prepared(&chat, cx); + chat.update(cx, |chat, cx| { + let (id, upload) = chat + .selected_upload() + .expect("prepared upload remains cancellable"); + assert!(!upload.fence.handed_off()); + let id = id.to_owned(); + chat.cancel_attachment_send(&id, cx); + }); + wait_for(cx, "upload cancellation RPC", |_| { + host.requests("cancel_attachment_upload").len() == 1 + }); + wait_for(cx, "cancelled draft restored", |cx| { + chat.update(cx, |chat, cx| { + chat.attachment_uploads.is_empty() + && chat.draft_attachments.len() == 1 + && chat.composer_text(cx).as_deref() == Some("with file") + }) + }); + host.reply(&text, Ok(json!("run-1"))); + wait_for(cx, "newer text acknowledgement", |cx| { + chat.update(cx, |chat, _| !chat.send_pending()) + }); + assert_eq!( + host.requests("send_message").len(), + 1, + "a cancelled prepared upload never hands off after the other send settles" + ); + chat.update(cx, |chat, _| { + assert!( + chat.sending_previews + .values() + .all(|preview| preview.prepared().is_none()) + ); + assert!(chat.draft_attachments[0].committed.is_none()); + }); +} + +#[gpui::test] +fn uploaded_prompt_waits_for_history_recovery_before_admission(cx: &mut TestAppContext) { + let host = UploadHost::new(); + let (chat, cx) = chat_window(&host, cx); + start_upload(&chat, cx, &host); + chat.update(cx, |chat, _| { + chat.host_recovery.unavailable = true; + chat.runtime_error = Some("Connected. Refreshing task history…".into()); + }); + host.finish_upload(); + wait_prepared(&chat, cx); + assert!(host.requests("send_message").is_empty()); + chat.update(cx, |chat, cx| { + // The recovery snapshot installs these states together before the + // production completion callback dispatches prepared uploads. + chat.host_recovery.unavailable = false; + chat.runtime_error = None; + chat.dispatch_prepared_uploads(cx); + }); + wait_for(cx, "admission after history recovery", |_| { + host.requests("send_message").len() == 1 + }); + let sent = host.requests("send_message")[0].clone(); + assert_eq!(sent.params[1]["text"], "with file"); + assert_eq!(sent.params[1]["attachments"].as_array().unwrap().len(), 1); + chat.update(cx, |chat, cx| chat.dispatch_prepared_uploads(cx)); + host.reply(&sent, Ok(json!("run-1"))); + wait_for(cx, "settled upload send", |cx| { + chat.update(cx, |chat, _| !chat.send_pending()) + }); + assert_eq!(host.requests("send_message").len(), 1); +} diff --git a/apps/maple-agent/app/src/ui/chat/uploads.rs b/apps/maple-agent/app/src/ui/chat/uploads.rs new file mode 100644 index 000000000..9b83d10b8 --- /dev/null +++ b/apps/maple-agent/app/src/ui/chat/uploads.rs @@ -0,0 +1,733 @@ +//! Upload admission precedes prompt admission. Cancelling and handing off the +//! prompt compete for one atomic transition, so a cancelled upload never sends. + +use super::*; +use crate::backend::{AttachmentUploadRequest, ImportAttachmentRequest}; +use maple_agent::store::{AttachmentUploadStatus, StoredAttachment}; +use std::sync::{ + Arc, + atomic::{AtomicU8, AtomicUsize, Ordering}, +}; + +const UPLOADING: u8 = 0; +const CANCELLED: u8 = 1; +const HANDED_OFF: u8 = 2; + +#[derive(Default)] +struct UploadState { + phase: AtomicU8, + active_operation: AtomicUsize, + cancellation: tokio::sync::Notify, +} + +#[derive(Clone, Default)] +pub(super) struct UploadFence(Arc); +impl UploadFence { + pub fn cancel(&self) -> bool { + let cancelled = self + .0 + .phase + .compare_exchange(UPLOADING, CANCELLED, Ordering::SeqCst, Ordering::SeqCst) + .is_ok(); + if cancelled { + self.0.cancellation.notify_waiters(); + } + cancelled + } + + async fn wait_cancelled(&self) { + let notified = self.0.cancellation.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if !self.cancelled() { + notified.await; + } + } + + pub fn hand_off(&self) -> bool { + self.0 + .phase + .compare_exchange(UPLOADING, HANDED_OFF, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + } + pub fn handed_off(&self) -> bool { + self.0.phase.load(Ordering::SeqCst) == HANDED_OFF + } + fn set_active_operation(&self, index: usize) { + self.0.active_operation.store(index, Ordering::SeqCst); + } + fn active_operation(&self) -> usize { + self.0.active_operation.load(Ordering::SeqCst) + } + pub fn cancelled(&self) -> bool { + self.0.phase.load(Ordering::SeqCst) == CANCELLED + } +} + +pub(super) struct PendingUpload { + pub session: String, + pub count: usize, + pub label: SharedString, + pub fence: UploadFence, + pub cancel_on_drop: bool, + backend: Arc, + user: String, + operations: Vec, +} + +pub(super) struct PreparedSend { + pub request: AgentSendMessageRequest, + pub send: UnresolvedSend, + pub fence: UploadFence, +} + +impl Drop for PendingUpload { + fn drop(&mut self) { + // Sign-out, account replacement and erase-success tear down the view. + // Detached import work must lose its permission to send immediately. + if self.cancel_on_drop && self.fence.cancel() { + drop(self.backend.spawn(cancel_operations( + self.backend.clone(), + self.user.clone(), + self.session.clone(), + self.operations.clone(), + self.fence.active_operation(), + ))); + } + } +} + +async fn cancel_operations( + backend: Arc, + user: String, + session: String, + mut operations: Vec, + active: usize, +) -> Result<(), String> { + if !operations.is_empty() { + let active = active.min(operations.len() - 1); + operations.rotate_left(active); + } + for (index, operation_id) in operations.into_iter().enumerate() { + let request = AttachmentUploadRequest { + session_id: session.clone(), + operation_id, + }; + // The active stream's cancellation goes first and can briefly retry a + // busy control slot. Cleanup of the other max-nine operations follows + // sequentially, preserving the separate Stop/question control capacity. + let attempts = if index == 0 { 3 } else { 1 }; + let _ = tokio::time::timeout(std::time::Duration::from_secs(3), async { + for attempt in 0..attempts { + if backend + .cancel_attachment_upload(&user, request.clone()) + .await + .is_ok() + { + break; + } + if attempt + 1 < attempts { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } + }) + .await; + } + Ok(()) +} + +impl ChatScreen { + /// Upload completion is not prompt admission. A newer text send can be + /// settling (including an unknown receipt) while these bytes finish. + pub(super) fn dispatch_prepared_uploads(&mut self, cx: &mut Context) { + if !self.accept_send_results { + return; + } + let ready: Vec<_> = self + .sending_previews + .iter() + .filter(|&(id, preview)| { + !self.discarded_sends.contains(id) && preview.prepared().is_some() + }) + .map(|(id, _)| id.clone()) + .collect(); + for id in ready { + let Some(prepared) = self + .sending_previews + .get(&id) + .and_then(|preview| preview.prepared()) + else { + continue; + }; + let session = &prepared.send.session; + let recovering = self.host_recovery.unavailable + || self.host_recovery.ended + || self.booting + || self.runtime_error.is_some(); + if !prepared.fence.cancelled() + && (recovering + || self.sending_previews.values().any(|preview| { + preview.session.as_ref() == Some(session) && preview.admitting() + }) + || self + .unresolved_sends + .values() + .any(|send| &send.session == session && send.awaiting_receipt())) + { + if let Some(upload) = self.attachment_uploads.get_mut(&id) { + upload.label = if recovering { + "Files ready · waiting for the host to recover" + } else { + "Files ready · waiting for the previous send" + } + .into(); + } + continue; + } + let prepared = self + .sending_previews + .get_mut(&id) + .and_then(|preview| preview.take_prepared()) + .unwrap(); + self.admit_prepared_send(id, prepared, cx); + } + } + + pub(super) fn admit_prepared_send( + &mut self, + id: String, + prepared: PreparedSend, + cx: &mut Context, + ) { + let PreparedSend { + request, + send, + fence, + } = prepared; + if !fence.hand_off() { + self.finish_attachment_send( + id, + send.clone(), + Ok((send, Err("Upload cancelled. No message was sent.".into()))), + cx, + ); + return; + } + let backend = send.backend.clone(); + let user = self.user_id.clone(); + let send_id = id.clone(); + let retained = send.clone(); + self.call( + async move { + let result = backend + .send_message_identified(&user, request, send_id) + .await; + Ok((send, result)) + }, + cx, + move |this, result, cx| { + this.finish_attachment_send(id, retained, result, cx); + }, + ); + } + + pub(super) fn finish_attachment_send( + &mut self, + send_id: String, + mut retained: UnresolvedSend, + result: Result<(UnresolvedSend, Result), String>, + cx: &mut Context, + ) { + let mut pending = self.attachment_uploads.remove(&send_id); + let cancelled = pending + .as_ref() + .is_some_and(|upload| upload.fence.cancelled()); + if result.is_ok() { + // A normal failed/unknown upload retains its operation for receipt + // reconciliation. Only view teardown cancels it on drop. + if let Some(upload) = &mut pending { + upload.cancel_on_drop = false; + } + } else if pending + .as_ref() + .is_some_and(|upload| !upload.fence.handed_off()) + { + for draft in &mut retained.attachments { + draft.reset_upload(); + } + } + drop(pending); + if !self.accept_send_results || self.discarded_sends.remove(&send_id) { + return; + } + match result { + Ok((mut send, result)) => { + if cancelled { + for draft in &mut send.attachments { draft.reset_upload(); } + } + self.finish_send(send_id, send, result, cx); + } + Err(_) => self.finish_send(send_id, retained, + Err("Send outcome unknown. The background request ended before reporting its result.".into()), cx), + } + } + + pub(super) fn begin_uploads( + &mut self, + send_id: &str, + session: &str, + drafts: &[DraftAttachment], + cx: &mut Context, + ) -> (UploadFence, tokio::sync::mpsc::UnboundedSender) { + let fence = UploadFence::default(); + let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); + if !drafts.is_empty() { + self.attachment_uploads.insert( + send_id.to_string(), + PendingUpload { + session: session.to_string(), + count: drafts.len(), + label: format!("Uploading files · 0 of {}", drafts.len()).into(), + fence: fence.clone(), + cancel_on_drop: true, + backend: self.backend.clone(), + user: self.user_id.clone(), + operations: drafts + .iter() + .map(|draft| draft.operation_id.clone()) + .collect(), + }, + ); + let id = send_id.to_string(); + let task = cx.spawn(async move |this, cx| { + while let Some(completed) = receiver.recv().await { + if this + .update(cx, |this, cx| { + if let Some(upload) = this.attachment_uploads.get_mut(&id) { + if upload.fence.cancelled() { + return; + } + upload.label = + format!("Uploading files · {completed} of {}", upload.count) + .into(); + cx.notify(); + } + }) + .is_err() + { + break; + } + } + }); + crate::ui::task::retain(&self.bridged_tasks, task); + } + (fence, sender) + } + + pub(super) fn cancel_attachment_send(&mut self, send_id: &str, cx: &mut Context) { + let Some(upload) = self.attachment_uploads.get_mut(send_id) else { + return; + }; + if !upload.fence.cancel() { + return; + } + upload.label = "Cancelling upload…".into(); + let backend = upload.backend.clone(); + let user = upload.user.clone(); + let session = upload.session.clone(); + let operations = upload.operations.clone(); + let active = upload.fence.active_operation(); + // This runs concurrently with import, not behind its blocking await. + self.call( + cancel_operations(backend, user, session, operations, active), + cx, + |_, _, _| {}, + ); + self.dispatch_prepared_uploads(cx); + cx.notify(); + } + + pub(super) fn selected_upload(&self) -> Option<(&str, &PendingUpload)> { + self.attachment_uploads + .iter() + .find(|(_, upload)| { + self.selected_session.as_deref() == Some(upload.session.as_str()) + && !upload.fence.handed_off() + }) + .map(|(id, upload)| (id.as_str(), upload)) + } +} + +pub(super) async fn commit_attachments( + backend: &AgentBackend, + user: &str, + session: &str, + drafts: &mut [DraftAttachment], + fence: &UploadFence, + progress: tokio::sync::mpsc::UnboundedSender, +) -> Result, String> { + let mut ids = Vec::with_capacity(drafts.len()); + for (index, draft) in drafts.iter_mut().enumerate() { + fence.set_active_operation(index); + if fence.cancelled() { + return Err("Upload cancelled. No message was sent.".into()); + } + if let Some((committed_session, _)) = &draft.committed + && committed_session == session + { + // Unbound committed files can expire while a failed draft stays + // open. Only an authoritative Missing permits another import; + // a lost status reply must retain the cached identity and bytes. + let query = tokio::time::timeout( + std::time::Duration::from_secs(15), + backend.attachment_upload_status( + user, + AttachmentUploadRequest { + session_id: session.to_owned(), + operation_id: draft.operation_id.clone(), + }, + ), + ); + let status = tokio::select! { + result = query => result.map_err(|_| "Could not confirm the retained file upload. Your draft is retained; try again when the connection recovers.".to_string())??, + () = fence.wait_cancelled() => return Err("Upload cancelled. No message was sent.".into()), + }; + match status { + AttachmentUploadStatus::Complete(attachment) => { + draft.committed = Some((session.to_owned(), attachment.id.clone())); + ids.push(attachment.id); + let _ = progress.send(index + 1); + continue; + } + AttachmentUploadStatus::Missing => { + // Reuse the immutable operation as well as the staged + // bytes, keeping cancellation and duplicate-ACK handling + // attached to the same operation throughout the retry. + draft.committed = None; + } + AttachmentUploadStatus::Cancelled => return Err("The file upload was cancelled. Your draft is retained.".into()), + AttachmentUploadStatus::Uploading => return Err("The retained file upload is still in progress. Your draft is retained; try again when it completes.".into()), + } + } + let staged = draft + .staged + .as_ref() + .ok_or("Attachment is still being prepared")?; + let import = backend.import_attachment( + user, + ImportAttachmentRequest { + session_id: session.to_string(), + operation_id: draft.operation_id.clone(), + path: staged.path.to_string_lossy().into_owned(), + name: draft.name.clone(), + mime_type: draft.mime_type.clone(), + }, + ); + let imported = tokio::select! { + result = import => result, + () = fence.wait_cancelled() => return Err("Upload cancelled. No message was sent.".into()), + }; + let imported = match imported { + Ok(attachment) => attachment, + Err(error) if !fence.cancelled() => { + // An ACK can be lost after publication. Query the same host and + // immutable operation once; never guess or resend a prompt. + let query = tokio::time::timeout( + std::time::Duration::from_secs(15), + backend.attachment_upload_status( + user, + AttachmentUploadRequest { + session_id: session.to_string(), + operation_id: draft.operation_id.clone(), + }, + ), + ); + let status = tokio::select! { + result = query => result.ok().and_then(Result::ok), + () = fence.wait_cancelled() => return Err("Upload cancelled. No message was sent.".into()), + }; + resolve_upload_receipt(error, status)? + } + Err(_) => return Err("Upload cancelled. No message was sent.".into()), + }; + if fence.cancelled() { + // The bounded cancellation worker handles publication races. Return + // the draft promptly; never wait for a second cancellation RPC. + return Err("Upload cancelled. No message was sent.".into()); + } + draft.committed = Some((session.to_string(), imported.id.clone())); + ids.push(imported.id); + let _ = progress.send(index + 1); + } + Ok(ids) +} + +fn resolve_upload_receipt( + error: String, + status: Option, +) -> Result { + match status { + Some(AttachmentUploadStatus::Complete(attachment)) => Ok(attachment), + Some(AttachmentUploadStatus::Cancelled) => { + Err("The file upload was cancelled. Your draft is retained.".into()) + } + Some(AttachmentUploadStatus::Uploading) | None => Err(format!( + "File upload could not be confirmed: {error}. Your draft is retained; try again when the connection recovers." + )), + Some(AttachmentUploadStatus::Missing) => Err(error), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Default)] + struct UploadHost { + complete: HashMap, + imports: Vec<(String, String, Vec)>, + fail_second: bool, + fail_status: bool, + } + + // Exercise the production client and upload orchestration. The fixture's + // missing receipt represents an unbound file removed by store maintenance; + // store tests separately exercise the 24-hour collection policy. + fn upload_host( + root: &std::path::Path, + ) -> ( + AgentBackend, + Arc>, + tokio::task::JoinHandle<()>, + ) { + use crate::host::{RemoteClient, ipc, protocol}; + let (client, mut server) = tokio::io::duplex(64 * 1024); + let state = Arc::new(std::sync::Mutex::new(UploadHost::default())); + let shared = state.clone(); + let task = tokio::spawn(async move { + while let Ok(request) = protocol::read::(&mut server).await { + let result = { + let mut state = shared.lock().unwrap(); + match request.method.as_str() { + "import_attachment" => { + let request: ImportAttachmentRequest = + serde_json::from_value(request.params[1].clone()).unwrap(); + if state.fail_second && request.name == "second.txt" { + Err("Synthetic second-file failure".into()) + } else { + let bytes = std::fs::read(&request.path).unwrap(); + state.imports.push(( + request.operation_id.clone(), + request.name.clone(), + bytes.clone(), + )); + let descriptor = StoredAttachment { + id: format!("file-{}", state.imports.len()), + name: request.name, + mime_type: request.mime_type, + source: "maple-attachment://fixture".into(), + size: bytes.len() as u64, + execution_path: request.path, + }; + state + .complete + .insert(request.operation_id, descriptor.clone()); + Ok(serde_json::to_value(descriptor).unwrap()) + } + } + "attachment_upload_status" => { + let request: AttachmentUploadRequest = + serde_json::from_value(request.params[1].clone()).unwrap(); + if state.fail_status { + Err("Synthetic unavailable receipt".into()) + } else { + Ok(serde_json::to_value( + state + .complete + .get(&request.operation_id) + .cloned() + .map(AttachmentUploadStatus::Complete) + .unwrap_or(AttachmentUploadStatus::Missing), + ) + .unwrap()) + } + } + _ => Err("Unexpected fixture request".into()), + } + }; + if protocol::write( + &mut server, + &protocol::Message::Response { + id: request.id, + result, + }, + ) + .await + .is_err() + { + break; + } + } + }); + let client = RemoteClient::attach_for_test( + ipc::Endpoint::for_test(root.to_owned()), + Box::new(client), + protocol::Welcome { + version: protocol::VERSION, + pid: std::process::id(), + build: protocol::BUILD.into(), + user_id: Some("user".into()), + store: maple_agent::store::StoreState::Open, + }, + ); + (AgentBackend::from_host_client_for_test(client), state, task) + } + + #[tokio::test] + async fn expired_partial_upload_reimports_staged_bytes_only_after_authoritative_missing() { + let root = tempfile::tempdir().unwrap(); + let original = root.path().join("original.txt"); + std::fs::write(&original, b"original bytes").unwrap(); + let (staged, size) = staging::copy_file(&root.path().join("drafts"), &original).unwrap(); + let mut first = DraftAttachment::preparing(1, "first.txt".into()); + first.staged = Some(staged); + first.size = size; + first.mime_type = "text/plain".into(); + let mut second = DraftAttachment::preparing(2, "second.txt".into()); + second.staged = + Some(staging::copy_bytes(&root.path().join("drafts"), b"second bytes").unwrap()); + second.mime_type = "text/plain".into(); + let mut drafts = vec![first, second]; + let (backend, host, task) = upload_host(root.path()); + host.lock().unwrap().fail_second = true; + let (progress, _) = tokio::sync::mpsc::unbounded_channel(); + let failure = commit_attachments( + &backend, + "user", + "task", + &mut drafts, + &UploadFence::default(), + progress, + ) + .await + .unwrap_err(); + assert!(failure.contains("second-file failure")); + let committed = drafts[0].committed.clone().unwrap(); + let operation = drafts[0].operation_id.clone(); + assert!(drafts[1].committed.is_none()); + + // The execution copy expires while the retained desktop draft stays + // alive. Changing the source must not change the retry's bytes. + std::fs::write(&original, b"changed original").unwrap(); + { + let mut host = host.lock().unwrap(); + host.complete.remove(&operation); + host.fail_status = true; + host.fail_second = false; + } + let (progress, _) = tokio::sync::mpsc::unbounded_channel(); + assert!( + commit_attachments( + &backend, + "user", + "task", + &mut drafts, + &UploadFence::default(), + progress + ) + .await + .is_err() + ); + assert_eq!(drafts[0].committed.as_ref(), Some(&committed)); + assert_eq!(drafts[0].operation_id, operation); + assert_eq!( + host.lock().unwrap().imports.len(), + 1, + "an unavailable receipt must never trigger reimport" + ); + + host.lock().unwrap().fail_status = false; + let (progress, _) = tokio::sync::mpsc::unbounded_channel(); + let ids = commit_attachments( + &backend, + "user", + "task", + &mut drafts, + &UploadFence::default(), + progress, + ) + .await + .unwrap(); + assert_eq!(ids.len(), 2); + assert_ne!(ids[0], committed.1); + assert_eq!( + drafts[0].operation_id, operation, + "retry keeps the cancellation identity" + ); + let imports = &host.lock().unwrap().imports; + assert_eq!(imports.len(), 3); + assert_eq!(imports[1].2, b"original bytes"); + assert_eq!(imports[2].2, b"second bytes"); + task.abort(); + } + #[test] + fn lost_upload_ack_resolves_only_a_committed_receipt() { + let descriptor = StoredAttachment { + id: "file-id".into(), + name: "report.zip".into(), + mime_type: "application/zip".into(), + source: "maple-attachment://file-id".into(), + size: 123, + execution_path: "/uploads/report.zip".into(), + }; + assert_eq!( + resolve_upload_receipt( + "connection lost".into(), + Some(AttachmentUploadStatus::Complete(descriptor.clone())) + ) + .unwrap(), + descriptor + ); + for status in [ + None, + Some(AttachmentUploadStatus::Uploading), + Some(AttachmentUploadStatus::Missing), + Some(AttachmentUploadStatus::Cancelled), + ] { + assert!(resolve_upload_receipt("connection lost".into(), status).is_err()); + } + } + + #[test] + fn cancellation_wins_before_prompt_handoff() { + let fence = UploadFence::default(); + assert!(fence.cancel()); + assert!(!fence.hand_off()); + assert!(fence.cancelled()); + } + + #[tokio::test] + async fn cancellation_wakes_pending_and_later_upload_waiters() { + let fence = UploadFence::default(); + let waiting = fence.clone(); + let waiter = tokio::spawn(async move { waiting.wait_cancelled().await }); + tokio::task::yield_now().await; + assert!(fence.cancel()); + tokio::time::timeout(std::time::Duration::from_secs(1), waiter) + .await + .expect("cancellation must wake the in-flight import") + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), fence.wait_cancelled()) + .await + .expect("an already-cancelled import must not miss the notification"); + assert!(!fence.hand_off()); + } + + #[test] + fn prompt_handoff_cannot_be_reported_as_cancelled() { + let fence = UploadFence::default(); + assert!(fence.hand_off()); + assert!(!fence.cancel()); + assert!(!fence.cancelled()); + } +} diff --git a/apps/maple-agent/app/src/ui/remote.rs b/apps/maple-agent/app/src/ui/remote.rs index c45621bfd..8ac2350ee 100644 --- a/apps/maple-agent/app/src/ui/remote.rs +++ b/apps/maple-agent/app/src/ui/remote.rs @@ -45,6 +45,7 @@ enum PeerInfoPhase { Loaded, Unsupported, Unreachable, + Incompatible, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -95,6 +96,7 @@ struct InstallIntent { } const UNCONFIRMED_INSTALL: &str = "The update request could not be confirmed. Do not resend it. Refresh this computer’s status, or check the remote computer directly."; +const PROTOCOL_UPDATE_REQUIRED: &str = "Update required. The computers use incompatible remote protocols. Update Maple on each computer before reconnecting."; const CLIENT_VERSION: &str = env!("CARGO_PKG_VERSION"); const CLIENT_VERSION_LABEL: &str = concat!( "Client v", @@ -155,25 +157,40 @@ impl ComputerUpdate { PeerInfoPhase::Loading => "Reading update status…", PeerInfoPhase::Unsupported => "Remote update status is unavailable. Use Open projects for ordinary remote access.", PeerInfoPhase::Unreachable => "Could not read update status. Check that this computer is online, then refresh.", + PeerInfoPhase::Incompatible => PROTOCOL_UPDATE_REQUIRED, _ => "Refresh to read this computer’s version and update status.", }.into(); self.blocked_label = None; self.version_warning = None; return; }; + let prefix = if self.phase == PeerInfoPhase::Loaded { + "Running" + } else { + "Last reported" + }; self.running_label = match &info.running.build { - Some(build) => format!("Running v{} · build {build}", info.running.version), - None => format!("Running v{}", info.running.version), + Some(build) => format!("{prefix} v{} · build {build}", info.running.version), + None => format!("{prefix} v{}", info.running.version), }; self.update_label = if self.phase == PeerInfoPhase::Loading { "Refreshing update status…".into() } else if self.phase == PeerInfoPhase::Unreachable { "Could not refresh update status. The last reported version is shown above.".into() + } else if self.phase == PeerInfoPhase::Incompatible { + PROTOCOL_UPDATE_REQUIRED.into() } else if let Some(offered) = &info.update.offered { format!("v{} · {}", offered.version, info.update.state.label()) } else { info.update.state.label().into() }; + // Cached observations remain useful as history, but must not look like + // a current offer, policy decision or version comparison after failure. + if self.phase != PeerInfoPhase::Loaded { + self.version_warning = None; + self.blocked_label = None; + return; + } self.version_warning = version_warning(info, CLIENT_VERSION); let blocked = install_blocked(info); let healthy_without_offer = info.update.offered.is_none() @@ -509,7 +526,13 @@ impl RemoteScreen { update.phase = PeerInfoPhase::Unsupported; update.info = None; } - Err(_) => update.phase = PeerInfoPhase::Unreachable, + Err(error) => { + update.phase = if crate::remote::wire::is_protocol_incompatibility_refusal(&error) { + PeerInfoPhase::Incompatible + } else { + PeerInfoPhase::Unreachable + }; + } } update.refresh_presentation(); self.invalidate_confirmation(); diff --git a/apps/maple-agent/app/src/ui/remote/tests.rs b/apps/maple-agent/app/src/ui/remote/tests.rs index f91d4a151..d620b83b7 100644 --- a/apps/maple-agent/app/src/ui/remote/tests.rs +++ b/apps/maple-agent/app/src/ui/remote/tests.rs @@ -427,6 +427,103 @@ fn version_advisory_clears_for_equal_unknown_or_unsupported_peer_versions(cx: &m assert!(cx.debug_bounds("remote-version-warning-a").is_none()); } +#[gpui::test] +fn peer_status_distinguishes_protocol_refusal_from_offline_revoked_and_invalid_frames( + cx: &mut TestAppContext, +) { + use crate::remote::wire::{Error, Refusal}; + + let screen = screen(cx); + let cases = [ + ( + Refusal::terminal(Error::Incompatible.to_string()).into_error(), + PeerInfoPhase::Incompatible, + ), + (Error::Disconnected.to_string(), PeerInfoPhase::Unreachable), + (Error::TimedOut.to_string(), PeerInfoPhase::Unreachable), + ( + Refusal::terminal("This computer is no longer paired").into_error(), + PeerInfoPhase::Unreachable, + ), + ( + Refusal::terminal(Error::InvalidFrame.to_string()).into_error(), + PeerInfoPhase::Unreachable, + ), + ]; + screen.update(cx, |screen, _| { + seed_update(screen, "b", [2; 16]); + for has_cached_info in [false, true] { + for (error, phase) in &cases { + seed_update(screen, "a", [1; 16]); + if !has_cached_info { + screen.computer_updates.get_mut("a").unwrap().info = None; + } + screen.apply_peer_info("a", [1; 16], 1, Err(error.clone())); + let update = &screen.computer_updates["a"]; + assert_eq!(update.phase, *phase, "{error}"); + assert_eq!( + update.update_label.starts_with("Update required"), + *phase == PeerInfoPhase::Incompatible, + "{error}" + ); + assert!(!update.install_available); + assert!(update.version_warning.is_none()); + assert!(update.blocked_label.is_none()); + if has_cached_info { + assert_eq!(update.running_label, "Last reported v0.0.6 · build 550"); + assert!(update.info.is_some(), "the old report remains historical"); + } else { + assert_eq!(update.running_label, "Version unavailable"); + } + assert!(screen.computer_updates["b"].install_available); + } + } + // Only a new successful report restores current status and install eligibility. + screen.apply_peer_info("a", [1; 16], 1, Ok(Some(ready_peer()))); + let update = &screen.computer_updates["a"]; + assert_eq!(update.phase, PeerInfoPhase::Loaded); + assert!(update.running_label.starts_with("Running v")); + assert!(update.install_available); + }); +} + +#[gpui::test] +fn protocol_refusal_cancels_consent_and_disables_cached_remote_update_offer( + cx: &mut TestAppContext, +) { + let (screen, cx) = show_update(cx); + let install = cx.debug_bounds("update-remote-a").unwrap(); + cx.simulate_click(install.center(), gpui::Modifiers::default()); + cx.run_until_parked(); + assert!(cx.debug_bounds("remote-update-confirmation").is_some()); + screen.update(cx, |screen, cx| { + screen.apply_peer_info( + "a", + [1; 16], + 1, + Err(crate::remote::wire::Refusal::terminal( + crate::remote::wire::Error::Incompatible.to_string(), + ) + .into_error()), + ); + assert!(screen.take_install_intent().is_none()); + assert!(screen.computer_updates["a"].attempt.is_none()); + cx.notify(); + }); + cx.run_until_parked(); + assert!(cx.debug_bounds("remote-update-confirmation").is_none()); + assert!(cx.debug_bounds("remote-version-warning-a").is_none()); + let install = cx.debug_bounds("update-remote-a").unwrap(); + cx.simulate_click(install.center(), gpui::Modifiers::default()); + cx.run_until_parked(); + assert!(cx.debug_bounds("remote-update-confirmation").is_none()); + screen.read_with(cx, |screen, _| { + assert!(!screen.computer_updates["a"].install_available); + assert!(screen.computer_updates["a"].attempt.is_none()); + assert!(screen.status.as_ref().unwrap().peers["a"].can_connect); + }); +} + #[gpui::test] fn installing_a_different_target_warns_without_requiring_the_client_to_update( cx: &mut TestAppContext, @@ -739,6 +836,16 @@ fn old_peer_probe_cannot_replace_newer_generation_or_pairing(cx: &mut TestAppCon screen.computer_updates.get_mut("a").unwrap().generation = 2; screen.apply_peer_info("a", [1; 16], 1, Ok(None)); assert_eq!(screen.computer_updates["a"].phase, PeerInfoPhase::Loaded); + screen.apply_peer_info( + "a", + [1; 16], + 1, + Err(crate::remote::wire::Refusal::terminal( + crate::remote::wire::Error::Incompatible.to_string(), + ) + .into_error()), + ); + assert_eq!(screen.computer_updates["a"].phase, PeerInfoPhase::Loaded); let mut status = screen.status.clone().unwrap(); status.peers.get_mut("a").unwrap().grant = [2; 16]; screen.update_status(status, cx); diff --git a/apps/maple-agent/app/src/ui/settings/advanced.rs b/apps/maple-agent/app/src/ui/settings/advanced.rs index fcd9fac95..3ff6fcd69 100644 --- a/apps/maple-agent/app/src/ui/settings/advanced.rs +++ b/apps/maple-agent/app/src/ui/settings/advanced.rs @@ -577,8 +577,8 @@ impl SettingsScreen { }), ); pane=pane.child(div().text_sm().text_color(gpui::rgb(theme::text_secondary())).child(if password { - "Maple's copy of your tasks, settings and sign-in on this computer is encrypted with your unlock password. Maple asks for it each time it starts." - } else {"Your tasks and settings are stored on this computer, readable only by your user account. Set an unlock password to encrypt them."})); + "Maple's copy of your conversation text, settings and sign-in on this computer is encrypted with your unlock password. Maple asks for it each time it starts. Attached files and images remain unencrypted." + } else {"Your conversation text, settings and sign-in are stored on this computer, readable only by your user account. Set an unlock password to encrypt them. Attached files and images remain unencrypted."})); if self.advanced.erase { pane = pane .child(div().text_sm().child(ERASE_COPY)) diff --git a/apps/maple-agent/app/src/ui/store_gate.rs b/apps/maple-agent/app/src/ui/store_gate.rs index 691a4638b..fb0595423 100644 --- a/apps/maple-agent/app/src/ui/store_gate.rs +++ b/apps/maple-agent/app/src/ui/store_gate.rs @@ -1419,7 +1419,7 @@ impl Render for PasswordForm { let (title, description, submit_label) = match self.action { PasswordAction::Set | PasswordAction::Setup => ( "Enable encryption", - "Protect the tasks, attachments, settings and sign-in saved on this computer.", + "Protect saved conversation text, settings and sign-in. Attached files and images remain unencrypted on this computer.", "Enable encryption", ), PasswordAction::Change => ( diff --git a/apps/maple-agent/crates/maple-agent/src/acp/mod.rs b/apps/maple-agent/crates/maple-agent/src/acp/mod.rs index 2bcbb62b9..06cf1ae46 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/mod.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/mod.rs @@ -162,11 +162,16 @@ where pub(super) struct PromptAdmission { pub(super) session_id: String, prompt: String, - images: Vec, + images: PromptAttachments, cancellation: CancellationToken, operation: Arc, } +enum PromptAttachments { + Images(Vec), + Committed(Vec), +} + /// What ended a turn's stream of events early. enum StreamStop { /// The run's events overflowed the stream. @@ -380,7 +385,8 @@ impl AcpConnectionContext { agent_client_protocol::Error::invalid_request() .data("This ACP connection has reached its session limit") })?; - let operation = AcpSessionOperation::new(&self.lifetime); + let history_epoch = self.session_history_epoch(session_id).await?; + let operation = AcpSessionOperation::new(&self.lifetime, history_epoch); let operation_guard = Arc::clone(&operation.gate).lock_owned().await; let replay_guard = Arc::clone(&operation.replay_gate).lock_owned().await; let _finalization = self.finalization.lock().await; @@ -397,6 +403,27 @@ impl AcpConnectionContext { operations.insert(session_id.to_owned(), operation.clone()); drop(operations); drop(sessions); + // A replacement may have been dispatched after the first stamp read + // but before insertion. Recheck with the operation now registered so + // either this read or the observer catches every newer epoch. + let validation = self + .session_history_epoch(session_id) + .await + .and_then(|epoch| { + if epoch == history_epoch { + Ok(()) + } else { + Err(paged::replay_error( + "This task changed while its session was loading; reload the session.", + )) + } + }); + if let Err(error) = validation { + drop(_finalization); + self.remove_session_operation_if_same(session_id, &operation) + .await; + return Err(error); + } let mut tasks = self.background_tasks.lock().await; while tasks.try_join_next().is_some() {} self.spawn_session_observer( @@ -409,6 +436,30 @@ impl AcpConnectionContext { Ok((operation, operation_guard, replay_guard)) } + async fn session_history_epoch( + &self, + session_id: &str, + ) -> Result { + let (account, store) = self + .agent + .history_replay_source() + .map_err(internal_acp_error)?; + let task = session_id.to_owned(); + let epoch = tokio::task::spawn_blocking(move || { + store + .task_history_stamp(account, &task) + .map(|stamp| stamp.rows_epoch) + }) + .await + .map_err(|_| internal_acp_error("ACP history stamp worker failed".to_string()))? + .map_err(|error| internal_acp_error(error.to_string()))?; + self.agent + .verify_generation() + .await + .map_err(internal_acp_error)?; + Ok(epoch) + } + /// Block until the lazy runtime start finished. The handshake answers /// before the runtime boots; the first session request pays the boot. async fn await_runtime_start(&self) -> Result<(), agent_client_protocol::Error> { @@ -551,6 +602,12 @@ impl AcpConnectionContext { } let session_servers = prepare_session_mcp(&request.mcp_servers)?; let protocol_session_id = SessionId::new(session_id.clone()); + // Rebuilding announces HistoryReplaced. Finish before registering this + // load; its epoch baseline also filters delayed replacement events. + self.agent + .ensure_task_history_ready(&session_id) + .await + .map_err(internal_acp_error)?; let (operation, operation_guard, replay_guard) = self.start_session_operation(cx, &session_id).await?; let attached = match self @@ -1084,7 +1141,7 @@ impl AcpConnectionContext { Ok(PromptAdmission { session_id, prompt, - images, + images: PromptAttachments::Images(images), cancellation, operation, }) @@ -1442,7 +1499,7 @@ impl AcpConnectionContext { cx: &ConnectionTo, session_id: &str, prompt: String, - images: Vec, + images: PromptAttachments, cancellation: &CancellationToken, ) -> Result { let protocol_session_id = SessionId::new(session_id.to_string()); @@ -1523,6 +1580,17 @@ impl AcpConnectionContext { prompt = expanded; } } + let attachments = match images { + PromptAttachments::Images(images) => self + .agent + .import_image_attachments(session_id.to_string(), images) + .await + .map_err(internal_acp_error)? + .into_iter() + .map(|attachment| attachment.id) + .collect(), + PromptAttachments::Committed(ids) => ids, + }; let run = match self .agent .send_surface_message( @@ -1537,7 +1605,7 @@ impl AcpConnectionContext { vision_capable: catalog.vision == Some(true), steer: false, queue_id: None, - attachments: images, + attachments, }, cancellation.clone(), ) diff --git a/apps/maple-agent/crates/maple-agent/src/acp/observer.rs b/apps/maple-agent/crates/maple-agent/src/acp/observer.rs index 190edda2c..92bb988fb 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/observer.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/observer.rs @@ -259,10 +259,11 @@ impl AcpConnectionContext { Event::HistoryReplaced { user_id, session_id, - .. + rows_epoch, } => { if user_id == self.agent.user_id() && let Some(operation) = self.session_operations.lock().await.get(&session_id) + && rows_epoch > operation.history_epoch { operation.fail_observer(ObserverFailure::HistoryChanged); } diff --git a/apps/maple-agent/crates/maple-agent/src/acp/observer_tests.rs b/apps/maple-agent/crates/maple-agent/src/acp/observer_tests.rs index 28460d779..4a4447e26 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/observer_tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/observer_tests.rs @@ -105,7 +105,7 @@ async fn a_stale_observer_cannot_retire_a_reloaded_session_with_the_same_id() { let stale = context.session_operations.lock().await[&task].clone(); // Install a replacement generation while retaining the old worker's Arc. // Its late cleanup must not remove the new maps, lease, or session count. - let replacement = AcpSessionOperation::new(&context.lifetime); + let replacement = AcpSessionOperation::new(&context.lifetime, stale.history_epoch); context .session_operations .lock() @@ -129,6 +129,73 @@ async fn a_stale_observer_cannot_retire_a_reloaded_session_with_the_same_id() { finish(serving).await; } +#[tokio::test] +async fn delayed_history_replacements_at_or_before_the_registered_epoch_are_ignored() { + let harness = Harness::new().await; + let task = harness.create_task().await; + append_test_messages( + &harness.handle, + &task, + harness.project.path(), + &["Saved history".to_owned()], + ); + let (store, account) = harness.handle.test_persistence().unwrap(); + store.retire_task_timeline_for_test(account, &task).unwrap(); + harness + .handle + .ensure_task_history_ready(&task) + .await + .unwrap(); + let (mut client, serving, context) = connect_with_context(&harness.handle, no_answers()); + client.initialize(json!({})).await; + client + .call( + "session/load", + json!({"sessionId":task,"cwd":harness.project.path(),"mcpServers":[]}), + ) + .await; + let operation = context.session_operations.lock().await[&task].clone(); + assert!(operation.history_epoch > 0); + assert_eq!( + operation.history_epoch, + store.task_history_stamp(account, &task).unwrap().rows_epoch + ); + let replay = Arc::clone(&operation.replay_gate).lock_owned().await; + operation.replaying.store(true, Ordering::Release); + // Deliberately dispatch after registration, while replay owns its gate. + // These are the notifications a completed preparation can leave queued. + for rows_epoch in [operation.history_epoch - 1, operation.history_epoch] { + let before = context.observer_dispatches.load(Ordering::Acquire); + harness.emit_service_event(crate::agent::AgentServiceEvent::HistoryReplaced { + user_id: context.agent.user_id().to_owned(), + session_id: task.clone(), + rows_epoch, + }); + eventually(|| context.observer_dispatches.load(Ordering::Acquire) > before).await; + assert!(!operation.cancellation.is_cancelled()); + assert_eq!(operation.failure_message(), None); + assert!( + context + .session_operation_is_current(&task, &operation) + .await + ); + } + operation.replaying.store(false, Ordering::Release); + drop(replay); + harness.faux.push_text("The session remains attached."); + assert_eq!( + client.prompt(&task, "Continue.").await["stopReason"], + "end_turn" + ); + assert_eq!( + client.texts("agent_message_chunk")[0].1, + "The session remains attached." + ); + assert!(!context.lifetime.is_cancelled()); + client.shutdown().await; + finish(serving).await; +} + #[tokio::test] async fn a_history_replacement_invalidates_only_its_own_operation() { let harness = Harness::new().await; @@ -153,7 +220,7 @@ async fn a_history_replacement_invalidates_only_its_own_operation() { let event = crate::agent::AgentServiceEvent::HistoryReplaced { user_id: context.agent.user_id().to_owned(), session_id: changed.clone(), - rows_epoch: 2, + rows_epoch: changed_operation.history_epoch + 1, }; // The private fixture emitter exposes timeline events only. Publish through // the same host dispatcher the runtime uses for this event variant. diff --git a/apps/maple-agent/crates/maple-agent/src/acp/session.rs b/apps/maple-agent/crates/maple-agent/src/acp/session.rs index 1d9275086..6c54ed23a 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/session.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/session.rs @@ -120,16 +120,20 @@ pub(super) struct AcpSessionOperation { pub(super) replay_gate: Arc>, pub(super) observer: ObserverMailbox, pub(super) replaying: AtomicBool, + /// Replacement events can wait in the connection's broadcast receiver + /// before this operation is registered. Only a newer epoch invalidates it. + pub(super) history_epoch: i64, pub(super) cancellation: CancellationToken, } impl AcpSessionOperation { - pub(super) fn new(connection_lifetime: &CancellationToken) -> Arc { + pub(super) fn new(connection_lifetime: &CancellationToken, history_epoch: i64) -> Arc { Arc::new(Self { gate: Arc::new(Mutex::new(())), replay_gate: Arc::new(Mutex::new(())), observer: ObserverMailbox::new(), replaying: AtomicBool::new(true), + history_epoch, cancellation: connection_lifetime.child_token(), }) } diff --git a/apps/maple-agent/crates/maple-agent/src/acp/shared.rs b/apps/maple-agent/crates/maple-agent/src/acp/shared.rs index 51b39b8d1..49bd6b742 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/shared.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/shared.rs @@ -49,6 +49,14 @@ impl AcpConnectionContext { .model_catalog_entry(&model) .await .map_err(internal_acp_error)?; + let attachments = self + .agent + .import_image_attachments(session_id.clone(), prompt_images(&request.prompt)) + .await + .map_err(internal_acp_error)? + .into_iter() + .map(|attachment| attachment.id) + .collect::>(); let injected = self .agent .steer_surface_message( @@ -56,7 +64,7 @@ impl AcpConnectionContext { AgentSendMessageRequest { session_id: session_id.clone(), text: prompt_text(&request.prompt)?, - attachments: prompt_images(&request.prompt), + attachments: attachments.clone(), model: Some(model), context_limit: None, vision_capable: catalog.and_then(|entry| entry.vision) == Some(true), @@ -75,7 +83,8 @@ impl AcpConnectionContext { // Legacy adapters start a new turn when steering races with completion. // Clients opting into promptRequired retain ownership of that fallback. let prompt = PromptRequest::new(request.session_id, request.prompt); - let admission = self.begin_prompt(&prompt).await?; + let mut admission = self.begin_prompt(&prompt).await?; + admission.images = PromptAttachments::Committed(attachments); let context = Arc::clone(self); let cx = cx.clone(); self.background_tasks.lock().await.spawn(async move { diff --git a/apps/maple-agent/crates/maple-agent/src/acp/tests.rs b/apps/maple-agent/crates/maple-agent/src/acp/tests.rs index e4d378641..d47455f74 100644 --- a/apps/maple-agent/crates/maple-agent/src/acp/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/acp/tests.rs @@ -303,6 +303,64 @@ async fn session_load_replays_more_than_two_pages_with_distinct_duplicate_key_me finish(serving).await; } +#[tokio::test] +async fn session_load_prepares_retired_history_before_registering_and_can_prompt() { + let harness = Harness::new().await; + let task = harness.create_task().await; + let texts = vec![ + "Saved question one".to_owned(), + "Saved question two".to_owned(), + ]; + append_test_messages(&harness.handle, &task, harness.project.path(), &texts); + let (store, account) = harness.handle.test_persistence().unwrap(); + let (mut client, serving, context) = connect_with_context(&harness.handle, no_answers()); + client.initialize(json!({})).await; + let old_epoch = store.task_history_stamp(account, &task).unwrap().rows_epoch; + store.retire_task_timeline_for_test(account, &task).unwrap(); + assert!(store.history_needs_rebuild(account, &task).unwrap()); + + // Block registration at its normal publication fence. Preparation must + // finish before reaching that fence, without a registered observer for + // its own replacement event to invalidate. + let registration = context.finalization.lock().await; + let request = client + .request( + "session/load", + json!({"sessionId":task,"cwd":harness.project.path(),"mcpServers":[]}), + ) + .await; + eventually(|| !store.history_needs_rebuild(account, &task).unwrap()).await; + let prepared_epoch = store.task_history_stamp(account, &task).unwrap().rows_epoch; + assert!(prepared_epoch > old_epoch); + assert!(!context.session_operations.lock().await.contains_key(&task)); + drop(registration); + + let response = client.response(request).await; + assert!( + response.get("error").is_none(), + "session/load failed: {response}" + ); + let replayed = client.texts("user_message_chunk"); + assert_eq!( + replayed.iter().map(|(_, text)| text).collect::>(), + texts.iter().collect::>() + ); + let operation = context.session_operations.lock().await[&task].clone(); + assert_eq!(operation.history_epoch, prepared_epoch); + assert!(!operation.cancellation.is_cancelled()); + harness.faux.push_text("The loaded session can answer."); + assert_eq!( + client.prompt(&task, "Continue after loading.").await["stopReason"], + "end_turn" + ); + assert_eq!( + client.texts("agent_message_chunk")[0].1, + "The loaded session can answer." + ); + client.shutdown().await; + finish(serving).await; +} + #[tokio::test] async fn oversized_replay_text_is_complete_and_every_encoded_frame_is_bounded() { let harness = Harness::new().await; @@ -960,7 +1018,7 @@ async fn prompt_images_travel_through_the_read_image_helper() { .find(|item| item.role.as_deref() == Some("user")) .unwrap(); assert_eq!(item.text.as_deref(), Some("What do you see?")); - let attachment = &item.input.as_ref().unwrap()["imageAttachments"][0]; + let attachment = &item.input.as_ref().unwrap()["attachments"][0]; assert_eq!(attachment["name"], "acp-image-1.png"); assert!( attachment["source"] @@ -1121,7 +1179,7 @@ fn session_load_preflight_rejects_missing_tasks() { #[tokio::test] async fn session_operation_cancellation_keeps_close_behind_the_active_fence() { let connection_lifetime = CancellationToken::new(); - let operation = AcpSessionOperation::new(&connection_lifetime); + let operation = AcpSessionOperation::new(&connection_lifetime, 1); let active = Arc::clone(&operation.gate).lock_owned().await; let waiting_gate = Arc::clone(&operation.gate); let waiting = tokio::spawn(async move { waiting_gate.lock_owned().await }); diff --git a/apps/maple-agent/crates/maple-agent/src/agent.rs b/apps/maple-agent/crates/maple-agent/src/agent.rs index 044e4fb64..8da117fa8 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent.rs @@ -42,7 +42,7 @@ use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::{Arc, RwLock}; use crate::maple_api::{MapleApiSession, account_scope}; -pub use attachments::AgentImageUpload; +pub use attachments::{AgentAttachment, AgentImageUpload}; use config::*; pub use config::{ MAPLE_WORKSPACE_DIRECTORY_NAME, MAPLE_WORKSPACE_DISPLAY_NAME, default_maple_workspace_path, diff --git a/apps/maple-agent/crates/maple-agent/src/agent/attachments.rs b/apps/maple-agent/crates/maple-agent/src/agent/attachments.rs index ee9a840d7..7c3e71d7d 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/attachments.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/attachments.rs @@ -2,10 +2,11 @@ use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD}; use serde::{Deserialize, Serialize}; use std::sync::Arc; -use crate::store::{AccountId, ImageUpload, Store}; +use crate::store::{AccountId, ImageUpload, Store, StoreError}; pub(super) const MAX_AGENT_IMAGE_BYTES: usize = 10 * 1024 * 1024; pub(super) const MAX_AGENT_IMAGES_PER_MESSAGE: usize = 10; +const MAX_STORED_IMAGE_BYTES: usize = 20 * 1024 * 1024; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -16,21 +17,17 @@ pub struct AgentImageUpload { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AgentImageAttachment { +pub struct AgentAttachment { pub id: String, pub name: String, pub mime_type: String, pub source: String, + pub size: u64, + pub execution_path: String, } -#[derive(Debug, Clone)] -pub(super) struct PreparedAgentImage { - pub attachment: AgentImageAttachment, - pub base64_data: String, -} - -/// The runtime's account-bound view of the host-owned SQL attachment store. -/// It never resolves an app data path or opens a previous attachment cache. +/// The runtime's account-bound view of the host-owned file attachment store. +/// Every resolution checks the execution account and task. #[derive(Clone)] pub(super) struct AgentAttachmentStore { store: Arc, @@ -46,15 +43,15 @@ impl AgentAttachmentStore { &self, session_id: &str, uploads: &[AgentImageUpload], - ) -> Result, String> { + ) -> Result, String> { if uploads.len() > MAX_AGENT_IMAGES_PER_MESSAGE { return Err(format!( "Attach at most {MAX_AGENT_IMAGES_PER_MESSAGE} images at a time" )); } - // Keep current validation messages and the model's original base64. The - // privileged store independently validates the same batch before writes. - let encoded = uploads + // Keep ACP's image validation messages. The privileged store independently + // validates the entire batch before writing files. + uploads .iter() .map(|upload| { let (declared_mime, base64_data) = parse_data_url(&upload.data_url)?; @@ -71,7 +68,7 @@ impl AgentAttachmentStore { if normalize_mime(declared_mime) != detected { return Err("Image content does not match its declared type".to_string()); } - Ok(base64_data.to_string()) + Ok(()) }) .collect::, String>>()?; let uploads = uploads @@ -85,19 +82,111 @@ impl AgentAttachmentStore { .store .store_image_uploads(self.account, session_id, &uploads) .map_err(|error| error.to_string())?; - Ok(stored - .into_iter() - .zip(encoded) - .map(|(attachment, base64_data)| PreparedAgentImage { - attachment: AgentImageAttachment { - id: attachment.id, - name: attachment.name, - mime_type: attachment.mime_type, - source: attachment.source, - }, - base64_data, - }) - .collect()) + Ok(stored.into_iter().map(AgentAttachment::from).collect()) + } + + /// Validate the entire message before publishing it. Pinned references + /// cannot expire while a message waits in the queue or is being edited. + pub(super) fn resolve_and_pin( + &self, + session_id: &str, + ids: &[String], + ) -> Result, String> { + let stored = self + .store + .resolve_attachments(self.account, session_id, ids) + .map_err(|error| error.to_string())?; + self.store + .pin_attachments(self.account, session_id, ids) + .map_err(|error| error.to_string())?; + Ok(stored.into_iter().map(AgentAttachment::from).collect()) + } + + pub(super) fn prepare_model_context( + &self, + session_id: &str, + messages: &mut [pi_ai::Message], + ) -> Result<(), String> { + self.refresh_metadata(session_id, messages)?; + self.store + .materialize_message_images(self.account, session_id, messages) + .map_err(|error| error.to_string()) + } + + /// A newly accepted image must reach Pi intact. Unlike retained history, + /// unavailable outgoing pixels are a send error, so the caller can keep the + /// user's message and attachments for correction or retry. + pub(super) fn prepare_outgoing_context( + &self, + session_id: &str, + messages: &mut [pi_ai::Message], + ) -> Result<(), String> { + self.refresh_metadata(session_id, messages)?; + self.store + .materialize_outgoing_images(self.account, session_id, messages) + .map_err(|error| error.to_string()) + } + + /// Pi's stranded steering queue contains text, not our typed Outgoing. + /// Recover attachment descriptors only when they still belong to this + /// task. A pasted listing must not gain image authority during requeue. + pub(super) fn split_owned_prompt( + &self, + session_id: &str, + prompt: &str, + ) -> Option<(String, Vec)> { + let (text, mut attachments) = split_attachment_prompt(prompt)?; + for attachment in &mut attachments { + *attachment = self + .store + .attachment_metadata(self.account, session_id, &attachment.id) + .ok()? + .into(); + } + Some((text, attachments)) + } + + fn refresh_metadata( + &self, + session_id: &str, + messages: &mut [pi_ai::Message], + ) -> Result<(), String> { + // Recovery can move the file-store generation. Resolve execution paths + // from scoped IDs each request instead of trusting persisted locations. + for message in messages.iter_mut() { + let pi_ai::Message::User(user) = message else { + continue; + }; + 'blocks: for block in &mut user.content { + let pi_ai::Content::Text(text) = block else { + continue; + }; + if !has_attachment_metadata(&text.text) { + continue; + } + let Some((typed, mut attachments, suffix)) = attachment_prompt_parts(&text.text) + else { + continue; + }; + for attachment in &mut attachments { + match self + .store + .attachment_metadata(self.account, session_id, &attachment.id) + { + Ok(stored) => *attachment = stored.into(), + // A listing can be pasted prose, including an example + // or a quote from another task. Its syntax alone grants + // no attachment authority and must not break inference. + Err(StoreError::NotFound | StoreError::InvalidInput) => continue 'blocks, + Err(error) => return Err(error.to_string()), + } + } + // Keep Pi's normalization notes (including coordinate scaling) + // and the original read_image guidance exactly as selected. + text.text = format!("{}{}", attachment_prompt(typed, &attachments, true), suffix); + } + } + Ok(()) } pub(super) fn read(&self, session_id: &str, attachment_id: &str) -> Result, String> { @@ -115,8 +204,8 @@ impl AgentAttachmentStore { crate::store::MAX_ATTACHMENT_RANGE_BYTES, ) .map_err(|error| error.to_string())?; - if range.size > MAX_AGENT_IMAGE_BYTES as u64 { - return Err("Agent image attachment exceeds the 10MB limit".to_string()); + if range.size > MAX_STORED_IMAGE_BYTES as u64 { + return Err("Agent image attachment exceeds the 20MB limit".to_string()); } if bytes.len().saturating_add(range.data.len()) as u64 > range.size { return Err("Agent image attachment is incomplete".to_string()); @@ -166,17 +255,43 @@ pub(super) fn attachment_id_from_source(source: &str) -> Option<&str> { .and_then(|id| (!id.is_empty() && !id.contains(['/', '\\', '?', '#'])).then_some(id)) } -const ATTACHED_IMAGE: &str = "The user attached the following image"; -const READ_IMAGE_HINT: &str = - "Use read_image with an attachment source when you need visual details from that image."; +impl From for AgentAttachment { + fn from(value: crate::store::StoredAttachment) -> Self { + Self { + id: value.id, + name: value.name, + mime_type: value.mime_type, + source: value.source, + size: value.size, + execution_path: value.execution_path, + } + } +} + +impl AgentAttachment { + pub fn is_image(&self) -> bool { + self.mime_type.starts_with("image/") + } +} + +const ATTACHED_FILES: &str = "The user attached these files on the execution computer:"; +const END_ATTACHMENTS: &str = "End of attachment metadata."; +const READ_IMAGE_HINT: &str = "Use read_image with an attachment source or executionPath when you need visual details from an image."; + +pub(super) fn has_attachment_metadata(text: &str) -> bool { + text.contains(ATTACHED_FILES) +} + +pub(super) fn prompt_has_inline_images(text: &str) -> bool { + has_attachment_metadata(text) && !text.lines().any(|line| line == READ_IMAGE_HINT) +} -/// The text a model gets for a message with images: what the user typed, -/// then each image's name and source, as Goose's runtime wrote it. A model -/// that gets the images beside the text (`inline`) is not told to look with -/// `read_image`; the source still lets it crop one. -pub(super) fn image_prompt( +/// Names and paths are JSON-quoted metadata, never instructions. Files are +/// stable copies on the execution computer; archives are not extracted. +/// The same metadata round-trips through queues, steering and the timeline. +pub(super) fn attachment_prompt( text: &str, - attachments: &[AgentImageAttachment], + attachments: &[AgentAttachment], inline: bool, ) -> String { let mut prompt = text.trim().to_string(); @@ -186,27 +301,68 @@ pub(super) fn image_prompt( if !prompt.is_empty() { prompt.push_str("\n\n"); } - prompt.push_str(ATTACHED_IMAGE); - prompt.push_str(if attachments.len() == 1 { - ":\n" - } else { - "s:\n" - }); + prompt.push_str(ATTACHED_FILES); + prompt.push('\n'); for attachment in attachments { - let name = serde_json::to_string(&attachment.name).unwrap_or_else(|_| "\"image\"".into()); - prompt.push_str(&format!("- {name}: {}\n", attachment.source)); + prompt.push_str("- "); + // A strongly typed descriptor only contains strings and an integer. + prompt + .push_str(&serde_json::to_string(attachment).expect("attachment metadata serializes")); + prompt.push('\n'); } - if !inline { + prompt.push_str(END_ATTACHMENTS); + if !inline && attachments.iter().any(AgentAttachment::is_image) { + prompt.push('\n'); prompt.push_str(READ_IMAGE_HINT); } prompt } -/// What the user typed and which images they attached, read back from the -/// text [`image_prompt`] wrote. `None` for a message without images. An -/// attachment read back has no media type. -pub(super) fn split_image_prompt(prompt: &str) -> Option<(String, Vec)> { - let start = prompt.rfind(ATTACHED_IMAGE)?; +pub(super) fn split_attachment_prompt(prompt: &str) -> Option<(String, Vec)> { + if has_attachment_metadata(prompt) { + return attachment_prompt_parts(prompt).map(|(typed, items, _)| (typed.to_owned(), items)); + } + // Historical prompts retain their labels after the byte migration. This + // only reads old text; it never reintroduces the retired image backend. + split_retired_image_prompt(prompt) +} + +fn attachment_prompt_parts(prompt: &str) -> Option<(&str, Vec, &str)> { + let start = if prompt.starts_with(ATTACHED_FILES) { + 0 + } else { + prompt.rfind(&format!("\n\n{ATTACHED_FILES}"))? + 2 + }; + let (typed, listing) = prompt.split_at(start); + let typed = match typed.strip_suffix("\n\n") { + Some(typed) => typed, + None if typed.is_empty() => "", + None => return None, + }; + let (listing, suffix) = listing.split_once(&format!("\n{END_ATTACHMENTS}"))?; + let mut lines = listing.lines(); + if lines.next()? != ATTACHED_FILES { + return None; + } + let mut attachments = Vec::new(); + for line in lines { + let attachment: AgentAttachment = serde_json::from_str(line.strip_prefix("- ")?).ok()?; + if attachment_id_from_source(&attachment.source)? != attachment.id { + return None; + } + attachments.push(attachment); + if attachments.len() > MAX_AGENT_IMAGES_PER_MESSAGE { + return None; + } + } + (!attachments.is_empty()).then_some((typed, attachments, suffix)) +} + +fn split_retired_image_prompt(prompt: &str) -> Option<(String, Vec)> { + const HEAD: &str = "The user attached the following image"; + const HINT: &str = + "Use read_image with an attachment source when you need visual details from that image."; + let start = prompt.rfind(HEAD)?; let (typed, listing) = prompt.split_at(start); let typed = match typed.strip_suffix("\n\n") { Some(typed) => typed, @@ -214,13 +370,13 @@ pub(super) fn split_image_prompt(prompt: &str) -> Option<(String, Vec return None, }; let mut lines = listing.lines(); - let head = lines.next()?.strip_prefix(ATTACHED_IMAGE)?; + let head = lines.next()?.strip_prefix(HEAD)?; if head != ":" && head != "s:" { return None; } let mut attachments = Vec::new(); for line in lines { - if line == READ_IMAGE_HINT { + if line == HINT { break; } let item = line.strip_prefix("- ")?; @@ -228,11 +384,13 @@ pub(super) fn split_image_prompt(prompt: &str) -> Option<(String, Vec user.content.iter().find_map(|block| match block { + pi_ai::Content::Image(image) => Some(image.data.clone()), + _ => None, + }), + _ => None, + }) + .unwrap(); + let id = + attachment_id_from_source(&source).expect("generated image stored as a file reference"); + assert_eq!(store.read(TASK_A, id).unwrap(), pixels); + } + #[test] fn attachment_adapter_captures_account_and_never_falls_back_to_files() { let temp = tempdir().unwrap(); @@ -404,8 +615,8 @@ mod tests { ) .unwrap(); let other = AgentAttachmentStore::new(Arc::clone(&store.store), other); - assert!(other.read(TASK_A, &prepared[0].attachment.id).is_err()); - assert!(other.read(TASK_C, &prepared[0].attachment.id).is_err()); + assert!(other.read(TASK_A, &prepared[0].id).is_err()); + assert!(other.read(TASK_C, &prepared[0].id).is_err()); // A previous-format fixture may remain on disk. An unknown SQL id // cannot select it, and the adapter never creates or removes its files. @@ -432,7 +643,7 @@ mod tests { }], ) .unwrap(); - let id = prepared[0].attachment.id.clone(); + let id = prepared[0].id.clone(); drop(store); let reopened = Arc::new(Store::open(crate::store::StorePaths::new(temp.path())).unwrap()); let account = reopened.account_id("attachment-fixture").unwrap().unwrap(); @@ -443,51 +654,156 @@ mod tests { ); } - fn attachment(id: char, name: &str) -> AgentImageAttachment { + fn attachment(id: char, name: &str, mime: &str) -> AgentAttachment { let id = id.to_string().repeat(32); - AgentImageAttachment { + AgentAttachment { source: format!("maple-attachment://{id}"), id, - name: name.to_string(), - mime_type: String::new(), + name: name.into(), + mime_type: mime.into(), + size: 42, + execution_path: format!("/uploads/session/{name}"), } } #[test] - fn an_image_prompt_names_the_images_and_reads_back() { - let one = [attachment('a', "dialog.png")]; - let prompt = image_prompt(" Why is this misaligned? ", &one, false); + fn mixed_attachment_metadata_round_trips_and_quotes_paths() { + let items = [ + attachment('a', "screen.png", "image/png"), + attachment('b', "report \"2026\"\n.zip", "application/zip"), + ]; + let prompt = attachment_prompt(" inspect these ", &items, false); + assert!(prompt.contains("executionPath")); + assert!(prompt.contains(READ_IMAGE_HINT)); assert_eq!( - prompt, - format!( - "Why is this misaligned?\n\nThe user attached the following image:\n- \"dialog.png\": maple-attachment://{}\n{READ_IMAGE_HINT}", - "a".repeat(32) - ) + split_attachment_prompt(&prompt), + Some(("inspect these".into(), items.to_vec())) ); + let inline = attachment_prompt("", &items, true); + assert!(!inline.contains(READ_IMAGE_HINT)); assert_eq!( - split_image_prompt(&prompt), - Some(("Why is this misaligned?".to_string(), one.to_vec())) + split_attachment_prompt(&inline), + Some((String::new(), items.to_vec())) ); + let files = attachment_prompt("", &items[1..], false); + assert!(!files.contains("read_image")); + assert_eq!(attachment_prompt(" plain ", &[], false), "plain"); + assert!(split_attachment_prompt("plain").is_none()); + } - // With the images beside the text, nothing asks for read_image; a - // name may hold anything, and the text may be empty. - let two = [attachment('b', "a \"b\"\nc.png"), attachment('c', "d.png")]; - let prompt = image_prompt("", &two, true); - assert!(prompt.starts_with("The user attached the following images:\n")); - assert!(!prompt.contains("read_image")); - assert_eq!( - split_image_prompt(&prompt), - Some((String::new(), two.to_vec())) + #[test] + fn selected_metadata_refreshes_execution_paths_without_reading_pixels() { + let temp = tempdir().unwrap(); + let store = fixture_attachment_store(temp.path(), &[TASK_A]); + let stored = store + .store_uploads( + TASK_A, + &[AgentImageUpload { + name: "screen.png".into(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .unwrap() + .remove(0); + let mut old = stored.clone(); + old.execution_path = "/an/old/store-generation/screen.png".into(); + let original = format!( + "{}\n[Image coordinates scaled by Pi.]", + attachment_prompt("look", &[old], false) ); + let mut messages = [pi_ai::Message::User(pi_ai::UserMessage { + content: vec![pi_ai::Content::text(original.clone())], + timestamp: 1, + })]; + std::fs::remove_file(&stored.execution_path).unwrap(); + store.prepare_model_context(TASK_A, &mut messages).unwrap(); + let pi_ai::Message::User(user) = &messages[0] else { + panic!("user") + }; + let (text, metadata) = + split_attachment_prompt(&pi_ai::content_text(&user.content)).unwrap(); + assert_eq!(text, "look"); + assert_eq!(metadata, [stored]); + assert!(pi_ai::content_text(&user.content).ends_with("[Image coordinates scaled by Pi.]")); + assert!(pi_ai::content_text(&user.content).contains(READ_IMAGE_HINT)); + assert!(original.contains("/an/old/store-generation/")); + } - // Text without images is left as it is. - assert_eq!(image_prompt(" plain ", &[], false), "plain"); - for text in [ - "plain", - "The user attached the following image: nothing", - "x\nThe user attached the following image:\n- \"a\": maple-attachment://1", - ] { - assert_eq!(split_image_prompt(text), None, "{text}"); + #[test] + fn unresolved_pasted_listings_stay_text_without_attachment_authority() { + let temp = tempdir().unwrap(); + let store = fixture_attachment_store(temp.path(), &[TASK_A, TASK_B]); + let foreign = store + .store_uploads( + TASK_B, + &[AgentImageUpload { + name: "other.png".into(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .unwrap() + .remove(0); + let missing = attachment('a', "example.zip", "application/zip"); + let mut invalid = missing.clone(); + invalid.id = "example-id".into(); + invalid.source = "maple-attachment://example-id".into(); + for listed in [foreign, missing, invalid] { + let text = attachment_prompt("Earlier notes quoted this listing:", &[listed], false); + let mut messages = [pi_ai::Message::User(pi_ai::UserMessage::text(&text))]; + let original = messages.clone(); + store.prepare_model_context(TASK_A, &mut messages).unwrap(); + assert_eq!(messages, original, "unresolved prose must remain unchanged"); + assert!(store.split_owned_prompt(TASK_A, &text).is_none()); } } + + #[test] + fn outgoing_images_fail_preparation_while_missing_history_remains_usable() { + let temp = tempdir().unwrap(); + let store = fixture_attachment_store(temp.path(), &[TASK_A]); + let image = store + .store_uploads( + TASK_A, + &[AgentImageUpload { + name: "screen.png".into(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .unwrap() + .remove(0); + let original = [pi_ai::Message::User(pi_ai::UserMessage { + content: vec![pi_ai::Content::image(&image.source, "image/png")], + timestamp: 1, + })]; + std::fs::remove_file(image.execution_path).unwrap(); + let mut outgoing = original.clone(); + assert!( + store + .prepare_outgoing_context(TASK_A, &mut outgoing) + .is_err() + ); + assert_eq!(outgoing, original); + let mut history = original; + store.prepare_model_context(TASK_A, &mut history).unwrap(); + let pi_ai::Message::User(user) = &history[0] else { + panic!("user") + }; + assert_eq!( + user.content, + [pi_ai::Content::text("Image no longer available.")] + ); + } + + #[test] + fn retired_images_keep_labels_but_no_payload_or_path() { + let prompt = format!( + "hello\n\nThe user attached the following image:\n- \"screen.png\": maple-attachment://{}\n", + "a".repeat(32) + ); + let (text, items) = split_attachment_prompt(&prompt).unwrap(); + assert_eq!(text, "hello"); + assert_eq!(items[0].name, "screen.png"); + assert!(items[0].is_image()); + assert!(items[0].execution_path.is_empty()); + } } diff --git a/apps/maple-agent/crates/maple-agent/src/agent/image_persistence_tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/image_persistence_tests.rs index 20d0f27b9..ca9dd4a64 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/image_persistence_tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/image_persistence_tests.rs @@ -59,17 +59,42 @@ async fn mcp_images_remain_saveable_and_do_not_block_lifecycle_operations() { let store = harness.service.host.store.as_ref().unwrap(); assert!(store.persistence_failure().is_none()); harness.service.flush_persistence().await.unwrap(); - assert_eq!( - store - .load_session( - runtime.store.account(), - &task, - crate::store::SessionLoad::Full - ) - .unwrap() - .entries, - retained - ); + let persisted = store + .load_session( + runtime.store.account(), + &task, + crate::store::SessionLoad::Full, + ) + .unwrap() + .entries; + assert_eq!(persisted.len(), retained.len()); + let images: Vec<_> = persisted.iter().filter_map(|entry| match &entry.kind { + pi_coding_agent::session::EntryKind::Message { + message: pi_coding_agent::SessionMessage::Llm(pi_ai::Message::ToolResult(result)) + } => Some(result), + _ => None, + }).flat_map(|result| { + assert!(!result.is_error, "unavailable media does not turn a successful tool into failure"); + result.content.iter().filter_map(|block| match block { + pi_ai::Content::Image(image) => Some(image.data.clone()), _ => None, + }) + }).collect(); + assert_eq!(images.len(), 1); + if data == "not base64!" { + assert_eq!(images[0], "maple-attachment-unavailable:retired"); + } else { + assert!(images[0].starts_with("maple-attachment://")); + } + // The argument that happens to contain base64 is ordinary text; + // only the typed image result is externalized. + assert!(persisted.iter().any(|entry| match &entry.kind { + pi_coding_agent::session::EntryKind::Message { + message: + pi_coding_agent::SessionMessage::Llm(pi_ai::Message::Assistant(message)), + } => message.content.iter().any(|block| matches!(block, + pi_ai::AssistantContent::ToolCall(call) if call.arguments["text"] == data)), + _ => false, + })); // Exercise each lifecycle barrier while the image-bearing context // is still loaded; deleting it first would trivialize the others. match finish { diff --git a/apps/maple-agent/crates/maple-agent/src/agent/maintenance.rs b/apps/maple-agent/crates/maple-agent/src/agent/maintenance.rs index a27516b43..8b71b9bc1 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/maintenance.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/maintenance.rs @@ -135,7 +135,7 @@ pub(super) fn start(service: &MapleAgentService) { fn admitted(state: &ServiceState) -> bool { state.admission.load(Ordering::Acquire) == AGENT_SERVICE_OPEN // Non-runnable does not mean unloaded: pending shutdown still owns - // source entries, blob leases and account/session identity. All idle + // source entries, attachment leases and account/session identity. All idle // effects yield until Retry or another completed shutdown saves them. && state .shutdown_pending @@ -155,7 +155,6 @@ struct Worker { account_cursor: Option, old_epoch_cursors: HashMap, rebuild: Option, - orphan_pass_done: bool, last_warning: Option, integrity: Option, last_integrity_attempt: Option, @@ -171,7 +170,7 @@ impl Worker { ..Self::default() }; } - /// At most four deleted-task steps, two account steps, one orphan-blob + /// At most four deleted-task steps, two account steps, one unbound-attachment /// step, one rebuild commit and one old-epoch step run in a tick. Every /// payload mutation uses the store's bounded idle-writer transactions. fn tick( @@ -192,11 +191,12 @@ impl Worker { if !admitted(state) { return; } - if !self.orphan_pass_done { - match while_service_open(state, || store.collect_orphan_blobs_step()) { - Ok(progress) => self.orphan_pass_done = progress.done, - Err(error) => self.report("orphan blob collection", error), - } + // A later cancellation or a newly expired staging upload must be + // discovered even after an earlier pass found no work. Each tick only + // queries one indexed candidate and removes at most one owned file. + match while_service_open(state, || store.collect_unbound_attachments_step()) { + Ok(_) => {} + Err(error) => self.report("unbound attachment collection", error), } if !admitted(state) { return; @@ -368,7 +368,7 @@ impl Worker { } fn collect_deleted_tasks(&mut self, state: &ServiceState, store: &Arc) { - let tasks = match store.deleted_tasks(TASK_GC_STEPS) { + let tasks = match store.deleted_tasks_for_collection(TASK_GC_STEPS) { Ok(tasks) => tasks, Err(error) => { self.report("deleted-task discovery", error); @@ -719,6 +719,42 @@ mod tests { ); } + #[test] + fn later_cancelled_attachments_are_collected_after_an_empty_pass() { + use sha2::{Digest, Sha256}; + let (_directory, store, service, _events, account) = fixture(); + create_task(&store, account); + let mut worker = Worker::default(); + worker.tick(&service.state, &store, None, &service.host.events); + let bytes = b"later cancelled upload"; + let metadata = crate::store::AttachmentUploadMetadata { + name: "notes.txt".into(), + mime_type: "text/plain".into(), + size: bytes.len() as u64, + sha256: format!("{:x}", Sha256::digest(bytes)), + }; + let crate::store::AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(account, TASK, "later-cancel", &metadata) + .unwrap() + else { + panic!("new upload") + }; + writer.write_chunk(bytes).unwrap(); + let attachment = writer.finish(&store).unwrap(); + assert!(std::path::Path::new(&attachment.execution_path).exists()); + store + .cancel_attachment_upload(account, TASK, "later-cancel") + .unwrap(); + worker.tick(&service.state, &store, None, &service.host.events); + assert!(!std::path::Path::new(&attachment.execution_path).exists()); + assert_eq!( + store + .attachment_upload_status(account, TASK, "later-cancel") + .unwrap(), + crate::store::AttachmentUploadStatus::Cancelled + ); + } + #[test] fn repair_emits_replaced_after_flip_then_collects_only_the_obsolete_epoch() { let (_directory, store, service, events, account) = fixture(); diff --git a/apps/maple-agent/crates/maple-agent/src/agent/provider.rs b/apps/maple-agent/crates/maple-agent/src/agent/provider.rs index ee8b2732b..23085e74f 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/provider.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/provider.rs @@ -343,6 +343,93 @@ pub(crate) fn maple_model( } } +/// A task-bound provider view. Pi selects effective context before this +/// boundary; Maple then applies its screenshot limit before reading any files. +/// The shared provider registry remains live, including scripted test providers. +pub(super) fn attachment_model_registry( + models: &ModelRegistry, + attachments: super::attachments::AgentAttachmentStore, + session_id: String, +) -> ModelRegistry { + let registry = ModelRegistry::new(Arc::new(NoKeys)); + registry.register_models(models.models()); + registry.register_api( + MAPLE_API, + Arc::new(AttachmentStreamFn { + inner: models.stream_fn(), + attachments, + session_id, + }), + ); + registry +} + +struct AttachmentStreamFn { + inner: Arc, + attachments: super::attachments::AgentAttachmentStore, + session_id: String, +} + +impl StreamFn for AttachmentStreamFn { + fn stream( + &self, + model: &Model, + mut context: Context, + options: StreamOptions, + ) -> AssistantMessageStream { + bound_tool_result_images(&mut context.messages); + let has_attachments = context.messages.iter().any(|message| { + let content = match message { + Message::User(user) => &user.content, + Message::ToolResult(result) => &result.content, + _ => return false, + }; + content.iter().any(|block| match block { + Content::Image(image) => image.data.starts_with("maple-attachment"), + Content::Text(text) => super::attachments::has_attachment_metadata(&text.text), + }) + }); + if !has_attachments { + return self.inner.stream(model, context, options); + } + let (sender, stream) = AssistantMessageStream::channel(); + let model = model.clone(); + let inner = Arc::clone(&self.inner); + let attachments = self.attachments.clone(); + let session_id = self.session_id.clone(); + tokio::spawn(async move { + let prepared = tokio::task::spawn_blocking(move || { + attachments.prepare_model_context(&session_id, &mut context.messages)?; + Ok::<_, String>(context) + }); + let prepared = tokio::select! { + _ = options.cancel.cancelled() => { + pi_ai::AssistantMessageBuilder::new(sender, &model).fail(pi_ai::StopReason::Aborted, "Request cancelled"); + return; + } + result = prepared => result, + }; + let context = match prepared { + Ok(Ok(context)) => context, + _ => { + pi_ai::AssistantMessageBuilder::new(sender, &model).fail( + pi_ai::StopReason::Error, + "Maple could not read an image attachment", + ); + return; + } + }; + let mut events = inner.stream(&model, context, options); + while let Some(event) = events.next().await { + if sender.send(event).is_err() { + break; + } + } + }); + stream + } +} + /// Pi's Chat Completions provider over the account's transport, plus the /// request and response adjustments Maple's models need. pub(crate) struct MapleStreamFn { diff --git a/apps/maple-agent/crates/maple-agent/src/agent/provider/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/provider/tests.rs index a3dc8ef9a..2b590639e 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/provider/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/provider/tests.rs @@ -687,6 +687,98 @@ fn only_the_newest_tool_result_images_stay_in_a_request() { assert_eq!(oldest.content[1], Content::text(SUPERSEDED_IMAGE_MARKER)); } +#[tokio::test] +async fn selected_images_materialize_after_screenshot_filtering_and_keep_source_references() { + use crate::agent::attachments::{AgentImageUpload, fixture_attachment_store}; + const TASK: &str = "018ddbcc-b80b-7c20-8123-000000000001"; + const OTHER: &str = "018ddbcc-b80b-7c20-8123-000000000002"; + const PNG: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="; + let root = tempfile::tempdir().unwrap(); + let attachments = fixture_attachment_store(root.path(), &[TASK, OTHER]); + let upload = AgentImageUpload { + name: "shot.png".into(), + data_url: format!("data:image/png;base64,{PNG}"), + }; + let foreign = attachments + .store_uploads(OTHER, std::slice::from_ref(&upload)) + .unwrap() + .remove(0); + let retained = attachments + .store_uploads(TASK, &[upload]) + .unwrap() + .remove(0); + let faux = pi_ai::faux::FauxProvider::new(); + faux.push_text("seen"); + let stream = AttachmentStreamFn { + inner: Arc::new(faux.clone()), + attachments, + session_id: TASK.into(), + }; + // Reading the oldest reference would fail task ownership. It must be + // removed before any materialization, while the retained image is real. + let original = vec![ + screenshot_result("old", &foreign.source), + screenshot_result("one", &retained.source), + screenshot_result("two", &retained.source), + screenshot_result("three", &retained.source), + ]; + let result = stream + .stream( + &model("glm-5-3"), + Context::from_messages(original.clone()), + StreamOptions::default(), + ) + .result() + .await; + assert_eq!(result.stop_reason, StopReason::Stop); + let request = faux.requests().pop().unwrap(); + let Message::ToolResult(old) = &request.context.messages[0] else { + panic!("tool result") + }; + assert_eq!(old.content[1], Content::text(SUPERSEDED_IMAGE_MARKER)); + assert!(!old.is_error); + for message in &request.context.messages[1..] { + let Message::ToolResult(tool) = message else { + panic!("tool result") + }; + assert!(matches!(&tool.content[1], Content::Image(image) if image.data == PNG)); + } + assert!( + matches!(&original[1], Message::ToolResult(tool) if matches!(&tool.content[1], Content::Image(image) if image.data == retained.source)) + ); + + // Missing retained pixels are harmless text; tool identity and success + // survive. A wrong-task retained reference still fails before inference. + std::fs::remove_file(&retained.execution_path).unwrap(); + faux.push_text("unavailable"); + let result = stream + .stream( + &model("glm-5-3"), + Context::from_messages(vec![screenshot_result("missing", &retained.source)]), + StreamOptions::default(), + ) + .result() + .await; + assert_eq!(result.stop_reason, StopReason::Stop); + let request = faux.requests().pop().unwrap(); + let Message::ToolResult(tool) = &request.context.messages[0] else { + panic!("tool result") + }; + assert_eq!(tool.tool_call_id, "missing"); + assert!(!tool.is_error); + assert_eq!(tool.content[1], Content::text("Image no longer available.")); + let result = stream + .stream( + &model("glm-5-3"), + Context::from_messages(vec![screenshot_result("foreign", &foreign.source)]), + StreamOptions::default(), + ) + .result() + .await; + assert_eq!(result.stop_reason, StopReason::Error); + assert_eq!(faux.requests().len(), 2); +} + const KIMI_TOOL_STREAM: [&str; 4] = [ "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"shell:0\",\"function\":{\"name\":\"shell\",\"arguments\":\"{}\"}}]}}]}\n\n", "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"chatcmpl-tool-abc\",\"function\":{\"name\":\"read\",\"arguments\":\"{}\"}}]}}]}\n\n", diff --git a/apps/maple-agent/crates/maple-agent/src/agent/runs.rs b/apps/maple-agent/crates/maple-agent/src/agent/runs.rs index 1c2ec2535..1f3a65f27 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/runs.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/runs.rs @@ -24,7 +24,7 @@ use tokio::sync::{Notify, mpsc, watch}; use tokio_util::sync::CancellationToken; use super::attachments::{ - AgentAttachmentStore, AgentImageAttachment, image_prompt, split_image_prompt, + AgentAttachment, AgentAttachmentStore, attachment_prompt, prompt_has_inline_images, }; use super::config::unix_ms; use super::runtime::AgentRuntime; @@ -47,8 +47,9 @@ const AGENT_RUN_EVENT_CAPACITY: usize = 256; const MAX_DESKTOP_QUEUE_ITEMS: usize = 16; const MAX_DESKTOP_QUEUE_TEXT_BYTES: usize = 32 * 1024; const QUEUED_MESSAGE_ATTACHMENTS_ERROR: &str = - "New images cannot be added while sending a queued message"; + "New attachments cannot be added while sending a queued message"; const EMPTY_PROMPT_ERROR: &str = "Prompt cannot be empty"; +const IMAGE_ATTACHMENT_PREPARATION_ERROR: &str = "Maple could not prepare an image attachment. Your message was not sent and is being kept for retry. Remove or replace the attachment, then send it again."; /// How often Stop repeats its abort until the run's prompt returns. const STOP_REPEAT_INTERVAL: Duration = Duration::from_millis(100); /// How long Stop and shutdown wait for a run to end before aborting it. @@ -84,7 +85,7 @@ fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { #[derive(Clone, Debug, Default)] struct Outgoing { text: String, - attachments: Vec, + attachments: Vec, images: Vec, /// Set for a message to the model alone, such as a background agent's /// result: Pi gets it as a custom message of this type, which the @@ -111,10 +112,18 @@ impl Outgoing { /// A message read back from the text Pi was given, as when steering /// goes back to the queue. Its images are then referred to by source. - fn from_prompt(prompt: String) -> Self { - match split_image_prompt(&prompt) { + fn from_prompt(prompt: String, store: &AgentAttachmentStore, session_id: &str) -> Self { + match store.split_owned_prompt(session_id, &prompt) { Some((text, attachments)) => Self { text, + images: attachments + .iter() + .filter(|attachment| attachment.is_image() && prompt_has_inline_images(&prompt)) + .map(|attachment| ImageContent { + data: attachment.source.clone(), + mime_type: attachment.mime_type.clone(), + }) + .collect(), attachments, ..Self::default() }, @@ -126,13 +135,45 @@ impl Outgoing { self.text.is_empty() && self.attachments.is_empty() } - /// The text and images Pi gets: the text names the attachments, and a - /// model without the images is told to look with `read_image`. - fn prompt(&self) -> (String, Vec) { - ( - image_prompt(&self.text, &self.attachments, !self.images.is_empty()), - self.images.clone(), - ) + /// Resolve only this message's selected images just before it enters Pi. + /// The queue keeps references, while Pi still receives ordinary pixels for + /// its normal image normalization and context-token accounting. + async fn prompt( + &self, + store: &AgentAttachmentStore, + session_id: &str, + ) -> Result<(String, Vec), AgentSessionError> { + let text = attachment_prompt(&self.text, &self.attachments, !self.images.is_empty()); + if self.images.is_empty() { + return Ok((text, Vec::new())); + } + let store = store.clone(); + let session_id = session_id.to_string(); + let images = self.images.clone(); + tokio::task::spawn_blocking(move || { + let mut messages = [Message::User(pi_ai::UserMessage { + content: std::iter::once(Content::text(text)) + .chain(images.into_iter().map(Content::Image)) + .collect(), + timestamp: 0, + })]; + store.prepare_outgoing_context(&session_id, &mut messages)?; + let Message::User(user) = &mut messages[0] else { + unreachable!() + }; + let mut text = Vec::new(); + let mut images = Vec::new(); + for block in std::mem::take(&mut user.content) { + match block { + Content::Text(block) => text.push(block.text), + Content::Image(image) => images.push(image), + } + } + Ok::<_, String>((text.join("\n\n"), images)) + }) + .await + .map_err(|error| attachment_preparation_error(&error.to_string()))? + .map_err(|error| attachment_preparation_error(&error)) } /// A hidden message as the custom message Pi gets; `None` for a @@ -148,12 +189,50 @@ impl Outgoing { /// Steer a message the user sent into the session's turn. A hidden /// message never is: it goes to the model as a turn of its own. - async fn steer(&self, session: &AgentSession) { + async fn steer( + &self, + session: &AgentSession, + store: &AgentAttachmentStore, + session_id: &str, + ) -> Result { debug_assert!(self.hidden.is_none(), "a hidden message is not steered"); - let (text, images) = self.prompt(); - if let Err(error) = session.steer(&text, images).await { - log::warn!("Failed to steer the Agent task: {error}"); + let (text, images) = self.prompt(store, session_id).await?; + session.steer(&text, images).await + } +} + +fn attachment_preparation_error(error: &str) -> AgentSessionError { + log::warn!("Failed to prepare Agent image attachments: {error}"); + AgentSessionError::Failed(IMAGE_ATTACHMENT_PREPARATION_ERROR.to_string()) +} + +/// After a steering failure, keep that message and every later one in order. +/// The prompt already in flight can finish without losing its accepted input. +#[derive(Default)] +struct DeferredSteering { + messages: Vec, + failure: Option, +} + +impl DeferredSteering { + async fn steer( + &mut self, + message: Outgoing, + session: &AgentSession, + store: &AgentAttachmentStore, + session_id: &str, + events: &RunEvents, + ) { + if self.failure.is_none() + && let Err(error) = message.steer(session, store, session_id).await + { + let item = error_item(error.to_string()); + events.publish(AgentRunEvent::Error(item.clone())); + self.failure = Some(item); + } else if self.failure.is_none() { + return; } + self.messages.push(message); } } @@ -801,6 +880,7 @@ impl AgentRuntimeHandle { ) -> Result { self.verify_generation().await?; self.ensure_accepting_new_work()?; + self.ensure_task_history_ready(&request.session_id).await?; let runtime = self.runtime().await?; runtime.send_desktop(request).await } @@ -930,6 +1010,11 @@ impl AgentRuntimeHandle { ) -> Result { self.verify_generation().await?; self.ensure_accepting_new_work()?; + tokio::select! { + biased; + _ = cancellation.cancelled() => return Err("The prompt was cancelled before admission".into()), + result = self.ensure_task_history_ready(&request.session_id) => result?, + } let runtime = self.runtime().await?; runtime.send_surface(access, request, cancellation).await } @@ -1053,8 +1138,9 @@ impl AgentRuntime { } } - /// The message a send carries. Its images are stored with the task - /// first; a model that sees images gets them beside the text. + /// Resolve committed IDs as one batch before accepting the message. The + /// transcript and queue carry file references; only selected model context + /// materializes image pixels at the provider boundary. async fn outgoing(&self, request: &AgentSendMessageRequest) -> Result { let text = request.text.trim().to_string(); if request.attachments.is_empty() { @@ -1065,17 +1151,18 @@ impl AgentRuntime { return Err(format!("Failed to find Agent task {session_id}")); } let store = AgentAttachmentStore::new(Arc::clone(self.store.store()), self.store.account()); - let uploads = request.attachments.clone(); - let stored = - tokio::task::spawn_blocking(move || store.store_uploads(&session_id, &uploads)) + let ids = request.attachments.clone(); + let attachments = + tokio::task::spawn_blocking(move || store.resolve_and_pin(&session_id, &ids)) .await - .map_err(|error| format!("Failed to store the images: {error}"))??; + .map_err(|error| format!("Failed to prepare attachments: {error}"))??; let images = if request.vision_capable { - stored + attachments .iter() - .map(|image| ImageContent { - data: image.base64_data.clone(), - mime_type: image.attachment.mime_type.clone(), + .filter(|attachment| attachment.is_image()) + .map(|attachment| ImageContent { + data: attachment.source.clone(), + mime_type: attachment.mime_type.clone(), }) .collect() } else { @@ -1083,7 +1170,7 @@ impl AgentRuntime { }; Ok(Outgoing { text, - attachments: stored.into_iter().map(|image| image.attachment).collect(), + attachments, images, hidden: None, }) @@ -1499,6 +1586,8 @@ impl AgentRuntime { usage, mut steers, } = run; + let attachments = + AgentAttachmentStore::new(Arc::clone(self.store.store()), self.store.account()); let watch = RunWatch::default(); let listener = session.subscribe({ let events = events.clone(); @@ -1521,6 +1610,8 @@ impl AgentRuntime { let mut failure = None; let mut prompted = false; + let mut unsubmitted = Vec::new(); + let mut deferred = DeferredSteering::default(); while !stopped.is_cancelled() { // Messages the user sent go together, as one turn; a hidden // message goes on its own. @@ -1534,13 +1625,41 @@ impl AgentRuntime { } let mut sending: Vec = prompts.drain(..batch).collect(); let first = sending.remove(0); + // Preparation can fail before Pi records the first message. Keep + // ownership of the entire batch until that preparation succeeds. + let prepared_first = if first.hidden.is_none() { + match first.prompt(&attachments, &session_id).await { + Ok(prepared) => Some(prepared), + Err(error) => { + failure = Some(error.to_string()); + unsubmitted.push(first); + unsubmitted.extend(sending); + break; + } + } + } else { + None + }; + if stopped.is_cancelled() { + unsubmitted.push(first); + unsubmitted.extend(sending); + break; + } // Steering queued before the prompt joins its first turn, so // several messages make one turn. while let Ok(message) = steers.try_recv() { - message.steer(&session).await; + deferred + .steer(message, &session, &attachments, &session_id, &events) + .await; + } + for message in sending { + deferred + .steer(message, &session, &attachments, &session_id, &events) + .await; } - for message in &sending { - message.steer(&session).await; + if stopped.is_cancelled() { + unsubmitted.push(first); + break; } prompted = true; let result = match first.custom_draft() { @@ -1548,12 +1667,20 @@ impl AgentRuntime { until_stopped( &session, &stopped, - with_steering(&session, &mut steers, session.send_custom_message(draft)), + with_steering( + &session, + &attachments, + &session_id, + &mut steers, + &mut deferred, + &events, + session.send_custom_message(draft), + ), ) .await } None => { - let (text, images) = first.prompt(); + let (text, images) = prepared_first.expect("a user message was prepared"); let options = PromptOptions { images, expand: false, @@ -1562,7 +1689,15 @@ impl AgentRuntime { until_stopped( &session, &stopped, - with_steering(&session, &mut steers, session.prompt(&text, options)), + with_steering( + &session, + &attachments, + &session_id, + &mut steers, + &mut deferred, + &events, + session.prompt(&text, options), + ), ) .await } @@ -1572,7 +1707,7 @@ impl AgentRuntime { failure = Some(error.to_string()); break; } - if stopped.is_cancelled() { + if stopped.is_cancelled() || deferred.failure.is_some() { break; } // A message steered in after Pi's last look would wait for the @@ -1582,11 +1717,23 @@ impl AgentRuntime { // everything accepted before that boundary before clearing Pi's queue. let chips = self.take_chips_or_close(&run_id, &session_id); while let Ok(message) = steers.try_recv() { - message.steer(&session).await; + deferred + .steer(message, &session, &attachments, &session_id, &events) + .await; + } + if deferred.failure.is_some() { + if let Some((chips, snapshot)) = chips { + events.publish(AgentRunEvent::QueueChanged(snapshot)); + prompts.extend(chips.into_iter().map(Chip::outgoing)); + } + break; } let (stranded, _) = session.clear_queue(); let waiting = std::mem::take(&mut prompts); - prompts = stranded.into_iter().map(Outgoing::from_prompt).collect(); + prompts = stranded + .into_iter() + .map(|prompt| Outgoing::from_prompt(prompt, &attachments, &session_id)) + .collect(); prompts.extend(waiting); match chips { Some((chips, snapshot)) => { @@ -1601,63 +1748,64 @@ impl AgentRuntime { if let Some(run) = self.runs.state().runs.get_mut(&run_id) { run.accepting = false; } - if failure.is_some() && !stopped.is_cancelled() { - let (stranded, _) = session.clear_queue(); - let mut unsent: Vec<_> = stranded.into_iter().map(Outgoing::from_prompt).collect(); - while let Ok(message) = steers.try_recv() { - unsent.push(message); - } - unsent.append(&mut prompts); - prompts = unsent; - } session.unsubscribe(listener); let cancelled = stopped.is_cancelled(); - if cancelled { - // Steering the session never took, and messages the run did not - // send yet, go back to the queue. + if cancelled || failure.is_some() || deferred.failure.is_some() { + // Only messages Pi has not accepted go back to the queue. In + // particular, a prompt error after recording its input must not + // duplicate that input on retry. let (stranded, _) = session.clear_queue(); - let mut unsent = if prompted { - Vec::new() - } else { - std::mem::take(&mut prompts) - }; - unsent.extend(stranded.into_iter().map(Outgoing::from_prompt)); + let mut unsent = unsubmitted; + if !prompted && unsent.is_empty() { + // No batch started. Otherwise its first message is already + // here, and its stranded steers precede the remaining batches. + unsent.append(&mut prompts); + } + unsent.extend( + stranded + .into_iter() + .map(|prompt| Outgoing::from_prompt(prompt, &attachments, &session_id)), + ); + unsent.append(&mut deferred.messages); while let Ok(message) = steers.try_recv() { unsent.push(message); } unsent.extend(std::mem::take(&mut prompts)); - // Hidden messages still waiting go into the history with the - // rest; the user's chips wait for their next send. - let waiting = self.runs.state().take_hidden_chips(&session_id); - unsent.extend(waiting.into_iter().map(Chip::outgoing)); - self.requeue(&session, &session_id, unsent, &events); - if prompted { - let id = format!("stopped-{}", pi_ai::now_ms()); - session.extension_context().append_entry( - MAPLE_NOTICE_ENTRY, - Some(notice_entry_data(&id, STOPPED_NOTICE_TEXT)), - ); - events.publish(AgentRunEvent::TimelineItem(notice_item( - id, - "Agent notice", - STOPPED_NOTICE_TEXT, - pi_ai::now_ms(), - ))); + if cancelled { + // Hidden messages still waiting go into the history with the + // rest; the user's chips wait for their next send. + let waiting = self.runs.state().take_hidden_chips(&session_id); + unsent.extend(waiting.into_iter().map(Chip::outgoing)); } + self.requeue(&session, &session_id, unsent, &events); } - if failure.is_some() && !prompts.is_empty() { - // A failed prompt leaves the messages after it for the next send. - self.requeue(&session, &session_id, std::mem::take(&mut prompts), &events); + if cancelled && prompted { + let id = format!("stopped-{}", pi_ai::now_ms()); + session.extension_context().append_entry( + MAPLE_NOTICE_ENTRY, + Some(notice_entry_data(&id, STOPPED_NOTICE_TEXT)), + ); + events.publish(AgentRunEvent::TimelineItem(notice_item( + id, + "Agent notice", + STOPPED_NOTICE_TEXT, + pi_ai::now_ms(), + ))); } - let terminal_state = if cancelled { - AgentRunTerminal::Cancelled + let failed = failure.is_some() || deferred.failure.is_some(); + if let Some(item) = deferred.failure { + // This error was already shown when steering failed, including + // when Stop raced with the failed attachment preparation. + self.remember_failure(&session_id, item); } else if let Some(error) = failure { let item = error_item(error); self.remember_failure(&session_id, item.clone()); events.publish(AgentRunEvent::Error(item)); - AgentRunTerminal::Failed - } else if last_reply_failed(&session) { + } + let terminal_state = if cancelled { + AgentRunTerminal::Cancelled + } else if failed || last_reply_failed(&session) { AgentRunTerminal::Failed } else { AgentRunTerminal::Completed @@ -1734,14 +1882,20 @@ impl AgentRuntime { /// drains the inbox, so async input handlers cannot reorder accepted steers. async fn with_steering( session: &AgentSession, + attachments: &AgentAttachmentStore, + session_id: &str, steers: &mut mpsc::UnboundedReceiver, + deferred: &mut DeferredSteering, + events: &RunEvents, prompt: impl std::future::Future>, ) -> Result { tokio::pin!(prompt); loop { tokio::select! { result = &mut prompt => return result, - Some(message) = steers.recv() => message.steer(session).await, + Some(message) = steers.recv() => { + deferred.steer(message, session, attachments, session_id, events).await; + } } } } @@ -1950,6 +2104,7 @@ impl AgentRuntime { #[cfg(test)] mod tests { use super::*; + use crate::agent::attachments::split_attachment_prompt; use crate::agent::tests::{Harness, eventually, finished}; #[test] @@ -2119,4 +2274,154 @@ mod tests { } } } + + #[tokio::test] + async fn failed_batch_image_preparation_preserves_order_and_retry_does_not_duplicate() { + const PNG: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="; + for failed_index in [0, 1] { + let harness = Harness::new().await; + let runtime = harness.runtime().await; + let task = harness.create_task().await; + let image = harness + .handle + .import_image_attachments( + task.clone(), + vec![super::super::attachments::AgentImageUpload { + name: "screen.png".into(), + data_url: format!("data:image/png;base64,{PNG}"), + }], + ) + .await + .unwrap() + .remove(0); + let request = AgentSendMessageRequest { + session_id: task.clone(), + text: String::new(), + model: None, + context_limit: Some(200_000), + vision_capable: true, + steer: false, + queue_id: None, + attachments: Vec::new(), + }; + let texts = ["First", "Second", "Third"]; + let mut prompts: Vec<_> = texts.iter().map(|text| Outgoing::text(*text)).collect(); + prompts[failed_index] = runtime + .outgoing(&AgentSendMessageRequest { + text: texts[failed_index].into(), + attachments: vec![image.id.clone()], + ..request.clone() + }) + .await + .unwrap(); + let original = std::fs::read(&image.execution_path).unwrap(); + let mut changed = original.clone(); + *changed.last_mut().unwrap() ^= 1; + std::fs::write(&image.execution_path, changed).unwrap(); + let run = runtime.register_run(&mut runtime.runs.state(), &task, prompts, Vec::new()); + if failed_index > 0 { + harness.faux.push_text("First accepted."); + } + let mut handle = runtime.start_run(&request, run).await.unwrap(); + assert_eq!(finished(&mut handle).await, AgentRunTerminal::Failed); + let queue = runtime.runs.queue_snapshot(&task); + let queued: Vec<_> = queue.items.iter().map(|item| item.text.as_str()).collect(); + assert_eq!(queued, texts[failed_index..]); + let session = runtime.loaded_session(&task).await.unwrap(); + let user_texts = || { + pi_coding_agent::messages::convert_to_llm(&session.messages()) + .into_iter() + .filter_map(|message| match message { + Message::User(user) => { + let text = pi_ai::content_text(&user.content); + Some(match split_attachment_prompt(&text) { + Some((typed, _)) => typed, + None => text, + }) + } + _ => None, + }) + .collect::>() + }; + assert_eq!(user_texts(), texts[..failed_index]); + let errors = harness + .recorder + .events() + .into_iter() + .filter(|event| { + matches!(event, AgentServiceEvent::Run { + event: AgentRunEvent::Error(item), .. + } if item.text.as_deref() == Some(IMAGE_ATTACHMENT_PREPARATION_ERROR)) + }) + .count(); + assert_eq!(errors, 1); + + std::fs::write(&image.execution_path, original).unwrap(); + harness.faux.push_text("Retry accepted."); + let mut retry = runtime.send_desktop(request).await.unwrap(); + assert_eq!(finished(&mut retry).await, AgentRunTerminal::Completed); + assert!(runtime.runs.queue_snapshot(&task).items.is_empty()); + assert_eq!(user_texts(), texts); + } + } + + #[tokio::test] + async fn stop_after_batch_steering_preserves_order_before_later_batches() { + let harness = Harness::new().await; + let runtime = harness.runtime().await; + let task = harness.create_task().await; + let request = AgentSendMessageRequest { + session_id: task.clone(), + text: String::new(), + model: None, + context_limit: Some(200_000), + vision_capable: false, + steer: false, + queue_id: None, + attachments: Vec::new(), + }; + let run = runtime.register_run( + &mut runtime.runs.state(), + &task, + vec![ + Outgoing::text("First"), + Outgoing::text("Second"), + Outgoing::text("Third"), + Outgoing::hidden("maple.test", "Hidden".into()), + Outgoing::text("Later"), + ], + Vec::new(), + ); + let session = runtime + .prepare_run_session(&request, &run.prompts, &run.events) + .await + .unwrap(); + let cancelled_at_boundary = Arc::new(AtomicBool::new(false)); + let listener = session.subscribe({ + let stopped = run.stopped.clone(); + let cancelled_at_boundary = cancelled_at_boundary.clone(); + move |event| { + if matches!(event, AgentSessionEvent::QueueUpdate { steering, .. } + if steering.iter().any(|text| text == "Third")) + { + cancelled_at_boundary.store(true, Ordering::SeqCst); + stopped.cancel(); + } + } + }); + let mut handle = runtime.launch_run(&task, run, session.clone()).unwrap(); + assert_eq!(finished(&mut handle).await, AgentRunTerminal::Cancelled); + session.unsubscribe(listener); + assert!(cancelled_at_boundary.load(Ordering::SeqCst)); + assert!(harness.faux.requests().is_empty()); + let queue = runtime.runs.queue_snapshot(&task); + let texts: Vec<_> = queue.items.iter().map(|item| item.text.as_str()).collect(); + assert_eq!(texts, ["First", "Second", "Third", "Later"]); + assert!( + session + .messages() + .iter() + .all(|message| !matches!(message, SessionMessage::Llm(Message::User(_)))) + ); + } } diff --git a/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs b/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs index 0a17ec8e1..020328c28 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/runtime.rs @@ -524,6 +524,19 @@ impl AgentRuntime { &self.models } + /// Every request containing task history, including side questions, uses + /// the same task-scoped attachment boundary. + pub(super) fn attachment_models(&self, session_id: &str) -> ModelRegistry { + super::provider::attachment_model_registry( + &self.models, + super::attachments::AgentAttachmentStore::new( + Arc::clone(self.store.store()), + self.store.account(), + ), + session_id.to_string(), + ) + } + /// A task's session for a side question: the loaded one as it is, or /// the task loaded on the model it last ran on. pub(super) async fn side_question_session( @@ -682,15 +695,17 @@ impl AgentRuntime { .map_err(|error| format!("Failed to read the Agent task: {error}"))??; // A model that cannot see images gets them described. let describe = !model.supports_images(); - let mut options = - AgentSessionOptions::new(&row.project_root, APP_NAME, manager, self.models.clone()); - options.model = Some(model); - options.settings = self.settings(); - options.resources = resources.unwrap_or_default(); let attachments = Arc::new(super::attachments::AgentAttachmentStore::new( Arc::clone(self.store.store()), self.store.account(), )); + let task_models = self.attachment_models(&row.id); + let mut options = + AgentSessionOptions::new(&row.project_root, APP_NAME, manager, task_models); + options.model = Some(model); + options.settings = self.settings(); + options.resources = resources.unwrap_or_default(); + let output_dir = crate::store::spill::SpillRoot::new(self.store.store().paths()) .and_then(|spills| spills.output_dir(&row.id)) .map_err(|error| error.to_string())?; diff --git a/apps/maple-agent/crates/maple-agent/src/agent/side_models.rs b/apps/maple-agent/crates/maple-agent/src/agent/side_models.rs index 85668904a..cde6fd6cf 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/side_models.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/side_models.rs @@ -563,7 +563,7 @@ impl AgentRuntimeHandle { } messages.extend(side_question_messages(&model, &prior, &question)); tokio::spawn(stream_side_question( - runtime.models().stream_fn(), + runtime.attachment_models(session_id).stream_fn(), model, Context::from_messages(messages), session_id.to_string(), diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs b/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs index 0334339e1..9dd819a99 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tasks.rs @@ -6,6 +6,7 @@ use std::path::Path; use std::sync::Arc; +use std::time::{Duration, Instant}; use icu_properties::{CodePointSetData, props::DefaultIgnorableCodePoint}; @@ -32,6 +33,9 @@ pub(super) const DEFAULT_AGENT_SESSION_TITLE: &str = "New task"; pub(super) const ACP_SESSION_FALLBACK_TITLE: &str = "Maple ACP"; /// Why `/compact` did nothing: the history is all recent enough to keep. pub const NOTHING_TO_COMPACT_ERROR: &str = "There is nothing to compact yet"; +const HISTORY_PREPARATION_BUDGET: Duration = Duration::from_secs(5); +const HISTORY_PREPARING: &str = + "Maple is preparing this task's saved history. Try opening it again shortly."; /// A title from a task's first prompt: its text on one line, shortened. pub(super) fn session_title_from_prompt(prompt: &str) -> String { @@ -171,7 +175,7 @@ impl AgentRuntimeHandle { request: crate::store::OpenTaskPageRequest, transport: crate::store::PageTransport, ) -> Result { - self.verify_generation().await?; + self.ensure_task_history_ready(&request.task).await?; let tasks = self.store()?; tokio::task::spawn_blocking(move || { tasks @@ -183,6 +187,137 @@ impl AgentRuntimeHandle { .map_err(|_| "Task page worker failed".to_string())? } + /// Migration deliberately removes obsolete derived history. Opening a task + /// gives its existing bounded projector immediate work instead of showing a + /// false empty conversation while it waits in the background queue. ACP uses + /// this same boundary before capturing its immutable history replay. + pub async fn ensure_task_history_ready(&self, task: &str) -> Result<(), String> { + self.verify_generation().await?; + let handle = self.clone(); + let task = task.to_owned(); + let deadline = Instant::now() + HISTORY_PREPARATION_BUDGET; + tokio::task::spawn_blocking(move || handle.prepare_history_for_read(&task, deadline)) + .await + .map_err(|_| "Task history preparation worker failed".to_string())? + } + + pub(super) fn prepare_history_for_read( + &self, + task: &str, + deadline: Instant, + ) -> Result<(), String> { + use crate::store::{RebuildProgress, StoreError}; + + let tasks = self.store()?; + let store = tasks.store(); + let account = tasks.account(); + let state = &self.service.state; + let preparing = |error: StoreError| match error { + StoreError::Busy => HISTORY_PREPARING.to_string(), + other => other.to_string(), + }; + 'readiness: loop { + self.verify_generation_now()?; + match store.history_needs_rebuild(account, task) { + Ok(false) => return Ok(()), + Ok(true) => {} + Err(StoreError::Busy) => { + self.retry_history_preparation(deadline)?; + continue; + } + Err(error) => return Err(preparing(error)), + } + if Instant::now() >= deadline { + return Err(HISTORY_PREPARING.to_string()); + } + let ticket = loop { + let result = super::maintenance::while_task_idle(state, account, task, || { + self.verify_generation_now() + .map_err(|_| StoreError::NotFound)?; + // Completion can race the initial read or a Busy retry. + // Under admission, a finished task needs no new epoch. + if !store.history_needs_rebuild(account, task)? { + return Ok(None); + } + store.begin_timeline_rebuild(account, task, false).map(Some) + }); + match result { + Ok(Some(ticket)) => break ticket, + Ok(None) => return Ok(()), + Err(StoreError::Busy) => self.retry_history_preparation(deadline)?, + Err(error) => return Err(preparing(error)), + } + }; + let prepared = loop { + match store.prepare_timeline_rebuild_before(&ticket, Some(deadline)) { + Ok(prepared) => break prepared, + Err(StoreError::Busy) => self.retry_history_preparation(deadline)?, + // Another explicit open or the background worker completed + // this epoch. Recheck instead of reporting a failed load. + Err(StoreError::Conflict) => continue 'readiness, + Err(error) => return Err(preparing(error)), + } + }; + loop { + self.verify_generation_now()?; + let chunk = match store.prepare_timeline_chunk_before(&prepared, Some(deadline)) { + Ok(chunk) => chunk, + Err(StoreError::Busy) => { + self.retry_history_preparation(deadline)?; + continue; + } + Err(StoreError::Conflict) => continue 'readiness, + Err(error) => return Err(preparing(error)), + }; + loop { + let progress = + super::maintenance::while_task_idle(state, account, task, || { + self.verify_generation_now() + .map_err(|_| StoreError::NotFound)?; + store.commit_timeline_chunk(&chunk, false) + }); + match progress { + Ok(RebuildProgress::Pending { .. }) => break, + Ok(RebuildProgress::NeedsSnapshot) | Err(StoreError::Conflict) => { + continue 'readiness; + } + Ok(RebuildProgress::Replaced { epoch }) => { + emit_agent_event( + &self.service.host.events, + AgentServiceEvent::HistoryReplaced { + user_id: self.user_id.to_string(), + session_id: task.to_owned(), + rows_epoch: epoch, + }, + ); + return Ok(()); + } + // Keep the encoded chunk while admission or the writer + // is busy. Commit still checks its source and cursor. + Err(StoreError::Busy) => self.retry_history_preparation(deadline)?, + Err(error) => return Err(preparing(error)), + } + } + } + } + } + + fn retry_history_preparation(&self, deadline: Instant) -> Result<(), String> { + self.verify_generation_now()?; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(HISTORY_PREPARING.to_string()); + } + // This runs only in the blocking worker, after all admission, writer + // and read guards are released. Brief contention consumes the existing + // budget rather than surfacing as an immediate failed task open. + std::thread::sleep(remaining.min(Duration::from_millis(10))); + if Instant::now() >= deadline { + return Err(HISTORY_PREPARING.to_string()); + } + Ok(()) + } + pub async fn history_forward_page( &self, request: crate::store::HistoryPageRequest, @@ -370,6 +505,7 @@ impl AgentRuntimeHandle { } pub async fn load_session(&self, session_id: String) -> Result { + self.ensure_task_history_ready(&session_id).await?; let runtime = self.current_runtime().await?; let store = self.store()?; let row = store @@ -694,6 +830,27 @@ impl AgentRuntimeHandle { Ok(()) } + /// ACP's protocol-native images enter the same file store as desktop + /// imports and remote streams. Sending subsequently uses committed IDs. + pub async fn import_image_attachments( + &self, + session_id: String, + images: Vec, + ) -> Result, String> { + self.verify_generation().await?; + let tasks = self.store()?; + if tasks.get(&session_id)?.is_none() { + return Err(format!("Failed to find Agent task {session_id}")); + } + let attachments = super::attachments::AgentAttachmentStore::new( + Arc::clone(tasks.store()), + tasks.account(), + ); + tokio::task::spawn_blocking(move || attachments.store_uploads(&session_id, &images)) + .await + .map_err(|error| format!("Attachment import failed: {error}"))? + } + /// Raw bytes of a stored image attachment, for display. pub async fn read_image_attachment( &self, diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/tests.rs index e56e03abc..1f7e258b6 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tests.rs @@ -1087,11 +1087,25 @@ async fn steering_joins_the_running_turn() { const PNG_1X1: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="; -fn with_image(mut request: AgentSendMessageRequest, name: &str) -> AgentSendMessageRequest { - request.attachments = vec![AgentImageUpload { - name: name.to_string(), - data_url: format!("data:image/png;base64,{PNG_1X1}"), - }]; +async fn with_image( + harness: &Harness, + mut request: AgentSendMessageRequest, + name: &str, +) -> AgentSendMessageRequest { + request.attachments = harness + .handle + .import_image_attachments( + request.session_id.clone(), + vec![AgentImageUpload { + name: name.to_string(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .await + .unwrap() + .into_iter() + .map(|attachment| attachment.id) + .collect(); request } @@ -1113,36 +1127,65 @@ fn last_user(request: &pi_ai::faux::FauxRequest) -> pi_ai::UserMessage { async fn attached_images_reach_the_model_and_show_with_their_message() { use base64::Engine as _; - let harness = - Harness::with_faux(FauxProvider::new().with_chunk_delay(Duration::from_millis(10))).await; - harness - .faux - .push_text("A dialog, streamed slowly enough to queue behind."); + let harness = Harness::new().await; + harness.faux.push_message(vec![faux_tool_call( + "request_user_input", + json!({"question":"Wait?","options":[{"label":"Continue"}]}), + )]); + harness.faux.push_text("A dialog."); harness.faux.push_text("Also a dialog."); let task = harness.create_task().await; - let first = with_image(harness.request(&task, "What is this?"), "dialog.png"); + let first = with_image( + &harness, + harness.request(&task, "What is this?"), + "dialog.png", + ) + .await; let mut run = harness.handle.send_message(first).await.unwrap(); + let question = || { + harness + .recorder + .events() + .iter() + .find_map(|event| match event { + AgentServiceEvent::Question { + session_id, + request_id, + .. + } if session_id == &task => Some(request_id.clone()), + _ => None, + }) + }; + eventually(|| question().is_some()).await; // A message sent during the run waits as a chip with its image. - let mut second = with_image(harness.request(&task, "And this?"), "second.png"); + // The question holds the first turn while the durable image import finishes. + let mut second = with_image(&harness, harness.request(&task, "And this?"), "second.png").await; second.vision_capable = true; let staged = harness.handle.send_message(second).await.unwrap(); + assert_eq!(staged.run_id, run.run_id); assert_eq!(staged.queued.unwrap().attachments[0].name, "second.png"); + release_attachment_question(&harness, &question().unwrap()).await; assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); // A model without vision gets the image by source, to look with // read_image; one with vision gets it beside the text. let requests = harness.faux.requests(); + assert_eq!( + requests.len(), + 3, + "first question, its answer, then the queued image" + ); let first = last_user(&requests[0]); let text = pi_ai::content_text(&first.content); assert!( text.starts_with( - "What is this?\n\nThe user attached the following image:\n- \"dialog.png\": maple-attachment://" + "What is this?\n\nThe user attached these files on the execution computer:\n- " ), "{text}" ); - assert!(text.ends_with("when you need visual details from that image.")); + assert!(text.ends_with("when you need visual details from an image.")); assert_eq!(first.content.len(), 1); - let second = last_user(&requests[1]); + let second = last_user(&requests[2]); assert!(!pi_ai::content_text(&second.content).contains("read_image")); assert!(matches!( &second.content[1], @@ -1158,7 +1201,7 @@ async fn attached_images_reach_the_model_and_show_with_their_message() { let images = row .input .as_ref() - .and_then(|input| input["imageAttachments"].as_array().cloned()) + .and_then(|input| input["attachments"].as_array().cloned()) .unwrap_or_default() .iter() .map(|image| { @@ -1196,6 +1239,569 @@ async fn attached_images_reach_the_model_and_show_with_their_message() { ); } +async fn committed_file( + harness: &Harness, + task: &str, + name: &str, + mime: &str, + bytes: &[u8], +) -> AgentAttachment { + use sha2::{Digest, Sha256}; + let runtime = harness.runtime().await; + let store = runtime.store.store(); + let metadata = crate::store::AttachmentUploadMetadata { + name: name.into(), + mime_type: mime.into(), + size: bytes.len() as u64, + sha256: format!("{:x}", Sha256::digest(bytes)), + }; + let started = store + .begin_attachment_upload( + runtime.store.account(), + task, + &format!("{:032x}", rand::random::()), + &metadata, + ) + .unwrap(); + let crate::store::AttachmentUploadStart::Writing(mut writer) = started else { + panic!("new upload") + }; + writer.write_chunk(bytes).unwrap(); + writer.finish(store).unwrap().into() +} + +#[tokio::test] +async fn mixed_files_and_images_use_committed_paths_and_preserve_refs_in_queue_edits_and_steering() +{ + let harness = Harness::new().await; + let task = harness.create_task().await; + let zip = committed_file( + &harness, + &task, + "report.zip", + "application/zip", + b"PK archive bytes", + ) + .await; + let image = harness + .handle + .import_image_attachments( + task.clone(), + vec![AgentImageUpload { + name: "screen.png".into(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .await + .unwrap() + .remove(0); + harness.faux.push_message(vec![faux_tool_call( + "request_user_input", + json!({"question":"Wait?","options":[{"label":"Continue"}]}), + )]); + harness.faux.push_text("Reviewed."); + let mut first = harness.request(&task, "Wait before reviewing"); + first.vision_capable = true; + let mut run = harness.handle.send_message(first).await.unwrap(); + let question = || { + harness + .recorder + .events() + .iter() + .find_map(|event| match event { + AgentServiceEvent::Question { + session_id, + request_id, + .. + } if session_id == &task => Some(request_id.clone()), + _ => None, + }) + }; + eventually(|| question().is_some()).await; + let mut upload = harness.request(&task, "Review both"); + upload.vision_capable = true; + upload.attachments = vec![zip.id.clone(), image.id.clone()]; + let chip = harness + .handle + .send_message(upload) + .await + .unwrap() + .queued + .unwrap(); + assert_eq!(chip.attachments, [zip.clone(), image.clone()]); + let control = AgentQueueControlRequest { + session_id: task.clone(), + queue_id: chip.queue_id.clone(), + }; + harness + .handle + .begin_queued_message_edit(control) + .await + .unwrap(); + let session = harness.runtime().await.loaded_session(&task).await.unwrap(); + let steered = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let listener = session.subscribe({ + let steered = steered.clone(); + move |event| { + if matches!(event, pi_coding_agent::AgentSessionEvent::QueueUpdate { steering, .. } + if steering.iter().any(|text| text.starts_with("Review the archive and screenshot"))) + { + steered.store(true, std::sync::atomic::Ordering::Release); + } + } + }); + // Editing the queued message only changes its text; the original committed + // file identities and image references survive the steering boundary. + let mut edit = harness.request(&task, "Review the archive and screenshot"); + edit.queue_id = Some(chip.queue_id); + edit.steer = true; + harness.handle.send_message(edit).await.unwrap(); + eventually(|| steered.load(std::sync::atomic::Ordering::Acquire)).await; + session.unsubscribe(listener); + assert!( + harness + .service + .answer_question( + &question().unwrap(), + r#"{"answers":{"question_0":{"answers":["Continue"]}}}"#.into() + ) + .await + ); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let request = harness.faux.requests().pop().unwrap(); + let user = last_user(&request); + let text = pi_ai::content_text(&user.content); + let (typed, shown) = attachments::split_attachment_prompt(&text).unwrap(); + assert_eq!(typed, "Review the archive and screenshot"); + assert_eq!(shown, [zip.clone(), image]); + assert!(matches!(&user.content[1], pi_ai::Content::Image(image) if image.data == PNG_1X1)); + assert_eq!( + std::fs::read(&zip.execution_path).unwrap(), + b"PK archive bytes" + ); + assert!(!harness.project.path().join("report.zip").exists()); + let detail = harness.handle.load_session(task).await.unwrap(); + let input = detail + .timeline + .iter() + .rev() + .find(|row| row.role.as_deref() == Some("user")) + .unwrap() + .input + .as_ref() + .unwrap(); + assert_eq!(input["attachments"][0]["executionPath"], zip.execution_path); +} + +#[tokio::test] +async fn an_invalid_or_foreign_attachment_rejects_the_whole_send_before_inference() { + let harness = Harness::new().await; + let task = harness.create_task().await; + let other = harness.create_task().await; + let valid = committed_file(&harness, &task, "notes.txt", "text/plain", b"notes").await; + let foreign = committed_file(&harness, &other, "foreign.zip", "application/zip", b"PK").await; + for bad in ["not-an-id".to_string(), foreign.id] { + let mut request = harness.request(&task, "Do not partially accept this"); + request.attachments = vec![valid.id.clone(), bad]; + assert!(harness.handle.send_message(request).await.is_err()); + } + assert!(harness.faux.requests().is_empty()); + let detail = harness.handle.load_session(task).await.unwrap(); + assert!( + !detail + .timeline + .iter() + .any(|row| row.role.as_deref() == Some("user")) + ); +} + +#[tokio::test] +async fn cold_native_image_history_materializes_only_for_inference_and_missing_files_are_nonfatal() +{ + let harness = Harness::new().await; + let task = harness.create_task().await; + let image = harness + .handle + .import_image_attachments( + task.clone(), + vec![AgentImageUpload { + name: "screen.png".into(), + data_url: format!("data:image/png;base64,{PNG_1X1}"), + }], + ) + .await + .unwrap() + .remove(0); + harness.faux.push_text("Seen."); + let mut first = harness.request(&task, "Inspect"); + first.vision_capable = true; + first.attachments = vec![image.id]; + let mut run = harness.handle.send_message(first).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let runtime = harness.runtime().await; + runtime.unload_task(&task).await.unwrap(); + let source = || { + runtime + .store + .store() + .load_session( + runtime.store.account(), + &task, + crate::store::SessionLoad::Full, + ) + .unwrap() + .entries + }; + let saved = source(); + let saved_image_ids = saved + .iter() + .flat_map(|entry| match &entry.kind { + pi_coding_agent::session::EntryKind::Message { + message: pi_coding_agent::SessionMessage::Llm(pi_ai::Message::User(user)), + } => user + .content + .iter() + .filter_map(|block| match block { + pi_ai::Content::Image(image) => { + attachments::attachment_id_from_source(&image.data).map(str::to_string) + } + _ => None, + }) + .collect::>(), + _ => Vec::new(), + }) + .collect::>(); + assert!(!saved_image_ids.is_empty()); + // Pi may normalize the incoming image into a separate persisted copy. + // Remove the pixels actually referenced by cold source history. + for id in saved_image_ids { + let image = runtime + .store + .store() + .attachment_metadata(runtime.store.account(), &task, &id) + .unwrap(); + std::fs::remove_file(image.execution_path).unwrap(); + } + harness + .faux + .push_text("I can continue without the missing image."); + let mut next = harness.request(&task, "Continue"); + next.vision_capable = true; + let mut run = harness.handle.send_message(next).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let request = harness.faux.requests().pop().unwrap(); + assert!(request.context.messages.iter().any(|message| matches!(message, + pi_ai::Message::User(user) if user.content.iter().any(|block| + matches!(block, pi_ai::Content::Text(text) if text.text == "Image no longer available."))))); + let after = source(); + assert_eq!(&after[..saved.len()], saved); +} + +#[tokio::test] +async fn cold_side_questions_materialize_images_and_degrade_missing_history() { + let harness = Harness::new().await; + let task = harness.create_task().await; + let mut first = with_image(&harness, harness.request(&task, "Inspect"), "screen.png").await; + first.vision_capable = true; + harness.faux.push_text("Seen."); + let mut run = harness.handle.send_message(first).await.unwrap(); + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let runtime = harness.runtime().await; + runtime.unload_task(&task).await.unwrap(); + let saved = runtime + .store + .store() + .load_session( + runtime.store.account(), + &task, + crate::store::SessionLoad::Full, + ) + .unwrap() + .entries; + let image_ids: Vec = saved + .iter() + .flat_map(|entry| match &entry.kind { + pi_coding_agent::session::EntryKind::Message { + message: pi_coding_agent::SessionMessage::Llm(pi_ai::Message::User(user)), + } => user + .content + .iter() + .filter_map(|block| match block { + pi_ai::Content::Image(image) => { + attachments::attachment_id_from_source(&image.data).map(str::to_string) + } + _ => None, + }) + .collect::>(), + _ => Vec::new(), + }) + .collect(); + assert!(!image_ids.is_empty()); + for missing in [false, true] { + runtime.unload_task(&task).await.unwrap(); + if missing { + for id in &image_ids { + let image = runtime + .store + .store() + .attachment_metadata(runtime.store.account(), &task, id) + .unwrap(); + std::fs::remove_file(image.execution_path).unwrap(); + } + } + let request_id = format!("side-image-{missing}"); + harness.faux.push_text("Side answer."); + harness + .handle + .ask_side_question( + &task, + request_id.clone(), + Vec::new(), + "What was in the image?".into(), + ) + .await + .unwrap(); + eventually(|| harness.recorder.events().iter().any(|event| matches!(event, + AgentServiceEvent::SideQuestion { session_id, request_id: id, event } + if session_id == &task && id == &request_id && !matches!(event, SideQuestionEvent::Chunk(_)) + ))).await; + assert!(harness.recorder.events().iter().any(|event| matches!(event, + AgentServiceEvent::SideQuestion { session_id, request_id: id, event: SideQuestionEvent::Finished } + if session_id == &task && id == &request_id + ))); + let request = harness.faux.requests().pop().unwrap(); + let content: Vec<_> = request + .context + .messages + .iter() + .flat_map(|message| match message { + pi_ai::Message::User(user) => user.content.iter().collect::>(), + _ => Vec::new(), + }) + .collect(); + if missing { + assert!(content.iter().any(|block| matches!(block, + pi_ai::Content::Text(text) if text.text == "Image no longer available."))); + assert!( + !content + .iter() + .any(|block| matches!(block, pi_ai::Content::Image(_))) + ); + } else { + assert!(content.iter().any(|block| matches!(block, + pi_ai::Content::Image(image) if image.data == PNG_1X1))); + } + assert_eq!( + runtime + .store + .store() + .load_session( + runtime.store.account(), + &task, + crate::store::SessionLoad::Full + ) + .unwrap() + .entries, + saved + ); + } +} + +#[tokio::test] +async fn pasted_foreign_file_listing_does_not_poison_later_requests() { + let harness = Harness::new().await; + let task = harness.create_task().await; + let other = harness.create_task().await; + let foreign = committed_file( + &harness, + &other, + "notes.txt", + "text/plain", + b"private notes", + ) + .await; + let pasted = attachments::attachment_prompt("Earlier notes said:", &[foreign], false); + for text in [pasted.as_str(), "hello", "are you there?"] { + harness.faux.push_text("Yes."); + let mut run = harness.send(&task, text).await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let request = harness.faux.requests().pop().unwrap(); + assert_eq!(last_user_text(&request), text); + assert_eq!(user_texts(&request)[0], pasted); + } +} + +async fn attachment_question_barrier(harness: &Harness, task: &str) -> (AgentRunHandle, String) { + harness.faux.push_message(vec![faux_tool_call( + "request_user_input", + json!({"question":"Wait?","options":[{"label":"Continue"}]}), + )]); + harness.faux.push_text("Released."); + let mut first = harness.request(task, "Wait before continuing"); + first.vision_capable = true; + let run = harness.handle.send_message(first).await.unwrap(); + let question = || { + harness + .recorder + .events() + .iter() + .find_map(|event| match event { + AgentServiceEvent::Question { + session_id, + request_id, + .. + } if session_id == task => Some(request_id.clone()), + _ => None, + }) + }; + eventually(|| question().is_some()).await; + (run, question().unwrap()) +} + +async fn release_attachment_question(harness: &Harness, question: &str) { + assert!( + harness + .service + .answer_question( + question, + r#"{"answers":{"question_0":{"answers":["Continue"]}}}"#.into(), + ) + .await + ); +} + +fn change_attachment_bytes(path: &str) { + let mut bytes = std::fs::read(path).unwrap(); + *bytes.last_mut().unwrap() ^= 1; + std::fs::write(path, bytes).unwrap(); +} + +#[tokio::test] +async fn failed_batch_image_preparation_preserves_unsent_messages_in_order() { + for failed_index in [0, 1] { + let harness = Harness::new().await; + let task = harness.create_task().await; + let (mut run, question) = attachment_question_barrier(&harness, &task).await; + let mut failed_image = None; + let texts = ["First queued", "Second queued", "Third queued"]; + for (index, text) in texts.iter().enumerate() { + let mut request = harness.request(&task, text); + request.vision_capable = true; + if index == failed_index { + request = with_image(&harness, request, "queued.png").await; + } + let chip = harness + .handle + .send_message(request) + .await + .unwrap() + .queued + .unwrap(); + if index == failed_index { + failed_image = Some(chip.attachments[0].clone()); + } + } + let image = failed_image.unwrap(); + change_attachment_bytes(&image.execution_path); + // The valid first message still runs when a later message cannot be prepared. + harness.faux.push_text("First queued completed."); + release_attachment_question(&harness, &question).await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Failed); + let detail = harness.handle.load_session(task.clone()).await.unwrap(); + let queued: Vec<_> = detail + .queue + .items + .iter() + .map(|chip| chip.text.as_str()) + .collect(); + assert_eq!(queued, texts[failed_index..]); + assert_eq!(detail.queue.items[0].attachments, [image]); + let users: Vec<_> = detail + .timeline + .iter() + .filter(|row| row.role.as_deref() == Some("user")) + .filter_map(|row| row.text.as_deref()) + .collect(); + let expected = if failed_index == 0 { + vec!["Wait before continuing"] + } else { + vec!["Wait before continuing", "First queued"] + }; + assert_eq!(users, expected); + assert!(harness.recorder.run_events(&run.run_id).iter().any(|event| matches!(event, + AgentRunEvent::Error(item) if item.text.as_deref().is_some_and(|text| text.contains("attachment")) + ))); + } +} + +#[tokio::test] +async fn failed_steered_image_preserves_accepted_messages_and_reports_an_error() { + let harness = Harness::new().await; + let task = harness.create_task().await; + let (mut run, question) = attachment_question_barrier(&harness, &task).await; + let mut request = with_image( + &harness, + harness.request(&task, "Inspect this image"), + "steered.png", + ) + .await; + request.vision_capable = true; + request.steer = true; + let runtime = harness.runtime().await; + let image: AgentAttachment = runtime + .store + .store() + .attachment_metadata(runtime.store.account(), &task, &request.attachments[0]) + .unwrap() + .into(); + // A same-size rewrite passes admission's file checks; image preparation + // must detect it without losing the accepted steer. + change_attachment_bytes(&image.execution_path); + let joined = harness.handle.send_message(request).await.unwrap(); + assert_eq!(joined.run_id, run.run_id); + let mut later = harness.request(&task, "Then use the result"); + later.steer = true; + later.vision_capable = true; + harness.handle.send_message(later).await.unwrap(); + eventually(|| { + harness + .recorder + .run_events(&run.run_id) + .iter() + .any(|event| matches!(event, AgentRunEvent::Error(_))) + }) + .await; + release_attachment_question(&harness, &question).await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Failed); + let detail = harness.handle.load_session(task).await.unwrap(); + assert_eq!( + detail + .queue + .items + .iter() + .map(|chip| chip.text.as_str()) + .collect::>(), + ["Inspect this image", "Then use the result"] + ); + assert_eq!(detail.queue.items[0].attachments, [image]); + let users: Vec<_> = detail + .timeline + .iter() + .filter(|row| row.role.as_deref() == Some("user")) + .filter_map(|row| row.text.as_deref()) + .collect(); + assert_eq!(users, ["Wait before continuing"]); + assert_eq!( + harness + .recorder + .run_events(&run.run_id) + .iter() + .filter(|event| matches!(event, AgentRunEvent::Error(_))) + .count(), + 1 + ); +} + /// The first tool result the model got in `request`. fn first_tool_result(request: &pi_ai::faux::FauxRequest) -> pi_ai::ToolResultMessage { request @@ -2687,3 +3293,218 @@ async fn viewing_unloaded_task_thinking_and_history_does_not_load_context() { assert_eq!(before.last_seq, after.last_seq); assert_eq!(before.facts, after.facts); } + +#[tokio::test] +async fn sending_before_open_rebuilds_retired_history_before_run_admission() { + let harness = Harness::new().await; + let task = harness.create_task().await; + harness.faux.push_text("Old reply stays."); + let mut first = harness.send(&task, "Old question stays.").await; + assert_eq!(finished(&mut first).await, AgentRunTerminal::Completed); + let (store, account) = harness.handle.test_persistence().unwrap(); + store.retire_task_timeline_for_test(account, &task).unwrap(); + assert!(store.history_needs_rebuild(account, &task).unwrap()); + + // Do not open/load this task after invalidation. Send must prepare its old + // history while idle, before registering an active run that defers rebuilds. + harness.faux.push_text("New reply appears."); + let mut second = harness.send(&task, "New question appears.").await; + assert_eq!(finished(&mut second).await, AgentRunTerminal::Completed); + assert!(!store.history_needs_rebuild(account, &task).unwrap()); + let page = harness + .handle + .open_task_page( + crate::store::OpenTaskPageRequest { + task: task.clone(), + limit_rows: None, + limit_bytes: None, + }, + crate::store::PageTransport::Remote, + ) + .await + .unwrap(); + let text = serde_json::to_string(&page.rows).unwrap(); + for expected in [ + "Old question stays.", + "Old reply stays.", + "New question appears.", + "New reply appears.", + ] { + assert!(text.contains(expected), "missing {expected}: {text}"); + } + let entries = store + .load_session(account, &task, crate::store::SessionLoad::Full) + .unwrap() + .entries; + let source = serde_json::to_string(&entries).unwrap(); + assert_eq!(source.matches("Old question stays.").count(), 1); + assert_eq!(source.matches("New question appears.").count(), 1); + let requests = harness.faux.requests(); + assert_eq!( + user_texts(requests.last().unwrap()), + ["Old question stays.", "New question appears."] + ); +} + +async fn history_preparation_fixture() -> ( + Harness, + String, + Arc, + crate::store::AccountId, +) { + let harness = Harness::new().await; + let task = harness.create_task().await; + harness.faux.push_text("Saved history stays visible."); + let mut run = harness.send(&task, "Keep this history.").await; + assert_eq!(finished(&mut run).await, AgentRunTerminal::Completed); + let (store, account) = harness.handle.test_persistence().unwrap(); + (harness, task, store, account) +} + +#[tokio::test] +async fn history_preparation_retries_brief_lifecycle_contention() { + let (harness, task, store, account) = history_preparation_fixture().await; + let admission = harness.service.state.lifecycle.lock().await; + store.retire_task_timeline_for_test(account, &task).unwrap(); + let handle = harness.handle.clone(); + let preparing_task = task.clone(); + let (started, entered) = tokio::sync::oneshot::channel(); + let preparation = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + handle.prepare_history_for_read( + &preparing_task, + std::time::Instant::now() + Duration::from_secs(5), + ) + }); + entered.await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !preparation.is_finished(), + "brief contention must not fail the open" + ); + drop(admission); + tokio::time::timeout(Duration::from_secs(6), preparation) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(!store.history_needs_rebuild(account, &task).unwrap()); + let page = store + .open_task_page( + account, + &crate::store::OpenTaskPageRequest { + task, + limit_rows: None, + limit_bytes: None, + }, + crate::store::PageTransport::Local, + ) + .unwrap(); + assert!( + serde_json::to_string(&page.rows) + .unwrap() + .contains("Saved history stays visible.") + ); +} + +#[tokio::test] +async fn history_preparation_contention_stops_at_its_existing_deadline() { + let (harness, task, store, account) = history_preparation_fixture().await; + let _admission = harness.service.state.lifecycle.lock().await; + store.retire_task_timeline_for_test(account, &task).unwrap(); + let stamp = store.task_history_stamp(account, &task).unwrap(); + let handle = harness.handle.clone(); + let preparing_task = task.clone(); + let budget = Duration::from_millis(150); + let preparation = tokio::task::spawn_blocking(move || { + let started = std::time::Instant::now(); + let result = handle.prepare_history_for_read(&preparing_task, started + budget); + (started.elapsed(), result) + }); + let (elapsed, result) = tokio::time::timeout(Duration::from_secs(2), preparation) + .await + .unwrap() + .unwrap(); + assert!(elapsed >= budget); + assert_eq!( + result.unwrap_err(), + "Maple is preparing this task's saved history. Try opening it again shortly." + ); + assert_eq!(store.task_history_stamp(account, &task).unwrap(), stamp); + assert!(store.history_needs_rebuild(account, &task).unwrap()); +} + +async fn history_preparation_idle_write( + mut operation: impl FnMut() -> Result, +) -> T { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + match operation() { + Ok(value) => return value, + Err(crate::store::StoreError::Busy) => { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Err(error) => panic!("history fixture idle write failed: {error}"), + } + } + }) + .await + .expect("fixture writer contention clears") +} + +#[tokio::test] +async fn history_preparation_rechecks_completion_before_starting_another_epoch() { + let (harness, task, store, account) = history_preparation_fixture().await; + let admission = harness.service.state.lifecycle.lock().await; + store.retire_task_timeline_for_test(account, &task).unwrap(); + let handle = harness.handle.clone(); + let preparing_task = task.clone(); + let (started, entered) = tokio::sync::oneshot::channel(); + let preparation = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + handle.prepare_history_for_read( + &preparing_task, + std::time::Instant::now() + Duration::from_secs(5), + ) + }); + entered.await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!preparation.is_finished()); + + // Complete the same work while the foreground request waits for admission. + // This test owns the lifecycle gate that normal background commits take. + let ticket = + history_preparation_idle_write(|| store.begin_timeline_rebuild(account, &task, false)) + .await; + let prepared = store.prepare_timeline_rebuild(&ticket).unwrap(); + let completed_epoch = loop { + let chunk = store.prepare_timeline_chunk(&prepared).unwrap(); + match history_preparation_idle_write(|| store.commit_timeline_chunk(&chunk, false)).await { + crate::store::RebuildProgress::Pending { .. } => {} + crate::store::RebuildProgress::Replaced { epoch } => break epoch, + crate::store::RebuildProgress::NeedsSnapshot => panic!("source did not change"), + } + }; + drop(admission); + tokio::time::timeout(Duration::from_secs(6), preparation) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + store.task_history_stamp(account, &task).unwrap().rows_epoch, + completed_epoch + ); + assert!( + !harness + .recorder + .events + .lock() + .unwrap() + .iter() + .any(|event| matches!( + event, AgentServiceEvent::HistoryReplaced { session_id, .. } if session_id == &task + )), + "waiting foreground open must not schedule a redundant rebuild" + ); +} diff --git a/apps/maple-agent/crates/maple-agent/src/agent/timeline.rs b/apps/maple-agent/crates/maple-agent/src/agent/timeline.rs index 3c4cdce29..de6aa273b 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/timeline.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/timeline.rs @@ -25,7 +25,7 @@ use pi_coding_agent::session::{EntryKind, SessionManager}; use serde::Serialize; use serde_json::{Value, json}; -use super::attachments::split_image_prompt; +use super::attachments::split_attachment_prompt; use super::external_agents; use super::types::AgentTimelineItem; @@ -271,20 +271,8 @@ pub(crate) fn user_item( merge: &str, ) -> AgentTimelineItem { let text = content_text(content); - let (text, input) = match split_image_prompt(&text) { - Some((typed, attachments)) => { - let attachments: Vec = attachments - .into_iter() - .map(|attachment| { - json!({ - "id": attachment.id, - "name": attachment.name, - "source": attachment.source, - }) - }) - .collect(); - (typed, Some(json!({ "imageAttachments": attachments }))) - } + let (text, input) = match split_attachment_prompt(&text) { + Some((typed, attachments)) => (typed, Some(json!({ "attachments": attachments }))), None => (text, None), }; AgentTimelineItem { diff --git a/apps/maple-agent/crates/maple-agent/src/agent/tools/read_image/tests.rs b/apps/maple-agent/crates/maple-agent/src/agent/tools/read_image/tests.rs index 5e4a474f4..eac299ed0 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/tools/read_image/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/tools/read_image/tests.rs @@ -164,7 +164,7 @@ async fn an_attachment_of_the_task_is_read() { }], ) .unwrap(); - let source = stored[0].attachment.source.clone(); + let source = stored[0].source.clone(); assert!(source.starts_with("maple-attachment://")); let result = run(&fixture.tool(None), json!({ "source": source })).await; assert!(!result.is_error, "{}", text(&result)); diff --git a/apps/maple-agent/crates/maple-agent/src/agent/types.rs b/apps/maple-agent/crates/maple-agent/src/agent/types.rs index 5d8ac3bad..8bfdf377d 100644 --- a/apps/maple-agent/crates/maple-agent/src/agent/types.rs +++ b/apps/maple-agent/crates/maple-agent/src/agent/types.rs @@ -12,8 +12,8 @@ pub use pi_ai::ThinkingLevel; use serde::{Deserialize, Serialize}; use serde_json::Value; -use super::attachments::AgentImageAttachment; -pub use super::attachments::AgentImageUpload; +use super::attachments::AgentAttachment; + use super::{DEFAULT_AGENT_MODEL, DEFAULT_MCP_TIMEOUT_SECONDS}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -526,7 +526,7 @@ pub struct AgentSendMessageRequest { #[serde(default)] pub queue_id: Option, #[serde(default)] - pub attachments: Vec, + pub attachments: Vec, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -576,7 +576,7 @@ pub struct AgentQueuedMessage { pub message_id: String, pub session_id: String, pub text: String, - pub attachments: Vec, + pub attachments: Vec, pub created_ms: u128, } diff --git a/apps/maple-agent/crates/maple-agent/src/store/attachment_cleanup.rs b/apps/maple-agent/crates/maple-agent/src/store/attachment_cleanup.rs new file mode 100644 index 000000000..729b3cf58 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/attachment_cleanup.rs @@ -0,0 +1,824 @@ +//! Bounded cleanup of an attachment directory whose owner has authorized deletion. +//! +//! Model-created adjacent files and extracted directories are part of that owned +//! directory too. Children need not have Maple's private creation permissions. +//! Symlinks are unlinked as entries and never traversed. Nested directories move +//! one level toward the root under generated names so each bounded step makes +//! durable progress without recursion or a persistent in-memory traversal cursor. + +use std::{io, path::Path}; + +use super::{StoreError, StoreResult}; + +const ENTRY_CAP: usize = 128; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(super) struct AttachmentCleanupProgress { + pub done: bool, + pub entries: usize, +} + +pub(super) fn collect_owned_attachment_tree_step( + root: &Path, +) -> StoreResult { + implementation::step(root) +} + +#[cfg(unix)] +mod implementation { + use super::*; + use std::{ + ffi::{CStr, CString}, + fs::File, + os::{ + fd::{AsRawFd, FromRawFd, RawFd}, + unix::ffi::OsStrExt, + }, + path::Component, + }; + + struct Directory(*mut libc::DIR); + impl Drop for Directory { + fn drop(&mut self) { + unsafe { + libc::closedir(self.0); + } + } + } + + fn entries(file: &File) -> StoreResult { + let descriptor = unsafe { libc::dup(file.as_raw_fd()) }; + if descriptor < 0 { + return Err(io::Error::last_os_error().into()); + } + let directory = unsafe { libc::fdopendir(descriptor) }; + if directory.is_null() { + let error = io::Error::last_os_error(); + unsafe { + libc::close(descriptor); + } + return Err(error.into()); + } + Ok(Directory(directory)) + } + + fn next(directory: &Directory) -> StoreResult> { + loop { + unsafe { + *errno() = 0; + } + let entry = unsafe { libc::readdir(directory.0) }; + if entry.is_null() { + let code = unsafe { *errno() }; + return if code == 0 { + Ok(None) + } else { + Err(io::Error::from_raw_os_error(code).into()) + }; + } + let name = unsafe { CStr::from_ptr((*entry).d_name.as_ptr()) }; + if name.to_bytes() == b"." || name.to_bytes() == b".." { + continue; + } + return Ok(Some(name.to_owned())); + } + } + + #[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))] + unsafe fn errno() -> *mut libc::c_int { + unsafe { libc::__error() } + } + #[cfg(not(any(target_os = "macos", target_os = "ios", target_os = "freebsd")))] + unsafe fn errno() -> *mut libc::c_int { + unsafe { libc::__errno_location() } + } + + fn open_at(parent: RawFd, name: &CStr) -> StoreResult { + let fd = unsafe { + libc::openat( + parent, + name.as_ptr(), + libc::O_RDONLY | libc::O_DIRECTORY | libc::O_NOFOLLOW | libc::O_CLOEXEC, + ) + }; + if fd < 0 { + return Err(io::Error::last_os_error().into()); + } + Ok(unsafe { File::from_raw_fd(fd) }) + } + + fn metadata_at(parent: RawFd, name: &CStr) -> StoreResult { + let mut metadata = std::mem::MaybeUninit::::uninit(); + if unsafe { + libc::fstatat( + parent, + name.as_ptr(), + metadata.as_mut_ptr(), + libc::AT_SYMLINK_NOFOLLOW, + ) + } != 0 + { + return Err(io::Error::last_os_error().into()); + } + Ok(unsafe { metadata.assume_init() }) + } + + fn descriptor_metadata(file: &File) -> StoreResult { + let mut opened = std::mem::MaybeUninit::::uninit(); + if unsafe { libc::fstat(file.as_raw_fd(), opened.as_mut_ptr()) } != 0 { + return Err(io::Error::last_os_error().into()); + } + Ok(unsafe { opened.assume_init() }) + } + + fn same_directory(parent: RawFd, name: &CStr, file: &File) -> StoreResult { + let named = metadata_at(parent, name)?; + let opened = descriptor_metadata(file)?; + if named.st_dev != opened.st_dev + || named.st_ino != opened.st_ino + || named.st_mode & libc::S_IFMT != libc::S_IFDIR + { + return Err(StoreError::UnsafePath); + } + Ok(opened) + } + + /// Extraction tools can restore read-only directory modes. Repair only a + /// verified, user-owned directory inside this tree, through its no-follow + /// handle; never chmod an ancestor or link target, and reject other devices. + fn prepare_owned_directory( + parent: RawFd, + name: &CStr, + file: &File, + device: libc::dev_t, + ) -> StoreResult<()> { + let metadata = same_directory(parent, name, file)?; + if metadata.st_uid != unsafe { libc::geteuid() } || metadata.st_dev != device { + return Err(StoreError::UnsafePath); + } + if metadata.st_mode & 0o700 != 0o700 + && unsafe { libc::fchmod(file.as_raw_fd(), (metadata.st_mode & 0o7777) | 0o700) } != 0 + { + return Err(io::Error::last_os_error().into()); + } + Ok(()) + } + + fn unlink(parent: RawFd, name: &CStr, directory: bool) -> StoreResult<()> { + if unsafe { + libc::unlinkat( + parent, + name.as_ptr(), + if directory { libc::AT_REMOVEDIR } else { 0 }, + ) + } != 0 + { + return Err(io::Error::last_os_error().into()); + } + Ok(()) + } + + fn parent_and_name(root: &Path) -> StoreResult<(File, CString)> { + if !root.is_absolute() { + return Err(StoreError::UnsafePath); + } + let mut components = root.components(); + if !matches!(components.next(), Some(Component::RootDir)) { + return Err(StoreError::UnsafePath); + } + let mut names = components + .map(|part| match part { + Component::Normal(name) => { + CString::new(name.as_bytes()).map_err(|_| StoreError::UnsafePath) + } + _ => Err(StoreError::UnsafePath), + }) + .collect::>>()?; + let name = names.pop().ok_or(StoreError::UnsafePath)?; + let mut parent = File::open("/")?; + for name in names { + parent = open_at(parent.as_raw_fd(), &name)?; + } + Ok((parent, name)) + } + + pub(super) fn step(root: &Path) -> StoreResult { + let (parent, name) = match parent_and_name(root) { + Ok(value) => value, + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => { + return Ok(AttachmentCleanupProgress { + done: true, + entries: 0, + }); + } + Err(error) => return Err(error), + }; + let root = match open_at(parent.as_raw_fd(), &name) { + Ok(value) => value, + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => { + parent.sync_all()?; + return Ok(AttachmentCleanupProgress { + done: true, + entries: 0, + }); + } + Err(error) => return Err(error), + }; + let meta = same_directory(parent.as_raw_fd(), &name, &root)?; + if meta.st_uid != unsafe { libc::geteuid() } + || meta.st_dev != descriptor_metadata(&parent)?.st_dev + { + return Err(StoreError::UnsafePath); + } + // This exact owned tree was authorized for deletion. Archive extraction + // may widen its root mode; tighten only that verified no-follow + // descriptor. Namespace ancestors retain their private-path checks. + if meta.st_mode & 0o7777 != 0o700 && unsafe { libc::fchmod(root.as_raw_fd(), 0o700) } != 0 { + return Err(io::Error::last_os_error().into()); + } + same_directory(parent.as_raw_fd(), &name, &root)?; + let mut progress = AttachmentCleanupProgress::default(); + let names = entries(&root)?; + while progress.entries < ENTRY_CAP { + let Some(child_name) = next(&names)? else { + break; + }; + same_directory(parent.as_raw_fd(), &name, &root)?; + let child_meta = metadata_at(root.as_raw_fd(), &child_name)?; + if child_meta.st_mode & libc::S_IFMT == libc::S_IFDIR { + let child = open_at(root.as_raw_fd(), &child_name)?; + prepare_owned_directory(root.as_raw_fd(), &child_name, &child, meta.st_dev)?; + let children = entries(&child)?; + while progress.entries < ENTRY_CAP { + let Some(nested) = next(&children)? else { + break; + }; + same_directory(parent.as_raw_fd(), &name, &root)?; + same_directory(root.as_raw_fd(), &child_name, &child)?; + let nested_meta = metadata_at(child.as_raw_fd(), &nested)?; + if nested_meta.st_mode & libc::S_IFMT == libc::S_IFDIR { + let nested_directory = open_at(child.as_raw_fd(), &nested)?; + prepare_owned_directory( + child.as_raw_fd(), + &nested, + &nested_directory, + meta.st_dev, + )?; + let generated = + CString::new(format!(".cleanup-{:032x}", rand::random::())) + .map_err(|_| StoreError::InvalidData)?; + match metadata_at(root.as_raw_fd(), &generated) { + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => {} + Ok(_) => return Err(StoreError::Conflict), + Err(error) => return Err(error), + } + if unsafe { + libc::renameat( + child.as_raw_fd(), + nested.as_ptr(), + root.as_raw_fd(), + generated.as_ptr(), + ) + } != 0 + { + return Err(io::Error::last_os_error().into()); + } + } else { + unlink(child.as_raw_fd(), &nested, false)?; + } + progress.entries += 1; + } + child.sync_all()?; + if progress.entries < ENTRY_CAP { + same_directory(root.as_raw_fd(), &child_name, &child)?; + match unlink(root.as_raw_fd(), &child_name, true) { + Ok(()) => progress.entries += 1, + Err(StoreError::Io { + kind: io::ErrorKind::DirectoryNotEmpty, + }) => {} + Err(error) => return Err(error), + } + } + } else { + unlink(root.as_raw_fd(), &child_name, false)?; + progress.entries += 1; + } + } + root.sync_all()?; + if progress.entries < ENTRY_CAP { + same_directory(parent.as_raw_fd(), &name, &root)?; + match unlink(parent.as_raw_fd(), &name, true) { + Ok(()) => { + progress.entries += 1; + progress.done = true; + } + Err(StoreError::Io { + kind: io::ErrorKind::DirectoryNotEmpty, + }) => {} + Err(error) => return Err(error), + } + } + parent.sync_all()?; + Ok(progress) + } +} + +#[cfg(windows)] +mod implementation { + use super::*; + use crate::store::paths::{sync_directory, verify_private_path}; + use std::{ + fs::{self, File}, + os::windows::{ + ffi::OsStrExt, + fs::MetadataExt, + io::{AsRawHandle, FromRawHandle}, + }, + path::PathBuf, + }; + use windows_sys::Win32::{ + Foundation::INVALID_HANDLE_VALUE, + Storage::FileSystem::{ + BY_HANDLE_FILE_INFORMATION, CreateFileW, DELETE, FILE_ATTRIBUTE_DIRECTORY, + FILE_ATTRIBUTE_REPARSE_POINT, FILE_DISPOSITION_FLAG_DELETE, + FILE_DISPOSITION_FLAG_IGNORE_READONLY_ATTRIBUTE, FILE_DISPOSITION_INFO_EX, + FILE_FLAG_BACKUP_SEMANTICS, FILE_FLAG_OPEN_REPARSE_POINT, FILE_READ_ATTRIBUTES, + FILE_SHARE_READ, FILE_SHARE_WRITE, FileDispositionInfoEx, GetFileInformationByHandle, + OPEN_EXISTING, SetFileInformationByHandle, + }, + }; + + /// A no-reparse directory handle without DELETE sharing keeps each ancestor + /// and the currently enumerated child from being renamed under path-based + /// std filesystem calls. Descendants may have ordinary inherited ACLs. + fn hold_directory(path: &Path) -> StoreResult { + let path: Vec = path.as_os_str().encode_wide().chain(Some(0)).collect(); + let handle = unsafe { + CreateFileW( + path.as_ptr(), + FILE_READ_ATTRIBUTES, + FILE_SHARE_READ | FILE_SHARE_WRITE, + std::ptr::null(), + OPEN_EXISTING, + FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, + std::ptr::null_mut(), + ) + }; + if handle == INVALID_HANDLE_VALUE { + return Err(io::Error::last_os_error().into()); + } + let file = unsafe { File::from_raw_handle(handle) }; + identity(&file)?; + Ok(file) + } + + fn information(file: &File) -> StoreResult { + let mut info = std::mem::MaybeUninit::::uninit(); + if unsafe { GetFileInformationByHandle(file.as_raw_handle(), info.as_mut_ptr()) } == 0 { + return Err(io::Error::last_os_error().into()); + } + Ok(unsafe { info.assume_init() }) + } + + fn identity(file: &File) -> StoreResult<(u32, u64)> { + let info = information(file)?; + if info.dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT != 0 + || info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY == 0 + { + return Err(StoreError::UnsafePath); + } + Ok(( + info.dwVolumeSerialNumber, + (u64::from(info.nFileIndexHigh) << 32) | u64::from(info.nFileIndexLow), + )) + } + + // Rust's FileType::is_dir excludes directory symlinks/junctions. Windows + // still requires RemoveDirectory for those reparse entries; it removes the + // link itself. Attribute classification never follows the target. + fn directory_attribute(attributes: u32) -> bool { + attributes & FILE_ATTRIBUTE_DIRECTORY != 0 + } + + /// Delete the opened entry without changing its attributes. Clearing + /// READONLY by path could follow a reparse point, and clearing it even by + /// handle would change the attributes of an external hard link too. + fn unlink_entry(path: &Path, directory: bool) -> StoreResult<()> { + // Keep the ordinary unlink behavior when it works; only a denied + // deletion needs the explicit READONLY override below. + let removed = if directory { + fs::remove_dir(path) + } else { + fs::remove_file(path) + }; + match removed { + Ok(()) => return Ok(()), + Err(error) if error.kind() == io::ErrorKind::PermissionDenied => {} + Err(error) => return Err(error.into()), + } + let path: Vec = path.as_os_str().encode_wide().chain(Some(0)).collect(); + let handle = unsafe { + CreateFileW( + path.as_ptr(), + DELETE | FILE_READ_ATTRIBUTES, + FILE_SHARE_READ | FILE_SHARE_WRITE, + std::ptr::null(), + OPEN_EXISTING, + FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, + std::ptr::null_mut(), + ) + }; + if handle == INVALID_HANDLE_VALUE { + return Err(io::Error::last_os_error().into()); + } + let file = unsafe { File::from_raw_handle(handle) }; + if directory_attribute(information(&file)?.dwFileAttributes) != directory { + return Err(StoreError::UnsafePath); + } + let disposition = FILE_DISPOSITION_INFO_EX { + Flags: FILE_DISPOSITION_FLAG_DELETE | FILE_DISPOSITION_FLAG_IGNORE_READONLY_ATTRIBUTE, + }; + if unsafe { + SetFileInformationByHandle( + file.as_raw_handle(), + FileDispositionInfoEx, + (&disposition as *const FILE_DISPOSITION_INFO_EX).cast(), + std::mem::size_of::() as u32, + ) + } == 0 + { + return Err(io::Error::last_os_error().into()); + } + // The disposition takes effect when this no-DELETE-sharing handle is + // closed. The caller still syncs the surviving parent directory. + drop(file); + Ok(()) + } + + fn hold_ancestors(root: &Path) -> StoreResult> { + if !root.is_absolute() { + return Err(StoreError::UnsafePath); + } + let mut paths = root + .ancestors() + .skip(1) + .filter(|path| !path.as_os_str().is_empty()) + .map(Path::to_path_buf) + .collect::>(); + paths.reverse(); + paths.iter().map(|path| hold_directory(path)).collect() + } + + pub(super) fn step(root: &Path) -> StoreResult { + let _ancestors = match hold_ancestors(root) { + Ok(value) => value, + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => { + return Ok(AttachmentCleanupProgress { + done: true, + entries: 0, + }); + } + Err(error) => return Err(error), + }; + match verify_private_path(root, true) { + Ok(()) => {} + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => { + return Ok(AttachmentCleanupProgress { + done: true, + entries: 0, + }); + } + Err(error) => return Err(error), + } + let guard = hold_directory(root)?; + let root_id = identity(&guard)?; + let mut progress = AttachmentCleanupProgress::default(); + for child in fs::read_dir(root)? { + if progress.entries >= ENTRY_CAP { + break; + } + let child = child?; + let metadata = fs::symlink_metadata(child.path())?; + let reparse = metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0; + if metadata.is_dir() && !reparse { + let child_guard = hold_directory(&child.path())?; + for nested in fs::read_dir(child.path())? { + if progress.entries >= ENTRY_CAP { + break; + } + let nested = nested?; + let info = fs::symlink_metadata(nested.path())?; + if info.is_dir() && info.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT == 0 { + let target = root.join(format!(".cleanup-{:032x}", rand::random::())); + match fs::symlink_metadata(&target) { + Err(error) if error.kind() == io::ErrorKind::NotFound => {} + Ok(_) => return Err(StoreError::Conflict), + Err(error) => return Err(error.into()), + } + fs::rename(nested.path(), target)?; + } else if directory_attribute(info.file_attributes()) { + unlink_entry(&nested.path(), true)?; + } else { + unlink_entry(&nested.path(), false)?; + } + progress.entries += 1; + } + let child_id = identity(&child_guard)?; + drop(child_guard); + if progress.entries < ENTRY_CAP { + let verified = hold_directory(&child.path())?; + if identity(&verified)? != child_id { + return Err(StoreError::UnsafePath); + } + drop(verified); + match unlink_entry(&child.path(), true) { + Ok(()) => progress.entries += 1, + Err(StoreError::Io { + kind: io::ErrorKind::DirectoryNotEmpty, + }) => {} + Err(error) => return Err(error), + } + } + } else if directory_attribute(metadata.file_attributes()) { + unlink_entry(&child.path(), true)?; + progress.entries += 1; + } else { + unlink_entry(&child.path(), false)?; + progress.entries += 1; + } + } + if identity(&guard)? != root_id { + return Err(StoreError::UnsafePath); + } + drop(guard); + if progress.entries < ENTRY_CAP { + let verified = hold_directory(root)?; + if identity(&verified)? != root_id { + return Err(StoreError::UnsafePath); + } + drop(verified); + match unlink_entry(root, true) { + Ok(()) => { + progress.entries += 1; + progress.done = true; + } + Err(StoreError::Io { + kind: io::ErrorKind::DirectoryNotEmpty, + }) => {} + Err(error) => return Err(error), + } + } + if let Some(parent) = root.parent() { + sync_directory(parent)?; + } + Ok(progress) + } + + #[cfg(test)] + mod tests { + use super::*; + + #[test] + fn directory_reparse_entries_select_directory_unlink() { + assert!(directory_attribute(FILE_ATTRIBUTE_DIRECTORY)); + assert!(directory_attribute( + FILE_ATTRIBUTE_DIRECTORY | FILE_ATTRIBUTE_REPARSE_POINT + )); + assert!(!directory_attribute(FILE_ATTRIBUTE_REPARSE_POINT)); + assert!(!directory_attribute(0)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::store::paths::private_directory; + use std::fs; + + #[test] + fn removes_read_only_files_in_owned_extractions() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap().join("owned"); + private_directory(&root).unwrap(); + let nested = root.join("extracted"); + fs::create_dir(&nested).unwrap(); + for path in [root.join("original"), nested.join("generated")] { + fs::write(&path, b"read-only attachment content").unwrap(); + let mut permissions = fs::metadata(&path).unwrap().permissions(); + permissions.set_readonly(true); + fs::set_permissions(&path, permissions).unwrap(); + } + while !collect_owned_attachment_tree_step(&root).unwrap().done {} + assert!(!root.exists()); + } + + #[cfg(unix)] + #[test] + fn read_only_owned_directories_are_repaired_without_following_symlinks() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + let directory = tempfile::tempdir().unwrap(); + let directory_path = directory.path().canonicalize().unwrap(); + let root = directory_path.join("owned"); + private_directory(&root).unwrap(); + let nested = root.join("extracted"); + let deeper = nested.join("deeper"); + fs::create_dir_all(&deeper).unwrap(); + for index in 0..300 { + fs::write(deeper.join(format!("file-{index}")), b"extracted content").unwrap(); + } + let outside = directory_path.join("outside"); + fs::create_dir(&outside).unwrap(); + fs::write(outside.join("sentinel"), b"retain external content").unwrap(); + fs::set_permissions(&outside, fs::Permissions::from_mode(0o500)).unwrap(); + symlink(&outside, nested.join("outside-link")).unwrap(); + fs::set_permissions(&deeper, fs::Permissions::from_mode(0o500)).unwrap(); + fs::set_permissions(&nested, fs::Permissions::from_mode(0o555)).unwrap(); + fs::set_permissions(&root, fs::Permissions::from_mode(0o500)).unwrap(); + + let mut completed = false; + for _ in 0..20 { + let progress = collect_owned_attachment_tree_step(&root).unwrap(); + assert!(progress.entries <= ENTRY_CAP); + if progress.done { + completed = true; + break; + } + assert!(progress.entries > 0); + } + assert!(completed); + assert!(!root.exists()); + assert_eq!( + fs::read(outside.join("sentinel")).unwrap(), + b"retain external content" + ); + let outside_mode = fs::metadata(&outside).unwrap().permissions().mode() & 0o777; + // Restore the test-owned target only after observing cleanup's result. + fs::set_permissions(&outside, fs::Permissions::from_mode(0o700)).unwrap(); + assert_eq!(outside_mode, 0o500); + } + + #[cfg(unix)] + #[test] + fn widened_owned_roots_are_tightened_before_bounded_cleanup() { + use std::os::unix::fs::PermissionsExt; + + for mode in [0o755, 0o711, 0o555] { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap().join("owned"); + private_directory(&root).unwrap(); + for index in 0..=ENTRY_CAP { + fs::write(root.join(format!("owned-{index}")), b"owned content").unwrap(); + } + fs::set_permissions(&root, fs::Permissions::from_mode(mode)).unwrap(); + let first = collect_owned_attachment_tree_step(&root).unwrap(); + assert!(!first.done); + assert_eq!(first.entries, ENTRY_CAP); + assert_eq!( + fs::metadata(&root).unwrap().permissions().mode() & 0o7777, + 0o700 + ); + assert!(collect_owned_attachment_tree_step(&root).unwrap().done); + assert!(!root.exists()); + } + } + + #[cfg(windows)] + #[test] + fn read_only_hard_link_cleanup_preserves_external_file_attributes() { + let directory = tempfile::tempdir().unwrap(); + let directory_path = directory.path().canonicalize().unwrap(); + let root = directory_path.join("owned"); + private_directory(&root).unwrap(); + let outside = directory_path.join("outside"); + fs::write(&outside, b"retain external hard link").unwrap(); + fs::hard_link(&outside, root.join("owned-link")).unwrap(); + let mut permissions = fs::metadata(&outside).unwrap().permissions(); + permissions.set_readonly(true); + fs::set_permissions(&outside, permissions).unwrap(); + + let result = collect_owned_attachment_tree_step(&root); + let unchanged = fs::metadata(&outside).unwrap().permissions().readonly(); + // Restore only the test-owned external fixture for TempDir cleanup. + let mut permissions = fs::metadata(&outside).unwrap().permissions(); + permissions.set_readonly(false); + fs::set_permissions(&outside, permissions).unwrap(); + assert!(result.unwrap().done); + assert!(!root.exists()); + assert!(unchanged); + assert_eq!(fs::read(&outside).unwrap(), b"retain external hard link"); + } + + #[test] + fn removes_adjacent_outputs_and_nested_extractions_in_bounded_steps() { + let directory = tempfile::tempdir().unwrap(); + let directory_path = directory.path().canonicalize().unwrap(); + let root = directory_path.join("owned"); + private_directory(&root).unwrap(); + fs::write(root.join(".DS_Store"), b"model adjacent metadata").unwrap(); + let nested = root.join("extracted"); + fs::create_dir(&nested).unwrap(); + let deeper = nested.join("nested"); + fs::create_dir(&deeper).unwrap(); + for index in 0..300 { + fs::write(deeper.join(format!("file-{index}")), b"ordinary user file").unwrap(); + } + let mut calls = 0; + loop { + let progress = collect_owned_attachment_tree_step(&root).unwrap(); + assert!(progress.entries <= ENTRY_CAP); + calls += 1; + if progress.done { + break; + } + assert!( + progress.entries > 0, + "each unfinished step must make progress" + ); + assert!(calls < 20, "cleanup must remain restartable and bounded"); + } + assert!(calls > 1); + assert!(!root.exists()); + assert!(collect_owned_attachment_tree_step(&root).unwrap().done); + } + + #[cfg(unix)] + #[test] + fn deeply_nested_extractions_make_restartable_progress_without_recursion() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap().join("owned"); + private_directory(&root).unwrap(); + let mut nested = root.clone(); + for _ in 0..96 { + nested = nested.join("d"); + fs::create_dir(&nested).unwrap(); + } + fs::write(nested.join("payload"), b"deep model output").unwrap(); + for _ in 0..100 { + let progress = collect_owned_attachment_tree_step(&root).unwrap(); + assert!(progress.entries <= ENTRY_CAP); + if progress.done { + assert!(!root.exists()); + return; + } + assert!(progress.entries > 0); + } + panic!("bounded cleanup failed to finish a deep tree"); + } + + #[cfg(unix)] + #[test] + fn unlinks_nested_symlinks_without_following_file_or_directory_targets() { + use std::os::unix::fs::symlink; + let directory = tempfile::tempdir().unwrap(); + let directory_path = directory.path().canonicalize().unwrap(); + let root = directory_path.join("owned"); + private_directory(&root).unwrap(); + let nested = root.join("extracted"); + fs::create_dir(&nested).unwrap(); + let outside = directory_path.join("outside"); + fs::create_dir(&outside).unwrap(); + fs::write(outside.join("sentinel"), b"never delete external content").unwrap(); + symlink(&outside, nested.join("directory-link")).unwrap(); + symlink(outside.join("sentinel"), root.join("file-link")).unwrap(); + while !collect_owned_attachment_tree_step(&root).unwrap().done {} + assert_eq!( + fs::read(outside.join("sentinel")).unwrap(), + b"never delete external content" + ); + } + + #[cfg(unix)] + #[test] + fn root_and_ancestor_symlinks_never_authorize_cleanup() { + use std::os::unix::fs::symlink; + let directory = tempfile::tempdir().unwrap(); + let directory_path = directory.path().canonicalize().unwrap(); + let outside = directory_path.join("outside"); + private_directory(&outside).unwrap(); + let root = outside.join("owned"); + private_directory(&root).unwrap(); + fs::write(root.join("sentinel"), b"leave this path intact").unwrap(); + let alias = directory_path.join("alias"); + symlink(&outside, &alias).unwrap(); + assert!(collect_owned_attachment_tree_step(&alias.join("owned")).is_err()); + let root_alias = directory_path.join("root-alias"); + symlink(&root, &root_alias).unwrap(); + assert!(collect_owned_attachment_tree_step(&root_alias).is_err()); + assert_eq!( + fs::read(root.join("sentinel")).unwrap(), + b"leave this path intact" + ); + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/attachment_migration.rs b/apps/maple-agent/crates/maple-agent/src/store/attachment_migration.rs new file mode 100644 index 000000000..059730ff5 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/attachment_migration.rs @@ -0,0 +1,1100 @@ +//! One-way retirement of the old SQLite image payloads. +//! +//! Schema v3 fences old binaries before this code changes any source. Every +//! step commits its progress with its effect, so a crash resumes at that step. +//! Nothing exports formerly protected bytes to a plaintext attachment file. + +use pi_ai::{Content, Message}; +use pi_coding_agent::{ + SessionMessage, + session::{EntryKind, SessionEntry}, +}; +use rusqlite::{Connection, TransactionBehavior, params}; + +use super::{DamageReason, StoreError, StoreResult, codec}; + +const PHASE_KEY: &str = "attachment_migration_phase"; +const CURSOR_KEY: &str = "attachment_migration_cursor"; +const RETIRED_IMAGE: &str = "maple-attachment-unavailable:retired"; +const ROW_BATCH: i64 = 128; +const CHUNK_BATCH: i64 = 4; +const SOURCE_BATCH_BYTES: usize = 4 * 1024 * 1024; + +pub(super) fn retire_legacy_images(connection: &mut Connection) -> StoreResult<()> { + let mut steps = 0usize; + while !retire_step(connection)? { + steps += 1; + if steps.is_multiple_of(16) { + // The ordinary checkpoint worker starts after migration. Keep a + // large retirement from retaining its entire rewrite in the WAL. + super::connection::checkpoint(connection, false)?; + } + } + super::connection::checkpoint(connection, false)?; + Ok(()) +} + +/// Each call is one short, durable SQL transaction. Source decoding is bounded +/// by the codec's existing per-entry cap, never the total history or blob size. +fn retire_step(connection: &mut Connection) -> StoreResult { + let phase: i64 = connection.query_row( + "SELECT value FROM store_meta WHERE key=?1", + [PHASE_KEY], + |row| row.get(0), + )?; + match phase { + 0 => retire_source_batch(connection), + 1 => invalidate_task_timelines(connection), + 2 => collect_payload_rows( + connection, + "timeline_rows", + 3, + ROW_BATCH, + SOURCE_BATCH_BYTES, + ), + 3 => collect_rows(connection, "attachments", 4, ROW_BATCH), + 4 => collect_rows(connection, "task_blobs", 5, ROW_BATCH), + 5 => collect_payload_rows(connection, "blob_chunks", 6, CHUNK_BATCH, 1024 * 1024), + 6 => collect_rows(connection, "blobs", 7, ROW_BATCH), + 7 => Ok(true), + _ => Err(StoreError::InvalidData), + } +} + +fn retire_source_batch(connection: &mut Connection) -> StoreResult { + let cursor: i64 = connection.query_row( + "SELECT value FROM store_meta WHERE key=?1", + [CURSOR_KEY], + |row| row.get(0), + )?; + if cursor < 0 + || connection.query_row( + "SELECT EXISTS(SELECT 1 FROM entries WHERE id<=0)", + [], + |row| row.get::<_, bool>(0), + )? + { + return Err(StoreError::InvalidData); + } + let mut prepared = Vec::new(); + let mut batch_bytes = 0usize; + { + let mut statement = connection.prepare( + "SELECT id,entry_id,flags,codec,bytes,length(payload), \ + CASE WHEN length(payload)<=?3 THEN payload END \ + FROM entries WHERE id>?1 ORDER BY id LIMIT ?2", + )?; + let mut rows = + statement.query(params![cursor, ROW_BATCH, codec::MAX_DECODED_BYTES as i64])?; + while let Some(row) = rows.next()? { + let id: i64 = row.get(0)?; + let entry_id: String = row.get(1)?; + let flags: i64 = row.get(2)?; + let encoded: i64 = row.get(3)?; + let size: i64 = row.get(4)?; + let actual_size: i64 = row.get(5)?; + if flags & !1 != 0 || size < 0 || size != actual_size { + return Err(damaged_source()); + } + let size = usize::try_from(size).map_err(|_| damaged_source())?; + if size > codec::MAX_DECODED_BYTES { + return Err(damaged_source()); + } + if !prepared.is_empty() && batch_bytes.saturating_add(size) > SOURCE_BATCH_BYTES { + break; + } + let payload: Vec = row.get(6)?; + let decoded = codec::decode(encoded, &payload).map_err(|error| match error { + StoreError::InvalidData => damaged_source(), + other => other, + })?; + let mut entry: SessionEntry = + serde_json::from_slice(&decoded).map_err(|_| damaged_source())?; + if entry.id != entry_id { + return Err(damaged_source()); + } + let replacement = if retire_entry_images(&mut entry) { + Some(codec::encode( + &serde_json::to_vec(&entry).map_err(|_| StoreError::InvalidData)?, + )?) + } else { + None + }; + prepared.push((id, flags, replacement)); + batch_bytes += size; + } + } + let Some((last_id, _, _)) = prepared.last() else { + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + transaction.execute("UPDATE store_meta SET value=0 WHERE key=?1", [CURSOR_KEY])?; + set_phase(&transaction, 1)?; + transaction.commit()?; + return Ok(false); + }; + let last_id = *last_id; + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + for (id, flags, replacement) in prepared { + if let Some((encoded, payload)) = replacement { + transaction.execute( + "UPDATE entries SET flags=0,codec=?2,bytes=?3,payload=?4 WHERE id=?1", + params![id, encoded, payload.len() as i64, payload], + )?; + } else if flags != 0 { + transaction.execute("UPDATE entries SET flags=0 WHERE id=?1", [id])?; + } + } + transaction.execute( + "UPDATE store_meta SET value=?2 WHERE key=?1", + params![CURSOR_KEY, last_id], + )?; + transaction.commit()?; + Ok(false) +} + +/// A verified invalid source record is damaged data, not an unavailable or +/// foreign store. Keep the migration fence and permit the existing explicit +/// recovery/reset actions; never retire arbitrary corrupt data as an image. +fn damaged_source() -> StoreError { + StoreError::Damaged { + reason: DamageReason::Database, + } +} + +fn invalidate_task_timelines(connection: &mut Connection) -> StoreResult { + let cursor: i64 = connection.query_row( + "SELECT value FROM store_meta WHERE key=?1", + [CURSOR_KEY], + |row| row.get(0), + )?; + let end: Option = connection.query_row( + "SELECT max(id) FROM (SELECT id FROM tasks WHERE id>?1 ORDER BY id LIMIT ?2)", + params![cursor, ROW_BATCH], + |row| row.get(0), + )?; + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + if let Some(end) = end { + // Old rows can contain cached attachment metadata. Rebuild from the + // rewritten source while retaining source identities and tool summaries. + transaction.execute( + "UPDATE tasks SET timeline_format=0,rebuild_epoch=NULL,rebuild_seq=NULL,rebuild_item_index=0,context_usage=NULL \ + WHERE id>?1 AND id<=?2", + params![cursor, end], + )?; + transaction.execute( + "UPDATE store_meta SET value=?2 WHERE key=?1", + params![CURSOR_KEY, end], + )?; + } else { + set_phase(&transaction, 2)?; + } + transaction.commit()?; + Ok(false) +} + +fn retire_entry_images(entry: &mut SessionEntry) -> bool { + let content = match &mut entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::User(message)), + } => Some(&mut message.content), + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(message)), + } => Some(&mut message.content), + EntryKind::Message { + message: SessionMessage::Custom(message), + } => Some(&mut message.content), + EntryKind::CustomMessage { content, .. } => Some(content), + EntryKind::ContextEdit { replacement, .. } => replacement.as_mut(), + _ => None, + }; + let mut changed = false; + if let Some(content) = content { + for block in content { + if let Content::Image(image) = block { + changed |= image.data != RETIRED_IMAGE; + image.data = RETIRED_IMAGE.into(); + } + } + } + changed +} + +fn collect_rows( + connection: &mut Connection, + table: &str, + next_phase: i64, + limit: i64, +) -> StoreResult { + // Names are internal constants above, never caller input. WITHOUT ROWID + // tables need their actual primary keys for bounded deletion. + let key = match table { + "blobs" => "id", + "attachments" => "(task_id,public_id)", + "task_blobs" => "(task_id,blob_id)", + _ => return Err(StoreError::InvalidData), + }; + let select = key.trim_start_matches('(').trim_end_matches(')'); + let sql = format!("DELETE FROM {table} WHERE {key} IN (SELECT {select} FROM {table} LIMIT ?1)"); + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + if transaction.execute(&sql, [limit])? == 0 { + set_phase(&transaction, next_phase)?; + } + transaction.commit()?; + Ok(false) +} + +fn collect_payload_rows( + connection: &mut Connection, + table: &str, + next_phase: i64, + limit: i64, + byte_limit: usize, +) -> StoreResult { + // Inspect only primary keys and actual BLOB lengths, never payload bytes or + // the cached `bytes` field. One oversized historical row is retired alone + // so malformed old chunk sizes cannot inflate the rest of a batch. + let (select, delete) = match table { + "timeline_rows" => ( + "SELECT id,0,length(item) FROM timeline_rows ORDER BY id LIMIT ?1", + "DELETE FROM timeline_rows WHERE id=?1 AND ?2=0", + ), + "blob_chunks" => ( + "SELECT blob_id,idx,length(data) FROM blob_chunks ORDER BY blob_id,idx LIMIT ?1", + "DELETE FROM blob_chunks WHERE blob_id=?1 AND idx=?2", + ), + _ => return Err(StoreError::InvalidData), + }; + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + let mut keys = Vec::new(); + let mut bytes = 0usize; + { + let mut statement = transaction.prepare(select)?; + let mut rows = statement.query([limit])?; + while let Some(row) = rows.next()? { + let size = + usize::try_from(row.get::<_, i64>(2)?).map_err(|_| StoreError::InvalidData)?; + if !keys.is_empty() && bytes.saturating_add(size) > byte_limit { + break; + } + keys.push((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)); + bytes = bytes.saturating_add(size); + if bytes >= byte_limit { + break; + } + } + } + if keys.is_empty() { + set_phase(&transaction, next_phase)?; + } else { + for key in keys { + transaction.execute(delete, params![key.0, key.1])?; + } + } + transaction.commit()?; + Ok(false) +} + +fn set_phase(connection: &Connection, phase: i64) -> StoreResult<()> { + if connection.execute( + "UPDATE store_meta SET value=?2 WHERE key=?1", + params![PHASE_KEY, phase], + )? != 1 + { + return Err(StoreError::InvalidData); + } + Ok(()) +} + +#[cfg(test)] +impl super::Store { + /// Reproduce the derived-history state after retirement without changing + /// source records, for runtime admission regressions. + pub(crate) fn retire_task_timeline_for_test( + &self, + account: super::AccountId, + task: &str, + ) -> StoreResult<()> { + self.write(super::Durability::Normal, |tx| { + let key = super::tasks::task_key_tx(tx, account, task)?; + tx.execute("DELETE FROM timeline_rows WHERE task_id=?1", [key.0])?; + tx.execute("UPDATE tasks SET timeline_format=0,rebuild_epoch=NULL,rebuild_seq=NULL,rebuild_item_index=0 WHERE id=?1", [key.0])?; + Ok(()) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::store::{ + AccountId, OpenTaskPageRequest, PageTransport, SessionLoad, Store, StorePaths, StoreState, + engine, keys, schema, + }; + use pi_ai::{ImageContent, ToolResultMessage, UserMessage}; + use pi_coding_agent::messages::CustomMessage; + use pi_coding_agent::session::{SESSION_VERSION, SessionHeader}; + use std::io::{Read, Write}; + use std::sync::{Arc, Mutex}; + + const TASK: &str = "00000000-0000-7000-8000-000000000001"; + const TEXT_TASK: &str = "00000000-0000-7000-8000-000000000002"; + + fn connection(path: &std::path::Path, keyed: bool) -> Connection { + let connection = Connection::open(path).unwrap(); + if keyed { + engine::key_connection( + &connection, + &keys::LockedKey::from_bytes([0xa5; 32]).unwrap(), + ) + .unwrap(); + } else { + engine::guard_plain(&connection).unwrap(); + } + connection + .pragma_update(None, "foreign_keys", true) + .unwrap(); + connection + } + + fn image(data: &str) -> Content { + Content::Image(ImageContent { + data: data.into(), + mime_type: "image/png".into(), + }) + } + + fn project_root() -> &'static str { + if cfg!(windows) { "C:/work" } else { "/work" } + } + + fn fixture(connection: &mut Connection, version: i64) -> Vec { + schema::migrate_to_version(connection, 1, version).unwrap(); + connection.execute("INSERT INTO accounts(id,user_key,user_id,created_ms,last_seen_ms) VALUES (1,'a','A',1,1)", []).unwrap(); + connection + .execute( + "INSERT INTO projects(id,account_id,root,last_used_ms) VALUES (1,1,?1,1)", + [project_root()], + ) + .unwrap(); + connection.execute("INSERT INTO tasks(id,uuid,account_id,project_id,kind,state,title,title_fold,created_ms,updated_ms,last_seq,timeline_format,rows_epoch,last_row_seq,pi_header,message_count) VALUES (1,?1,1,1,'desktop','settled','Keep title','keep title',1,1,130,1,3,1,?2,1)", params![TASK, header(TASK)]).unwrap(); + connection + .execute("UPDATE tasks SET context_usage=1234", []) + .unwrap(); + let mut kinds = vec![ + EntryKind::Message { + message: SessionMessage::Llm(Message::User(UserMessage { + content: vec![Content::text("keep user text"), image("maple-blob:123")], + timestamp: 100, + })), + }, + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(ToolResultMessage { + tool_call_id: "call-1".into(), + tool_name: "screenshot".into(), + content: vec![ + Content::text("keep OCR"), + image("maple-inline-image:opaque"), + ], + details: Some(serde_json::json!({"keep":"data:image/png;base64,opaque"})), + usage: None, + is_error: false, + timestamp: 101, + })), + }, + EntryKind::Message { + message: SessionMessage::Custom(CustomMessage { + custom_type: "extension".into(), + content: vec![image("raw-base64-image")], + display: true, + details: None, + timestamp: 102, + }), + }, + EntryKind::CustomMessage { + custom_type: "custom".into(), + content: vec![ + Content::text("keep custom text"), + image("data:image/png;base64,AA=="), + ], + display: true, + details: Some(serde_json::json!({"image":"opaque extension data stays"})), + }, + EntryKind::ContextEdit { + target_id: "entry0".into(), + replacement: Some(vec![ + Content::text("keep replacement"), + image("maple-blob:456"), + ]), + }, + EntryKind::Compaction { + summary: "keep summary".into(), + first_kept_entry_id: "entry1".into(), + tokens_before: 100, + details: None, + usage: None, + from_extension: false, + system_message: None, + }, + ]; + // More than a source batch proves its cursor resumes a partial history. + kinds.extend((6..130).map(|index| EntryKind::Label { + target_id: "entry0".into(), + label: Some(format!("preserved label {index}")), + })); + let entries = kinds + .into_iter() + .enumerate() + .map(|(index, kind)| SessionEntry { + id: format!("entry{index}"), + parent_id: index.checked_sub(1).map(|parent| format!("entry{parent}")), + timestamp: 100 + index as i64, + kind, + }) + .collect::>(); + for (index, entry) in entries.iter().enumerate() { + let payload = serde_json::to_vec(entry).unwrap(); + // Exercise both codecs even for this small typed fixture. + let (encoded, payload) = if index % 2 == 0 { + (1, zstd::stream::encode_all(payload.as_slice(), 3).unwrap()) + } else { + (0, payload) + }; + connection.execute("INSERT INTO entries(task_id,seq,entry_id,kind,created_ms,flags,codec,bytes,payload) VALUES (1,?1,?2,'message',1,?3,?4,?5,?6)", params![index as i64 + 1, entry.id, i64::from(index < 2), encoded, payload.len() as i64, payload]).unwrap(); + } + connection.execute("INSERT INTO blobs(id,sha256,mime,size,complete,created_ms) VALUES (1,?1,'image/png',5120,1,1)", [vec![1u8;32]]).unwrap(); + for index in 0..5 { + connection + .execute( + "INSERT INTO blob_chunks(blob_id,idx,data) VALUES (1,?1,?2)", + params![index, vec![42u8; 1024]], + ) + .unwrap(); + } + connection + .execute("INSERT INTO task_blobs(task_id,blob_id) VALUES (1,1)", []) + .unwrap(); + connection.execute("INSERT INTO attachments(task_id,public_id,blob_id,created_ms) VALUES (1,'old-image',1,1)", []).unwrap(); + connection.execute("INSERT INTO timeline_rows(task_id,epoch,row_seq,row_key,entry_seq,item_index,update_seq,bytes,item) VALUES (1,3,1,'row',1,0,1,3,?1)", [b"old".as_slice()]).unwrap(); + connection.execute("INSERT INTO tool_summaries(task_id,entry_seq,item_index,bytes,summary) VALUES (1,2,0,12,'keep summary')", []).unwrap(); + entries + } + + fn header(task: &str) -> String { + serde_json::to_string(&SessionHeader { + version: SESSION_VERSION, + id: task.to_owned(), + timestamp: 1, + cwd: project_root().into(), + parent_session: None, + }) + .unwrap() + } + + /// The expected fixture is independent of the retirement visitor: only the + /// five declared typed-image fields change. Opaque details and all source + /// identities, text, tree edges, statuses and timestamps remain identical. + fn expected_retired_entries(original: &[SessionEntry]) -> Vec { + let mut values = original + .iter() + .map(|entry| serde_json::to_value(entry).unwrap()) + .collect::>(); + for (index, path) in [ + (0, "/message/content/1/data"), + (1, "/message/content/1/data"), + (2, "/message/content/0/data"), + (3, "/content/1/data"), + (4, "/replacement/1/data"), + ] { + *values[index].pointer_mut(path).unwrap() = + serde_json::json!("maple-attachment-unavailable:retired"); + } + values + .into_iter() + .map(|value| serde_json::from_value(value).unwrap()) + .collect() + } + + #[test] + fn plain_and_keyed_retirement_resumes_preserves_source_and_removes_sql_bytes() { + for keyed in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("migration.db"); + let mut database = connection(&path, keyed); + let original = fixture(&mut database, 3); + // Stop after one source batch: schema v3 and the cursor both survive + // reopening. Released v2 readers reject this user_version first. + assert!(!retire_step(&mut database).unwrap()); + assert_eq!( + database + .pragma_query_value::(None, "user_version", |row| row.get(0)) + .unwrap(), + 3 + ); + database.close().unwrap(); + let mut database = connection(&path, keyed); + assert_eq!(schema::preflight(&database, false).unwrap(), 3); + assert_eq!( + database + .query_row::( + "SELECT value FROM store_meta WHERE key=?1", + [CURSOR_KEY], + |row| row.get(0) + ) + .unwrap(), + 128 + ); + // Step until the bounded chunk phase, then prove a step retires at + // most four chunks and resumes the remainder after another reopen. + while database + .query_row::( + "SELECT value FROM store_meta WHERE key=?1", + [PHASE_KEY], + |row| row.get(0), + ) + .unwrap() + != 5 + { + assert!(!retire_step(&mut database).unwrap()); + } + assert!(!retire_step(&mut database).unwrap()); + assert_eq!( + database + .query_row::("SELECT count(*) FROM blob_chunks", [], |row| row + .get(0)) + .unwrap(), + 1 + ); + database.close().unwrap(); + let mut database = connection(&path, keyed); + schema::migrate(&mut database, 200).unwrap(); + schema::validate(&database).unwrap(); + let mut statement = database + .prepare("SELECT flags,codec,bytes,payload FROM entries ORDER BY seq") + .unwrap(); + let stored = statement + .query_map([], |row| { + Ok(( + row.get::<_, i64>(0)?, + row.get::<_, i64>(1)?, + row.get::<_, i64>(2)?, + row.get::<_, Vec>(3)?, + )) + }) + .unwrap() + .collect::>>() + .unwrap(); + for (expected, (flags, encoded, bytes, payload)) in + expected_retired_entries(&original).into_iter().zip(stored) + { + assert_eq!(flags, 0); + assert_eq!(bytes as usize, payload.len()); + let actual = crate::store::session::decode_entry(encoded, &payload).unwrap(); + assert_eq!(actual, expected); + } + for table in ["blobs", "blob_chunks", "task_blobs", "attachments"] { + assert_eq!( + database + .query_row::( + "SELECT count(*) FROM sqlite_schema WHERE name=?1", + [table], + |row| row.get(0) + ) + .unwrap(), + 0 + ); + } + assert_eq!( + database + .query_row::("SELECT summary FROM tool_summaries", [], |row| row + .get(0)) + .unwrap(), + "keep summary" + ); + assert_eq!( + database + .query_row::("SELECT count(*) FROM timeline_rows", [], |row| row + .get(0)) + .unwrap(), + 0 + ); + assert_eq!( + database + .query_row::<(String, i64, i64), _, _>( + "SELECT title,last_seq,timeline_format FROM tasks", + [], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)) + ) + .unwrap(), + ("Keep title".into(), 130, 0) + ); + assert_eq!( + database + .query_row::, _, _>("SELECT context_usage FROM tasks", [], |row| { + row.get(0) + }) + .unwrap(), + None + ); + assert!(!temp.path().join("attachments").exists()); + } + } + + fn add_text_task(database: &Connection) { + database.execute("INSERT INTO tasks(id,uuid,account_id,project_id,kind,state,title,title_fold,created_ms,updated_ms,last_seq,timeline_format,rows_epoch,last_row_seq,pi_header,message_count) VALUES (2,?1,1,1,'desktop','settled','Plain task','plain task',1,1,1,1,3,1,?2,1)", params![TEXT_TASK, header(TEXT_TASK)]).unwrap(); + let entry = SessionEntry { + id: "00000001".into(), + parent_id: None, + timestamp: 1, + kind: EntryKind::Message { + message: SessionMessage::Llm(Message::User(UserMessage { + content: vec![Content::text("plain conversation survives")], + timestamp: 1, + })), + }, + }; + let payload = serde_json::to_vec(&entry).unwrap(); + database.execute("INSERT INTO entries(task_id,seq,entry_id,kind,created_ms,flags,codec,bytes,payload) VALUES (2,1,'00000001','message',1,0,0,?1,?2)", params![payload.len() as i64, payload]).unwrap(); + } + + /// Build an on-disk released store, including its private generation and + /// password slot. Startup/unlock below must use the production lifecycle. + fn released_store( + root: &std::path::Path, + keyed: bool, + corrupt: bool, + ) -> (StorePaths, Option, Vec) { + let paths = StorePaths::new(root); + paths.prepare().unwrap(); + paths.precreate_database().unwrap(); + let mut database = connection(&paths.database(), keyed); + let original = fixture(&mut database, 2); + add_text_task(&database); + let store_id: Vec = database + .query_row( + "SELECT value FROM store_meta WHERE key='store_id'", + [], + |row| row.get(0), + ) + .unwrap(); + let mut generation = super::super::paths::private_new_file(&paths.generation()).unwrap(); + generation + .write_all( + serde_json::to_string(&serde_json::json!({ + "format": "maple-store-generation/1", + "store_id": store_id.iter().map(|byte| format!("{byte:02x}")).collect::(), + })) + .unwrap() + .as_bytes(), + ) + .unwrap(); + generation.sync_all().unwrap(); + if corrupt { + database + .execute( + "UPDATE entries SET codec=0,bytes=1,payload=x'7b' WHERE id=1", + [], + ) + .unwrap(); + } + database.close().unwrap(); + let password = keyed.then(|| { + let password = keys::Password::normalize("Z7jv!9qP#3nMx@8Ks^2dH").unwrap(); + let mut salt = [0; 16]; + std::fs::File::open(paths.database()) + .unwrap() + .read_exact(&mut salt) + .unwrap(); + let keyfile = keys::KeyFile::create( + &password, + &[], + &keys::LockedKey::from_bytes([0xa5; 32]).unwrap(), + salt, + keys::KdfParams::floor(), + ) + .unwrap(); + keys::KeyFilePair::write(&paths.db(), &keyfile).unwrap(); + password + }); + (paths, password, original) + } + + #[derive(Default)] + struct HistoryEvents(Mutex>); + + impl crate::agent::AgentEventSink for HistoryEvents { + fn emit(&self, event: &crate::agent::AgentServiceEvent) { + self.0.lock().unwrap().push(event.clone()); + } + } + + #[test] + fn startup_retirement_opens_plain_and_keyed_history_on_demand() { + use crate::agent::{AgentPathLayout, MapleAgentHostResources, MapleAgentService}; + + for keyed in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let (paths, password, original) = released_store(temp.path(), keyed, false); + let store = Arc::new(Store::unavailable(paths, StoreState::Maintenance).unwrap()); + store.initialize_recovered().unwrap(); + if let Some(password) = password.as_ref() { + assert_eq!(store.state(), StoreState::Locked); + store.unlock_store(password, None).unwrap(); + } + assert_eq!(store.state(), StoreState::Open); + let loaded = store + .load_session(AccountId(1), TASK, SessionLoad::Full) + .unwrap(); + assert_eq!(loaded.entries, expected_retired_entries(&original)); + assert_eq!(loaded.header.unwrap().id, TASK); + assert!(store.history_needs_rebuild(AccountId(1), TASK).unwrap()); + assert!( + store + .history_needs_rebuild(AccountId(1), TEXT_TASK) + .unwrap() + ); + let request = OpenTaskPageRequest { + task: TEXT_TASK.into(), + limit_rows: None, + limit_bytes: None, + }; + assert!(matches!( + store.open_task_page(AccountId(1), &request, PageTransport::Local), + Err(StoreError::Busy) + )); + + // A deadline expires between bounded source batches, leaves the old + // epoch unexposed, and permits the foreground open to resume work. + let ticket = store + .begin_timeline_rebuild(AccountId(1), TEXT_TASK, false) + .unwrap(); + assert!(matches!( + store.prepare_timeline_rebuild_before(&ticket, Some(std::time::Instant::now())), + Err(StoreError::Busy) + )); + let events = Arc::new(HistoryEvents::default()); + // Construct outside Tokio so no background maintenance worker can + // make this pass before the production foreground-open path runs. + let service = MapleAgentService::new( + MapleAgentHostResources::new( + AgentPathLayout::from_resource_root(temp.path().join("resources")) + .with_home(Some(temp.path().join("home"))), + events.clone(), + String::new(), + ) + .with_store(Arc::clone(&store)), + ); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let handle = service.handle_for_user("A").await.unwrap(); + let page = handle + .open_task_page(request, PageTransport::Local) + .await + .unwrap(); + assert_eq!( + page.timeline_format, + super::super::projector::TIMELINE_FORMAT + ); + assert!(!page.rows.is_empty()); + assert!( + serde_json::to_string(&page.rows) + .unwrap() + .contains("plain conversation survives") + ); + assert!( + !store + .history_needs_rebuild(AccountId(1), TEXT_TASK) + .unwrap() + ); + let stamp = store.task_history_stamp(AccountId(1), TEXT_TASK).unwrap(); + let snapshot = store + .capture_replay( + AccountId(1), + TEXT_TASK, + stamp, + Arc::new(std::sync::atomic::AtomicBool::new(false)), + std::time::Duration::from_secs(5), + ) + .unwrap(); + let mut replay = Vec::new(); + while let Some((item, _)) = snapshot.next_row().unwrap() { + replay.push(item); + } + assert!( + serde_json::to_string(&replay) + .unwrap() + .contains("plain conversation survives") + ); + // Opening one task does not synchronously rebuild every task. + assert!(store.history_needs_rebuild(AccountId(1), TASK).unwrap()); + // ACP calls this same readiness boundary before its immutable + // replay; remote paging uses the same prepared rows and caps. + handle.ensure_task_history_ready(TASK).await.unwrap(); + let remote = handle + .open_task_page( + OpenTaskPageRequest { + task: TASK.into(), + limit_rows: None, + limit_bytes: None, + }, + PageTransport::Remote, + ) + .await + .unwrap(); + assert!(!remote.rows.is_empty()); + assert!( + !serde_json::to_string(&remote) + .unwrap() + .contains("maple-blob:") + ); + }); + assert!(events.0.lock().unwrap().iter().any(|event| matches!( + event, + crate::agent::AgentServiceEvent::HistoryReplaced { session_id, .. } if session_id == TEXT_TASK + ))); + } + } + + #[test] + fn corrupt_retirement_is_damaged_and_allows_explicit_reset() { + use crate::store::lifecycle::{ERASE_CONFIRMATION, LifecycleError}; + + for keyed in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let (paths, password, _) = released_store(temp.path(), keyed, true); + let old_generation: serde_json::Value = + serde_json::from_slice(&std::fs::read(paths.generation()).unwrap()).unwrap(); + let store = Store::unavailable(paths.clone(), StoreState::Maintenance).unwrap(); + if let Some(password) = password.as_ref() { + store.initialize_recovered().unwrap(); + assert_eq!(store.state(), StoreState::Locked); + assert!(matches!( + store.unlock_store(password, None), + Err(LifecycleError::Store(StoreError::Damaged { + reason: DamageReason::Database + })) + )); + } else { + assert_eq!(store.initialize_recovered(), Err(damaged_source())); + } + assert_eq!( + store.state(), + StoreState::Damaged { + reason: DamageReason::Database + } + ); + assert_eq!( + store.erase_local_data(""), + Err(LifecycleError::Confirmation) + ); + assert_eq!( + store.state(), + StoreState::Damaged { + reason: DamageReason::Database + } + ); + store.erase_local_data(ERASE_CONFIRMATION).unwrap(); + if keyed { + assert_eq!(store.state(), StoreState::NeedsSetup); + store + .setup_store(None, &[], keys::KdfParams::floor()) + .unwrap(); + } + // Plain reset immediately creates a plain replacement. Resetting + // a password store waits for the user's explicit new-store choice. + assert_eq!(store.state(), StoreState::Open); + let replacement_id = store.store_id().unwrap(); + let replacement_hex = replacement_id + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + assert_ne!( + old_generation["store_id"].as_str().unwrap(), + replacement_hex + ); + assert!(store.task(AccountId(1), TASK).unwrap().is_none()); + drop(store); + let reopened = Store::unavailable(paths, StoreState::Maintenance).unwrap(); + reopened.initialize_recovered().unwrap(); + assert_eq!(reopened.state(), StoreState::Open); + assert_eq!(reopened.store_id().unwrap(), replacement_id); + assert!(reopened.task(AccountId(1), TASK).unwrap().is_none()); + } + } + + #[test] + fn malformed_source_does_not_become_an_unavailable_image() { + for corruption in [ + "flags", + "size", + "codec", + "source", + "entry_id", + "later_source", + ] { + let temp = tempfile::tempdir().unwrap(); + let mut database = connection(&temp.path().join("migration.db"), false); + // Enter from the shipped schema: even an invalid first source row + // must commit the version fence before any retirement is attempted. + fixture(&mut database, 2); + let sql = match corruption { + "flags" => "UPDATE entries SET flags=2 WHERE id=1", + "size" => "UPDATE entries SET bytes=bytes+1 WHERE id=1", + "codec" => "UPDATE entries SET codec=2 WHERE id=1", + "source" => "UPDATE entries SET codec=0,bytes=1,payload=x'7b' WHERE id=1", + "entry_id" => "UPDATE entries SET entry_id='wrong' WHERE id=1", + "later_source" => "UPDATE entries SET codec=2 WHERE id=2", + _ => unreachable!(), + }; + database.execute(sql, []).unwrap(); + assert_eq!( + schema::migrate(&mut database, 200), + Err(damaged_source()), + "{corruption}" + ); + assert_eq!( + database + .pragma_query_value::(None, "user_version", |row| row.get(0)) + .unwrap(), + 3 + ); + assert_eq!( + database + .query_row::("SELECT count(*) FROM blob_chunks", [], |row| row + .get(0)) + .unwrap(), + 5 + ); + assert_eq!( + database + .query_row::( + "SELECT value FROM store_meta WHERE key=?1", + [CURSOR_KEY], + |row| row.get(0) + ) + .unwrap(), + 0 + ); + } + } + + #[test] + fn collection_uses_actual_payload_sizes_and_retires_oversized_rows_alone() { + for keyed in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let mut database = connection(&temp.path().join("migration.db"), keyed); + fixture(&mut database, 3); + database.execute("DELETE FROM timeline_rows", []).unwrap(); + for (id, size) in [(1, 3 * 1024 * 1024), (2, 5 * 1024 * 1024), (3, 1024 * 1024)] { + // Intentionally wrong cached bytes: retirement must use length(item). + database.execute("INSERT INTO timeline_rows(id,task_id,epoch,row_seq,row_key,entry_seq,item_index,update_seq,bytes,item) VALUES (?1,1,3,?1,'row',1,0,1,1,zeroblob(?2))", params![id, size]).unwrap(); + } + set_phase(&database, 2).unwrap(); + for expected in [2, 1, 0] { + assert!(!retire_step(&mut database).unwrap()); + assert_eq!( + database + .query_row::("SELECT count(*) FROM timeline_rows", [], |row| row + .get(0)) + .unwrap(), + expected + ); + } + database.execute("DELETE FROM blob_chunks", []).unwrap(); + for (index, size) in [(0, 2 * 1024 * 1024), (1, 256 * 1024), (2, 256 * 1024)] { + database + .execute( + "INSERT INTO blob_chunks(blob_id,idx,data) VALUES (1,?1,zeroblob(?2))", + params![index, size], + ) + .unwrap(); + } + set_phase(&database, 5).unwrap(); + assert!(!retire_step(&mut database).unwrap()); + assert_eq!( + database + .query_row::("SELECT count(*) FROM blob_chunks", [], |row| row + .get(0)) + .unwrap(), + 2 + ); + assert!(!retire_step(&mut database).unwrap()); + assert_eq!( + database + .query_row::("SELECT count(*) FROM blob_chunks", [], |row| row + .get(0)) + .unwrap(), + 0 + ); + } + } + + #[test] + fn retired_plain_and_keyed_blob_pages_are_reclaimed_by_bounded_idle_vacuum() { + for keyed in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("migration.db"); + let mut database = connection(&path, keyed); + database + .pragma_update(None, "auto_vacuum", "INCREMENTAL") + .unwrap(); + database.pragma_update(None, "journal_mode", "WAL").unwrap(); + database.pragma_update(None, "synchronous", "FULL").unwrap(); + fixture(&mut database, 2); + { + let transaction = database + .transaction_with_behavior(TransactionBehavior::Immediate) + .unwrap(); + for index in 5..17 { + transaction.execute("INSERT INTO blob_chunks(blob_id,idx,data) VALUES (1,?1,zeroblob(262144))", [index]).unwrap(); + } + transaction + .execute("UPDATE blobs SET size=3150848", []) + .unwrap(); + transaction.commit().unwrap(); + } + super::super::connection::checkpoint(&database, true).unwrap(); + let before = std::fs::metadata(&path).unwrap().len(); + schema::migrate(&mut database, 200).unwrap(); + let mut pages: i64 = database + .pragma_query_value(None, "freelist_count", |row| row.get(0)) + .unwrap(); + assert!(pages > 64); + for _ in 0..64 { + if pages == 0 { + break; + } + { + let transaction = database + .transaction_with_behavior(TransactionBehavior::Immediate) + .unwrap(); + let mut statement = transaction + .prepare("PRAGMA incremental_vacuum(64)") + .unwrap(); + let mut rows = statement.query([]).unwrap(); + while rows.next().unwrap().is_some() {} + drop(rows); + drop(statement); + transaction.commit().unwrap(); + } + let remaining: i64 = database + .pragma_query_value(None, "freelist_count", |row| row.get(0)) + .unwrap(); + assert!((1..=64).contains(&(pages - remaining))); + pages = remaining; + super::super::connection::checkpoint(&database, false).unwrap(); + } + assert_eq!(pages, 0); + super::super::connection::checkpoint(&database, true).unwrap(); + assert!(std::fs::metadata(&path).unwrap().len() < before); + schema::validate(&database).unwrap(); + assert_eq!( + database + .query_row::("SELECT count(*) FROM entries", [], |row| row.get(0)) + .unwrap(), + 130 + ); + } + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/attachment_regression_tests.rs b/apps/maple-agent/crates/maple-agent/src/store/attachment_regression_tests.rs new file mode 100644 index 000000000..46b20cbf4 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/attachment_regression_tests.rs @@ -0,0 +1,530 @@ +use super::*; + +const ALICE_TASK: &str = "00000000-0000-7000-8000-000000000001"; +const BOB_TASK: &str = "00000000-0000-7000-8000-000000000002"; + +fn regression_task(index: usize) -> String { + format!("00000000-0000-7000-8000-{index:012x}") +} + +fn reserve_upload(store: &Store, account: AccountId, task: &str, operation: &str, size: u64) { + let mut upload = metadata(b"declared upload fixture"); + upload.size = size; + match store + .begin_attachment_upload(account, task, operation, &upload) + .unwrap() + { + AttachmentUploadStart::Writing(writer) => drop(writer), + AttachmentUploadStart::Complete(_) => panic!("new reservation was already complete"), + } +} + +fn stored_attachment_count(store: &Store, id: &str) -> i64 { + store + .read(|tx| { + Ok(tx.query_row( + "SELECT count(*) FROM file_attachments WHERE public_id=?1", + [id], + |row| row.get(0), + )?) + }) + .unwrap() +} + +fn stored_task_count(store: &Store, uuid: &str) -> i64 { + store + .read(|tx| { + Ok( + tx.query_row("SELECT count(*) FROM tasks WHERE uuid=?1", [uuid], |row| { + row.get(0) + })?, + ) + }) + .unwrap() +} + +fn set_legacy_name(store: &Store, id: &str) { + store + .write(Durability::Durable, |tx| { + // This was a possible persisted result of the old normalizer. It must + // not be accepted by content readers, but deletion still owns its tree. + tx.execute( + "UPDATE file_attachments SET name=' notes.txt' WHERE public_id=?1", + [id], + )?; + Ok(()) + }) + .unwrap(); +} + +#[test] +fn normalized_attachment_names_are_bounded_fixed_points() { + let mut names = vec![ + "* notes.txt".to_owned(), + "? FAQ.md".to_owned(), + "| draft.md".to_owned(), + "../* \u{3000}notes.txt. \u{3000}".to_owned(), + "\u{3000}\t<>:*?\u{3000}".to_owned(), + ".part".to_owned(), + format!("{}\u{3000}suffix", "あ".repeat(59)), + format!("CON.{}", "x".repeat(176)), + ]; + for prefix in ["", "* ", "? \u{3000}", "CON.", "LPT9.", ".part."] { + for length in [0, 1, 168, 169, 179, 180, 181, 240] { + for suffix in ["", ".", " ", "\u{3000}", "\u{3000}. "] { + names.push(format!("{prefix}{}{suffix}", "a".repeat(length))); + names.push(format!("{prefix}{}{suffix}", "あ".repeat(length))); + } + } + } + for name in names { + let normalized = normalized_name(&name); + assert!(!normalized.is_empty(), "empty normalization of {name:?}"); + assert!( + normalized.len() <= 180, + "oversized normalization of {name:?}" + ); + assert_eq!( + normalized_name(&normalized), + normalized, + "not a fixed point: {name:?}" + ); + } + assert_eq!(normalized_name("* notes.txt"), "notes.txt"); + assert_eq!(normalized_name("? FAQ.md"), "FAQ.md"); +} + +#[test] +fn normalized_attachment_names_support_upload_status_retry_and_cancel() { + let (_directory, store, alice, _) = fixture(); + for (index, name) in [ + "* notes.txt".to_owned(), + "? FAQ.md".to_owned(), + "| draft.md".to_owned(), + format!("{}\u{3000}suffix", "あ".repeat(59)), + format!("CON.{}", "x".repeat(176)), + ] + .into_iter() + .enumerate() + { + let operation = format!("normalized-{index}"); + let bytes = b"normalization lifecycle fixture"; + let mut upload = metadata(bytes); + upload.name = name; + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(alice, ALICE_TASK, &operation, &upload) + .unwrap() + else { + panic!("new upload was already complete"); + }; + writer.write_chunk(bytes).unwrap(); + let attachment = writer.finish(&store).unwrap(); + assert_eq!(attachment.name, normalized_name(&upload.name)); + assert_eq!(fs::read(&attachment.execution_path).unwrap(), bytes); + assert_eq!( + store + .attachment_upload_status(alice, ALICE_TASK, &operation) + .unwrap(), + AttachmentUploadStatus::Complete(attachment.clone()), + ); + let AttachmentUploadStart::Complete(retried) = store + .begin_attachment_upload(alice, ALICE_TASK, &operation, &upload) + .unwrap() + else { + panic!("completed retry started another upload"); + }; + assert_eq!(retried, attachment); + store + .cancel_attachment_upload(alice, ALICE_TASK, &operation) + .unwrap(); + assert_eq!( + store + .attachment_upload_status(alice, ALICE_TASK, &operation) + .unwrap(), + AttachmentUploadStatus::Cancelled, + ); + assert!( + !Path::new(&attachment.execution_path) + .parent() + .unwrap() + .exists() + ); + } +} + +#[test] +fn legacy_invalid_names_do_not_prevent_cancel_expiry_or_task_collection() { + let (_directory, store, alice, bob) = fixture(); + let cancelled = write_upload( + &store, + alice, + ALICE_TASK, + "legacy-cancel", + b"cancel fixture", + ); + let expired = write_upload( + &store, + alice, + ALICE_TASK, + "legacy-expire", + b"expired fixture", + ); + let other = write_upload(&store, bob, BOB_TASK, "unrelated", b"other account fixture"); + set_legacy_name(&store, &cancelled.id); + set_legacy_name(&store, &expired.id); + assert!( + store + .attachment_upload_status(alice, ALICE_TASK, "legacy-cancel") + .is_err() + ); + + store + .cancel_attachment_upload(alice, ALICE_TASK, "legacy-cancel") + .unwrap(); + assert!( + !Path::new(&cancelled.execution_path) + .parent() + .unwrap() + .exists() + ); + let cancelled_state: (String, i64) = store + .read(|tx| { + Ok(tx.query_row( + "SELECT state,size FROM file_attachments WHERE public_id=?1", + [&cancelled.id], + |row| Ok((row.get(0)?, row.get(1)?)), + )?) + }) + .unwrap(); + assert_eq!(cancelled_state, ("cancelled".to_owned(), 0)); + + expire(&store, &expired.id); + store.collect_unbound_attachments_step().unwrap(); + assert_eq!(stored_attachment_count(&store, &expired.id), 0); + assert!( + !Path::new(&expired.execution_path) + .parent() + .unwrap() + .exists() + ); + + store.mark_task_deleted(alice, ALICE_TASK, 2).unwrap(); + let mut done = false; + for _ in 0..8 { + if store.collect_task_step(alice, ALICE_TASK).unwrap().done { + done = true; + break; + } + } + assert!(done, "invalid cancellation metadata stranded the task"); + assert_eq!(stored_task_count(&store, ALICE_TASK), 0); + assert_eq!(stored_attachment_count(&store, &cancelled.id), 0); + assert_eq!( + fs::read(&other.execution_path).unwrap(), + b"other account fixture" + ); +} + +#[cfg(unix)] +#[test] +fn widened_owned_roots_do_not_strand_session_or_confirmed_account_deletion() { + use std::os::unix::fs::PermissionsExt; + + for mode in [0o755, 0o711] { + for delete_account in [false, true] { + let (_directory, store, alice, bob) = fixture(); + let owned = write_upload(&store, alice, ALICE_TASK, "widened-root", b"owned archive"); + store + .pin_attachments(alice, ALICE_TASK, std::slice::from_ref(&owned.id)) + .unwrap(); + let other = write_upload( + &store, + bob, + BOB_TASK, + "other-account", + b"keep other account", + ); + let root = Path::new(&owned.execution_path).parent().unwrap(); + fs::set_permissions(root, fs::Permissions::from_mode(mode)).unwrap(); + if delete_account { + store + .begin_account_deletion(alice, "widened-root-deletion", 2) + .unwrap(); + store + .confirm_account_deletion(alice, "widened-root-deletion") + .unwrap(); + } else { + store.mark_task_deleted(alice, ALICE_TASK, 2).unwrap(); + } + let mut task_done = false; + for _ in 0..8 { + if store.collect_task_step(alice, ALICE_TASK).unwrap().done { + task_done = true; + break; + } + } + assert!(task_done, "mode {mode:o} stranded session cleanup"); + assert!(!root.exists()); + assert_eq!(stored_task_count(&store, ALICE_TASK), 0); + if delete_account { + let mut account_done = false; + for _ in 0..16 { + if store.collect_account_step(alice).unwrap().done { + account_done = true; + break; + } + } + assert!(account_done, "mode {mode:o} stranded account cleanup"); + } + assert_eq!( + fs::read(&other.execution_path).unwrap(), + b"keep other account" + ); + } + } +} + +#[test] +fn pending_file_count_is_per_account_even_during_pending_account_deletion() { + let (_directory, store, alice, bob) = fixture(); + for index in 0..MAX_PENDING_FILES { + let task = regression_task(100 + index / MAX_ATTACHMENTS); + if index % MAX_ATTACHMENTS == 0 { + create_task(&store, alice, &task); + } + reserve_upload(&store, alice, &task, &format!("pending-count-{index}"), 1); + } + let overflow = regression_task(200); + create_task(&store, alice, &overflow); + assert!(matches!( + store.begin_attachment_upload(alice, &overflow, "count-overflow", &metadata(b"x")), + Err(StoreError::Busy), + )); + reserve_upload(&store, bob, BOB_TASK, "bob-first", 1); + + store + .begin_account_deletion(alice, "pending-count-deletion", 3) + .unwrap(); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET updated_ms=0 WHERE account_id=?1", + [alice.0], + )?; + Ok(()) + }) + .unwrap(); + assert!(store.collect_unbound_attachments_step().unwrap().done); + let retained: i64 = store + .read(|tx| { + Ok(tx.query_row( + "SELECT count(*) FROM file_attachments WHERE account_id=?1", + [alice.0], + |row| row.get(0), + )?) + }) + .unwrap(); + assert_eq!(retained, MAX_PENDING_FILES as i64); + reserve_upload(&store, bob, BOB_TASK, "bob-while-alice-pending", 1); +} + +#[test] +fn pending_declared_byte_limit_is_per_account() { + let (_directory, store, alice, bob) = fixture(); + assert_eq!(MAX_PENDING_BYTES / MAX_ATTACHMENT_BYTES as u64, 10); + for index in 0..10 { + reserve_upload( + &store, + alice, + ALICE_TASK, + &format!("pending-bytes-{index}"), + MAX_ATTACHMENT_BYTES as u64, + ); + } + // Use another task to distinguish the account byte limit from the pending + // per-task count. Reservations write no declared payload bytes to disk. + let overflow = regression_task(201); + create_task(&store, alice, &overflow); + assert!(matches!( + store.begin_attachment_upload(alice, &overflow, "bytes-overflow", &metadata(b"x")), + Err(StoreError::Busy), + )); + reserve_upload(&store, bob, BOB_TASK, "bob-with-full-other-account", 1); +} + +#[cfg(unix)] +struct UnsafeTaskDirectory { + path: PathBuf, + saved: PathBuf, +} + +#[cfg(unix)] +impl UnsafeTaskDirectory { + fn new(attachment: &StoredAttachment, outside: &Path) -> Self { + let path = Path::new(&attachment.execution_path) + .parent() + .unwrap() + .parent() + .unwrap() + .to_owned(); + let saved = path.with_file_name(format!( + "retained-{}", + path.file_name().unwrap().to_string_lossy() + )); + fs::rename(&path, &saved).unwrap(); + std::os::unix::fs::symlink(outside, &path).unwrap(); + Self { path, saved } + } +} + +#[cfg(unix)] +impl Drop for UnsafeTaskDirectory { + fn drop(&mut self) { + let _ = fs::remove_file(&self.path); + let _ = fs::rename(&self.saved, &self.path); + } +} + +#[cfg(unix)] +#[test] +fn unbound_collection_rotates_past_multiple_unsafe_trees() { + let (_directory, store, alice, bob) = fixture(); + let outside = tempfile::tempdir().unwrap(); + fs::write(outside.path().join("keep"), b"unrelated directory fixture").unwrap(); + let mut blocked = Vec::new(); + let mut guards = Vec::new(); + for index in 0..2 { + let task = regression_task(300 + index); + create_task(&store, alice, &task); + let attachment = write_upload( + &store, + alice, + &task, + &format!("unsafe-unbound-{index}"), + b"retained fixture", + ); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET updated_ms=?2 WHERE public_id=?1", + params![attachment.id, index as i64], + )?; + Ok(()) + }) + .unwrap(); + guards.push(UnsafeTaskDirectory::new(&attachment, outside.path())); + blocked.push(attachment); + } + let healthy = write_upload(&store, bob, BOB_TASK, "healthy-expired", b"healthy fixture"); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET updated_ms=2 WHERE public_id=?1", + [&healthy.id], + )?; + Ok(()) + }) + .unwrap(); + let mut failures = 0; + for _ in 0..6 { + match store.collect_unbound_attachments_step() { + Err(StoreError::UnsafePath) => failures += 1, + Ok(_) => {} + Err(error) => panic!("unexpected collection error: {error:?}"), + } + if stored_attachment_count(&store, &healthy.id) == 0 { + break; + } + } + assert_eq!(failures, 2); + assert_eq!(stored_attachment_count(&store, &healthy.id), 0); + assert!( + !Path::new(&healthy.execution_path) + .parent() + .unwrap() + .exists() + ); + for attachment in &blocked { + assert_eq!(stored_attachment_count(&store, &attachment.id), 1); + } + assert_eq!( + fs::read(outside.path().join("keep")).unwrap(), + b"unrelated directory fixture" + ); + drop(guards); + for _ in 0..4 { + store.collect_unbound_attachments_step().unwrap(); + } + for attachment in blocked { + assert_eq!(stored_attachment_count(&store, &attachment.id), 0); + assert!( + !Path::new(&attachment.execution_path) + .parent() + .unwrap() + .exists() + ); + } +} + +#[cfg(unix)] +#[test] +fn deleted_task_collection_rotates_past_four_unsafe_trees() { + let (_directory, store, alice, bob) = fixture(); + let outside = tempfile::tempdir().unwrap(); + fs::write(outside.path().join("keep"), b"unrelated directory fixture").unwrap(); + let mut blocked = Vec::new(); + let mut guards = Vec::new(); + for index in 0..4 { + let task = regression_task(400 + index); + create_task(&store, alice, &task); + let attachment = write_upload( + &store, + alice, + &task, + &format!("unsafe-deleted-{index}"), + b"retained fixture", + ); + guards.push(UnsafeTaskDirectory::new(&attachment, outside.path())); + store.mark_task_deleted(alice, &task, 1).unwrap(); + blocked.push(task); + } + let healthy = write_upload(&store, bob, BOB_TASK, "healthy-deleted", b"healthy fixture"); + store.mark_task_deleted(bob, BOB_TASK, 2).unwrap(); + let mut failures = 0; + for _ in 0..6 { + for task in store.deleted_tasks_for_collection(4).unwrap() { + match store.collect_task_step(task.account, &task.uuid) { + Err(StoreError::UnsafePath) => failures += 1, + Ok(_) => {} + Err(error) => panic!("unexpected task collection error: {error:?}"), + } + } + if stored_task_count(&store, BOB_TASK) == 0 { + break; + } + } + assert!(failures >= 4); + assert_eq!(stored_task_count(&store, BOB_TASK), 0); + assert!( + !Path::new(&healthy.execution_path) + .parent() + .unwrap() + .exists() + ); + for task in &blocked { + assert_eq!(stored_task_count(&store, task), 1); + } + assert_eq!( + fs::read(outside.path().join("keep")).unwrap(), + b"unrelated directory fixture" + ); + drop(guards); + for _ in 0..4 { + for task in store.deleted_tasks_for_collection(4).unwrap() { + store.collect_task_step(task.account, &task.uuid).unwrap(); + } + } + for task in blocked { + assert_eq!(stored_task_count(&store, &task), 0); + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/attachment_tests.rs b/apps/maple-agent/crates/maple-agent/src/store/attachment_tests.rs new file mode 100644 index 000000000..138866f3a --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/attachment_tests.rs @@ -0,0 +1,1037 @@ +use super::*; +use crate::store::StorePaths; + +#[path = "attachment_regression_tests.rs"] +mod regression_tests; + +fn fixture() -> (tempfile::TempDir, Arc, AccountId, AccountId) { + let directory = tempfile::tempdir().unwrap(); + let store = Arc::new(Store::open(StorePaths::new(directory.path())).unwrap()); + let alice = store.upsert_account("alice", 1).unwrap(); + let bob = store.upsert_account("bob", 1).unwrap(); + create_task(&store, alice, "00000000-0000-7000-8000-000000000001"); + create_task(&store, bob, "00000000-0000-7000-8000-000000000002"); + (directory, store, alice, bob) +} + +fn create_task(store: &Store, account: AccountId, uuid: &str) -> i64 { + store.write(Durability::Durable, |transaction| { + transaction.execute( + "INSERT INTO projects(account_id,root,last_used_ms) VALUES(?1,?2,1) ON CONFLICT DO NOTHING", + params![account.0, format!("/fixture/{}",account.0)], + )?; + let project: i64 = transaction.query_row("SELECT id FROM projects WHERE account_id=?1",[account.0],|row| row.get(0))?; + transaction.execute( + "INSERT INTO tasks(uuid,account_id,project_id,kind,state,title,title_fold,created_ms,updated_ms) \ + VALUES(?1,?2,?3,'desktop','active','fixture','fixture',1,1)", + params![uuid,account.0,project], + )?; + Ok(transaction.last_insert_rowid()) + }).unwrap() +} + +fn png(size: usize) -> Vec { + let mut bytes = vec![17; size.max(8)]; + bytes[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); + bytes +} + +fn upload(bytes: &[u8]) -> ImageUpload { + ImageUpload { + name: "../photo.png".to_string(), + data_url: format!("data:image/png;base64,{}", STANDARD.encode(bytes)), + } +} + +fn metadata(bytes: &[u8]) -> AttachmentUploadMetadata { + AttachmentUploadMetadata { + name: "../report.zip".into(), + mime_type: "application/zip".into(), + size: bytes.len() as u64, + sha256: hex(&Sha256::digest(bytes)), + } +} +fn write_upload( + store: &Store, + account: AccountId, + task: &str, + op: &str, + bytes: &[u8], +) -> StoredAttachment { + match store + .begin_attachment_upload(account, task, op, &metadata(bytes)) + .unwrap() + { + AttachmentUploadStart::Complete(value) => value, + AttachmentUploadStart::Writing(mut writer) => { + for chunk in bytes.chunks(MAX_ATTACHMENT_RANGE_BYTES) { + writer.write_chunk(chunk).unwrap(); + } + writer.finish(store).unwrap() + } + } +} +fn expire(store: &Store, id: &str) { + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET updated_ms=0 WHERE public_id=?1", + [id], + )?; + Ok(()) + }) + .unwrap(); +} + +#[test] +fn cleanup_metadata_rejects_negative_sizes_and_accepts_zero_and_file_cap() { + let (_directory, store, _, _) = fixture(); + for size in [-1, 0, MAX_ATTACHMENT_BYTES as i64] { + let result = store.read(|tx| { + Ok(tx + .query_row( + "SELECT '0123456789abcdef0123456789abcdef','cleanup-size',?1", + [size], + attachment_cleanup_row, + )? + .size) + }); + if size < 0 { + assert_eq!(result, Err(StoreError::InvalidData)); + } else { + assert_eq!(result, Ok(size as u64)); + } + } +} + +#[test] +fn generic_upload_is_a_stable_file_with_only_metadata_in_sql() { + let (_directory, store, alice, bob) = fixture(); + let bytes = b"a ZIP fixture with its own bytes"; + let first = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "operation-1", + bytes, + ); + let second = write_upload( + &store, + bob, + "00000000-0000-7000-8000-000000000002", + "operation-2", + bytes, + ); + assert_eq!(first.name, "report.zip"); + assert_ne!(first.id, second.id); + assert_ne!(first.execution_path, second.execution_path); + assert_eq!(fs::read(&first.execution_path).unwrap(), bytes); + assert_eq!( + store + .read_attachment( + alice, + "00000000-0000-7000-8000-000000000001", + &first.id, + 2, + 5 + ) + .unwrap() + .data, + &bytes[2..7] + ); + assert!( + store + .read_attachment(bob, "00000000-0000-7000-8000-000000000002", &first.id, 0, 1) + .is_err() + ); + assert!( + store + .resolve_attachments( + alice, + "00000000-0000-7000-8000-000000000002", + std::slice::from_ref(&first.id) + ) + .is_err() + ); + store + .read(|tx| { + for legacy in ["blobs", "blob_chunks", "task_blobs", "attachments"] { + assert!(!tx.query_row( + "SELECT EXISTS(SELECT 1 FROM sqlite_schema WHERE name=?1)", + [legacy], + |row| row.get::<_, bool>(0) + )?); + } + let names: Vec = tx + .prepare("PRAGMA table_info(file_attachments)")? + .query_map([], |row| row.get(1))? + .collect::>()?; + assert!(!names.iter().any(|name| name == "data" || name == "payload")); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn operation_replay_is_idempotent_and_bound_to_task_and_metadata() { + let (_directory, store, alice, _) = fixture(); + create_task(&store, alice, "00000000-0000-7000-8000-000000000003"); + let first = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "operation", + b"hello", + ); + assert_eq!( + first, + write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "operation", + b"hello" + ) + ); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "operation", + &metadata(b"other") + ), + Err(StoreError::Conflict) + )); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000003", + "operation", + &metadata(b"hello") + ), + Err(StoreError::Conflict) + )); + assert!( + matches!(store.attachment_upload_status(alice,"00000000-0000-7000-8000-000000000001","operation").unwrap(),AttachmentUploadStatus::Complete(value) if value==first) + ); +} + +#[test] +fn partial_upload_can_restart_but_never_publish_wrong_size_or_digest() { + let (_directory, store, alice, _) = fixture(); + let meta = metadata(b"expected"); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "partial", + &meta, + ) + .unwrap() + else { + panic!() + }; + writer.write_chunk(b"ex").unwrap(); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "partial", + &meta + ), + Err(StoreError::Busy) + )); + drop(writer); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "partial", + &meta, + ) + .unwrap() + else { + panic!() + }; + writer.write_chunk(b"xxxxxxxx").unwrap(); + assert!(matches!( + writer.finish(&store), + Err(StoreError::InvalidInput) + )); + assert_eq!( + store + .attachment_upload_status(alice, "00000000-0000-7000-8000-000000000001", "partial") + .unwrap(), + AttachmentUploadStatus::Uploading + ); + let result = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "partial", + b"expected", + ); + assert_eq!(fs::read(result.execution_path).unwrap(), b"expected"); +} + +#[test] +fn published_file_without_ready_ack_is_reconciled_after_reopen() { + let (directory, store, alice, _) = fixture(); + let attachment = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "lost-ack", + b"published", + ); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET state='uploading' WHERE public_id=?1", + [&attachment.id], + )?; + Ok(()) + }) + .unwrap(); + drop(store); + let store = Store::open(StorePaths::new(directory.path())).unwrap(); + assert_eq!( + attachment, + write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "lost-ack", + b"published" + ) + ); +} + +#[test] +fn cancellation_before_begin_and_during_stream_are_durable_fences() { + let (directory, store, alice, _) = fixture(); + store + .cancel_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "before-begin", + ) + .unwrap(); + assert_eq!( + store + .attachment_upload_status( + alice, + "00000000-0000-7000-8000-000000000001", + "before-begin" + ) + .unwrap(), + AttachmentUploadStatus::Cancelled + ); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "before-begin", + &metadata(b"late") + ), + Err(StoreError::Conflict) + )); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "during", + &metadata(b"hello"), + ) + .unwrap() + else { + panic!() + }; + store + .cancel_attachment_upload(alice, "00000000-0000-7000-8000-000000000001", "during") + .unwrap(); + assert!(matches!( + writer.write_chunk(b"hello"), + Err(StoreError::Conflict) + )); + drop(writer); + drop(store); + let store = Store::open(StorePaths::new(directory.path())).unwrap(); + assert_eq!( + store + .attachment_upload_status(alice, "00000000-0000-7000-8000-000000000001", "during") + .unwrap(), + AttachmentUploadStatus::Cancelled + ); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "before-begin", + &metadata(b"late") + ), + Err(StoreError::Conflict) + )); +} + +#[test] +fn pinning_is_all_or_nothing_and_unresolved_references_survive_staging_gc() { + let (_directory, store, alice, _) = fixture(); + let attachment = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "pin", + b"durable", + ); + assert!( + store + .pin_attachments( + alice, + "00000000-0000-7000-8000-000000000001", + &[attachment.id.clone(), "f".repeat(32)] + ) + .is_err() + ); + assert!( + !store + .read(|tx| Ok(tx.query_row( + "SELECT bound FROM file_attachments WHERE public_id=?1", + [&attachment.id], + |row| row.get::<_, bool>(0) + )?)) + .unwrap() + ); + store + .pin_attachments( + alice, + "00000000-0000-7000-8000-000000000001", + std::slice::from_ref(&attachment.id), + ) + .unwrap(); + expire(&store, &attachment.id); + assert!(store.collect_unbound_attachments_step().unwrap().done); + assert!(Path::new(&attachment.execution_path).exists()); + assert!(matches!( + store.cancel_attachment_upload(alice, "00000000-0000-7000-8000-000000000001", "pin"), + Err(StoreError::Conflict) + )); +} + +#[test] +fn deleting_one_session_removes_only_its_exact_files_and_survives_missing_bytes() { + let (_directory, store, alice, bob) = fixture(); + let first = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "delete", + b"same", + ); + let other = write_upload( + &store, + bob, + "00000000-0000-7000-8000-000000000002", + "keep", + b"same", + ); + store + .mark_task_deleted(alice, "00000000-0000-7000-8000-000000000001", 5) + .unwrap(); + fs::remove_file(&first.execution_path).unwrap(); + for _ in 0..20 { + if store + .collect_task_step(alice, "00000000-0000-7000-8000-000000000001") + .unwrap() + .done + { + break; + } + } + assert!( + store + .task(alice, "00000000-0000-7000-8000-000000000001") + .unwrap() + .is_none() + ); + assert!(!Path::new(&first.execution_path).parent().unwrap().exists()); + assert_eq!(fs::read(&other.execution_path).unwrap(), b"same"); +} + +#[test] +fn missing_effective_image_becomes_unavailable_but_bad_source_stays_an_error() { + let (_directory, store, alice, _) = fixture(); + let attachment = store + .store_image_uploads( + alice, + "00000000-0000-7000-8000-000000000001", + &[upload(&png(100))], + ) + .unwrap() + .remove(0); + let mut messages = vec![Message::User(pi_ai::UserMessage { + content: vec![ + Content::text("keep this text"), + Content::image(&attachment.source, "image/png"), + ], + timestamp: 1, + })]; + fs::remove_file(&attachment.execution_path).unwrap(); + store + .materialize_message_images(alice, "00000000-0000-7000-8000-000000000001", &mut messages) + .unwrap(); + let Message::User(message) = &messages[0] else { + panic!() + }; + assert_eq!( + message.content, + vec![ + Content::text("keep this text"), + Content::text("Image no longer available.") + ] + ); + let mut invalid = vec![Message::User(pi_ai::UserMessage { + content: vec![Content::image( + format!("{ATTACHMENT_SOURCE_PREFIX}{}", "f".repeat(32)), + "image/png", + )], + timestamp: 1, + })]; + assert!( + store + .materialize_message_images(alice, "00000000-0000-7000-8000-000000000001", &mut invalid) + .is_err() + ); +} + +#[test] +fn historical_file_changes_are_unavailable_but_outgoing_images_remain_strict() { + for change in 0..3 { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let bytes = png(100); + let attachment = store + .store_image_uploads(alice, task, &[upload(&bytes)]) + .unwrap() + .remove(0); + let original = vec![Message::User(pi_ai::UserMessage { + content: vec![ + Content::text("keep the caption"), + Content::image(&attachment.source, "image/png"), + ], + timestamp: 1, + })]; + match change { + 0 => { + let mut changed = bytes.clone(); + changed[99] ^= 1; + fs::write(&attachment.execution_path, changed).unwrap(); + } + 1 => fs::write(&attachment.execution_path, b"resized").unwrap(), + 2 => fs::remove_file(&attachment.execution_path).unwrap(), + _ => unreachable!(), + } + let mut outgoing = original.clone(); + assert!( + store + .materialize_outgoing_images(alice, task, &mut outgoing) + .is_err() + ); + let Message::User(outgoing) = &outgoing[0] else { + panic!() + }; + assert_eq!( + outgoing.content[1], + Content::image(&attachment.source, "image/png") + ); + let mut historical = original; + store + .materialize_message_images(alice, task, &mut historical) + .unwrap(); + let Message::User(historical) = &historical[0] else { + panic!() + }; + assert_eq!( + historical.content, + vec![ + Content::text("keep the caption"), + Content::text("Image no longer available.") + ] + ); + } +} + +#[cfg(unix)] +#[test] +fn unsafe_historical_images_are_not_read_or_allowed_in_new_sends() { + use std::os::unix::fs::{PermissionsExt, symlink}; + for change in 0..3 { + let (directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let attachment = store + .store_image_uploads(alice, task, &[upload(&png(100))]) + .unwrap() + .remove(0); + let outside = directory.path().join("outside-image"); + match change { + 0 => fs::set_permissions( + &attachment.execution_path, + fs::Permissions::from_mode(0o644), + ) + .unwrap(), + 1 => { + fs::remove_file(&attachment.execution_path).unwrap(); + fs::write(&outside, b"unrelated image bytes").unwrap(); + symlink(&outside, &attachment.execution_path).unwrap(); + } + 2 => fs::hard_link(&attachment.execution_path, &outside).unwrap(), + _ => unreachable!(), + } + let original = vec![Message::User(pi_ai::UserMessage { + content: vec![Content::image(&attachment.source, "image/png")], + timestamp: 1, + })]; + assert!( + store + .materialize_outgoing_images(alice, task, &mut original.clone()) + .is_err() + ); + let mut historical = original; + store + .materialize_message_images(alice, task, &mut historical) + .unwrap(); + let Message::User(message) = &historical[0] else { + panic!() + }; + assert_eq!( + message.content, + vec![Content::text("Image no longer available.")] + ); + if change == 1 { + assert_eq!(fs::read(&outside).unwrap(), b"unrelated image bytes"); + } + } +} + +#[test] +fn historical_metadata_corruption_is_not_hidden_as_unavailable_media() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let attachment = store + .store_image_uploads(alice, task, &[upload(&png(100))]) + .unwrap() + .remove(0); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE file_attachments SET name=' not-normalized.png' WHERE public_id=?1", + [&attachment.id], + )?; + Ok(()) + }) + .unwrap(); + let mut messages = vec![Message::User(pi_ai::UserMessage { + content: vec![Content::image(&attachment.source, "image/png")], + timestamp: 1, + })]; + assert_eq!( + store.materialize_message_images(alice, task, &mut messages), + Err(StoreError::InvalidData) + ); +} + +#[cfg(unix)] +#[test] +fn confirmed_account_deletion_reclaims_readonly_owned_trees_and_keeps_other_accounts() { + use std::os::unix::fs::PermissionsExt; + let (_directory, store, alice, bob) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let owned = write_upload(&store, alice, task, "readonly-account", b"owned archive"); + store + .pin_attachments(alice, task, std::slice::from_ref(&owned.id)) + .unwrap(); + let other = write_upload( + &store, + bob, + "00000000-0000-7000-8000-000000000002", + "other-account", + b"keep other account", + ); + let root = Path::new(&owned.execution_path).parent().unwrap(); + let extracted = root.join("extracted"); + fs::create_dir(&extracted).unwrap(); + fs::write(extracted.join("plaintext"), b"owned extraction").unwrap(); + fs::set_permissions(&extracted, fs::Permissions::from_mode(0o555)).unwrap(); + store + .begin_account_deletion(alice, "confirmed-readonly", 2) + .unwrap(); + store + .confirm_account_deletion(alice, "confirmed-readonly") + .unwrap(); + let mut done = false; + for _ in 0..32 { + for deleted in store.deleted_tasks_for_collection(4).unwrap() { + store + .collect_task_step(deleted.account, &deleted.uuid) + .unwrap(); + } + match store.collect_account_step(alice) { + Ok(progress) if progress.done => { + done = true; + break; + } + Ok(_) | Err(StoreError::Busy) => {} + Err(error) => panic!("account cleanup failed: {error:?}"), + } + } + assert!(done); + assert!(!root.exists()); + let retained: i64 = store + .read(|tx| { + Ok(tx.query_row( + "SELECT count(*) FROM accounts WHERE id=?1", + [alice.0], + |row| row.get(0), + )?) + }) + .unwrap(); + assert_eq!(retained, 0); + assert_eq!( + fs::read(&other.execution_path).unwrap(), + b"keep other account" + ); +} + +#[test] +fn file_caps_and_image_caps_are_separate_and_chunk_writes_are_bounded() { + let (_directory, store, alice, _) = fixture(); + let bytes = vec![42; MAX_UPLOAD_IMAGE_BYTES + 1]; + let mut large_image_metadata = metadata(&bytes); + large_image_metadata.mime_type = "image/png".into(); + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "image-through-file-api", + &large_image_metadata + ), + Err(StoreError::InvalidInput) + )); + let generic = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "large-file", + &bytes, + ); + assert_eq!(generic.size, bytes.len() as u64); + assert!(matches!( + store.store_image_uploads( + alice, + "00000000-0000-7000-8000-000000000001", + &[upload(&png(MAX_UPLOAD_IMAGE_BYTES + 1))] + ), + Err(StoreError::InvalidInput) + )); + let mut excessive = metadata(b""); + excessive.size = MAX_ATTACHMENT_BYTES as u64 + 1; + assert!(matches!( + store.begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "too-large", + &excessive + ), + Err(StoreError::InvalidInput) + )); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload( + alice, + "00000000-0000-7000-8000-000000000001", + "chunk", + &metadata(&bytes), + ) + .unwrap() + else { + panic!() + }; + assert!(matches!( + writer.write_chunk(&bytes), + Err(StoreError::InvalidInput) + )); +} + +#[test] +fn cleanup_never_follows_a_substituted_attachment_symlink() { + #[cfg(unix)] + { + let (directory, store, alice, _) = fixture(); + let attachment = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "symlink", + b"private", + ); + let outside = directory.path().join("do-not-delete"); + fs::write(&outside, b"outside").unwrap(); + fs::remove_file(&attachment.execution_path).unwrap(); + std::os::unix::fs::symlink(&outside, &attachment.execution_path).unwrap(); + store + .mark_task_deleted(alice, "00000000-0000-7000-8000-000000000001", 5) + .unwrap(); + let progress = store + .collect_task_attachments_step(alice, "00000000-0000-7000-8000-000000000001") + .unwrap(); + assert_eq!(progress.links, 1); + assert_eq!(fs::read(outside).unwrap(), b"outside"); + } +} + +#[test] +fn provisional_session_is_kept_while_upload_has_no_message_yet() { + let (_directory, store, alice, _) = fixture(); + store + .write(Durability::Durable, |tx| { + tx.execute( + "UPDATE tasks SET provisional=1 WHERE uuid='00000000-0000-7000-8000-000000000001'", + [], + )?; + Ok(()) + }) + .unwrap(); + let _attachment = write_upload( + &store, + alice, + "00000000-0000-7000-8000-000000000001", + "draft", + b"still needed", + ); + assert!( + !store + .mark_untouched_provisional_deleted(alice, "00000000-0000-7000-8000-000000000001", 10) + .unwrap() + ); +} + +#[test] +fn repeated_cancel_does_not_consume_active_upload_quota_and_cleans_bytes() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + for index in 0..80 { + let op = format!("cancel-{index}"); + let file = write_upload(&store, alice, task, &op, b"cancelled bytes"); + store.cancel_attachment_upload(alice, task, &op).unwrap(); + assert!(!Path::new(&file.execution_path).exists()); + assert!(store.collect_unbound_attachments_step().unwrap().done); + assert_eq!( + store.attachment_upload_status(alice, task, &op).unwrap(), + AttachmentUploadStatus::Cancelled + ); + } + let file = write_upload(&store, alice, task, "still-usable", b"oops"); + assert_eq!(file.size, 4); +} + +#[test] +fn cancelled_active_writers_reserve_bytes_until_files_are_collected() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let mut metadata = metadata(b"partial"); + metadata.size = MAX_ATTACHMENT_BYTES as u64; + let mut writers = Vec::new(); + for index in 0..MAX_ATTACHMENTS { + let operation = format!("held-cancel-{index}"); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(alice, task, &operation, &metadata) + .unwrap() + else { + panic!("new writer"); + }; + writer.write_chunk(b"partial").unwrap(); + store + .cancel_attachment_upload(alice, task, &operation) + .unwrap(); + assert!(writer.is_cancelled()); + writers.push(writer); + } + assert_eq!(metadata.size * writers.len() as u64, MAX_PENDING_BYTES); + assert!(matches!( + store.begin_attachment_upload(alice, task, "budget-still-reserved", &metadata), + Err(StoreError::Busy) + )); + assert_eq!(store.collect_unbound_attachments_step().unwrap().files, 0); + drop(writers); + let collected = store.collect_unbound_attachments_step().unwrap(); + assert_eq!(collected.bytes, MAX_ATTACHMENT_BYTES); + assert_eq!(collected.links, 1); + assert!(matches!( + store.begin_attachment_upload(alice, task, "budget-released", &metadata), + Ok(AttachmentUploadStart::Writing(_)) + )); +} + +#[test] +fn source_batch_over_composer_count_has_independent_file_storage() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + for index in 0..MAX_ATTACHMENTS { + write_upload(&store, alice, task, &format!("pending-{index}"), b"pending"); + } + let entries = (0..15) + .map(|index| SessionEntry { + id: format!("{index:08x}"), + parent_id: None, + timestamp: 1, + kind: EntryKind::Message { + message: SessionMessage::Llm(Message::User(pi_ai::UserMessage { + content: vec![Content::image(STANDARD.encode(png(20)), "image/png")], + timestamp: 1, + })), + }, + }) + .collect::>(); + let prepared = store.prepare_entries_images(alice, task, &entries).unwrap(); + assert_eq!(prepared.entries.len(), 15); + for entry in prepared.entries { + let EntryKind::Message { + message: SessionMessage::Llm(Message::User(message)), + } = entry.entry.kind + else { + panic!() + }; + let Content::Image(image) = &message.content[0] else { + panic!() + }; + assert!(image.data.starts_with(ATTACHMENT_SOURCE_PREFIX)); + } +} + +#[test] +fn source_image_operation_identity_is_scoped_to_each_session() { + let (_directory, store, alice, _) = fixture(); + let first = "00000000-0000-7000-8000-000000000001"; + let second = "00000000-0000-7000-8000-000000000003"; + create_task(&store, alice, second); + let entry = SessionEntry { + id: "entry-with-a-long-external-identity".repeat(4), + parent_id: None, + timestamp: 1, + kind: EntryKind::Message { + message: SessionMessage::Llm(Message::User(pi_ai::UserMessage { + content: vec![Content::image(STANDARD.encode(png(20)), "image/png")], + timestamp: 1, + })), + }, + }; + let a = store.prepare_entry_images(alice, first, &entry).unwrap(); + let b = store.prepare_entry_images(alice, second, &entry).unwrap(); + assert_ne!(a.attachments, b.attachments); + let apath = store + .attachment_metadata(alice, first, &a.attachments[0]) + .unwrap() + .execution_path; + let bpath = store + .attachment_metadata(alice, second, &b.attachments[0]) + .unwrap() + .execution_path; + assert_ne!(apath, bpath); + assert_eq!(fs::read(apath).unwrap(), fs::read(bpath).unwrap()); +} + +#[test] +fn stored_source_rejects_inline_images_while_incoming_media_is_externalized() { + let entry = SessionEntry { + id: "01020304".into(), + parent_id: None, + timestamp: 1, + kind: EntryKind::Message { + message: SessionMessage::Llm(Message::User(pi_ai::UserMessage { + content: vec![Content::image(STANDARD.encode(png(20)), "image/png")], + timestamp: 1, + })), + }, + }; + let encoded = super::super::codec::encode(&serde_json::to_vec(&entry).unwrap()).unwrap(); + assert!(matches!( + super::super::session::decode_entry(encoded.0, &encoded.1), + Err(StoreError::InvalidData) + )); +} + +#[test] +fn partial_filename_is_reserved_case_insensitively() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let mut metadata = metadata(b"data"); + metadata.name = ".PART".into(); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(alice, task, "reserved-name", &metadata) + .unwrap() + else { + panic!() + }; + writer.write_chunk(b"data").unwrap(); + let file = writer.finish(&store).unwrap(); + assert_eq!(file.name, "attachment-.PART"); + assert_eq!(fs::read(file.execution_path).unwrap(), b"data"); +} + +#[test] +fn image_mime_mismatch_is_refused_but_arbitrary_file_bytes_are_allowed() { + let (_directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let bytes = b"PK\x03\x04a zip, not a png"; + let mut meta = metadata(bytes); + meta.mime_type = "image/png".into(); + let AttachmentUploadStart::Writing(mut writer) = store + .begin_attachment_upload(alice, task, "false-image", &meta) + .unwrap() + else { + panic!() + }; + writer.write_chunk(bytes).unwrap(); + assert!(matches!( + writer.finish(&store), + Err(StoreError::InvalidInput) + )); + let file = write_upload(&store, alice, task, "ordinary-zip", bytes); + assert_eq!(fs::read(file.execution_path).unwrap(), bytes); +} + +#[test] +fn session_collection_reclaims_adjacent_outputs_and_nested_folders_in_bounded_steps() { + let (directory, store, alice, _) = fixture(); + let task = "00000000-0000-7000-8000-000000000001"; + let attachment = write_upload(&store, alice, task, "expanded-archive", b"archive"); + let root = Path::new(&attachment.execution_path).parent().unwrap(); + fs::write(root.join(".DS_Store"), b"finder metadata").unwrap(); + let nested = root.join("unpacked").join("nested"); + fs::create_dir_all(&nested).unwrap(); + for index in 0..300 { + fs::write( + nested.join(format!("result-{index}")), + b"model-created output", + ) + .unwrap(); + } + let outside = directory.path().join("outside-folder"); + fs::create_dir(&outside).unwrap(); + fs::write(outside.join("keep"), b"unrelated bytes").unwrap(); + #[cfg(unix)] + std::os::unix::fs::symlink(&outside, root.join("external-link")).unwrap(); + store.mark_task_deleted(alice, task, 2).unwrap(); + let mut steps = 0; + loop { + let progress = store.collect_task_attachments_step(alice, task).unwrap(); + assert!(progress.files <= 128); + steps += 1; + if progress.done { + break; + } + assert!(steps < 100, "bounded attachment cleanup did not finish"); + } + assert!(steps > 1); + assert!(!root.exists()); + assert_eq!(fs::read(outside.join("keep")).unwrap(), b"unrelated bytes"); +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/attachments.rs b/apps/maple-agent/crates/maple-agent/src/store/attachments.rs new file mode 100644 index 000000000..f4cc40a49 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/attachments.rs @@ -0,0 +1,1498 @@ +//! Session-owned attachment files. SQLite contains only bounded metadata. +//! +//! Upload operations are immutable and idempotent. A private file is flushed and +//! published before its metadata becomes ready; a retry reconciles that narrow +//! crash window. No database transaction is retained during file I/O. + +use std::collections::{HashMap, HashSet}; +use std::fs::{self, File, OpenOptions}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, +}; + +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use pi_ai::{Content, Message}; +use pi_coding_agent::{ + SessionMessage, + session::{EntryKind, SessionEntry}, +}; +use rusqlite::{OptionalExtension, Transaction, params}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use super::paths::{private_directory, private_new_file, sync_directory, verify_private_path}; +use super::{AccountId, Durability, Store, StoreError, StoreResult}; + +pub const MAX_ATTACHMENT_BYTES: usize = 30 * 1024 * 1024; +pub const MAX_ATTACHMENT_RANGE_BYTES: usize = 1024 * 1024; +pub const MAX_ATTACHMENTS: usize = 10; +pub const MAX_UPLOAD_IMAGE_BYTES: usize = 10 * 1024 * 1024; +const MAX_GENERATED_IMAGE_BYTES: usize = 20 * 1024 * 1024; +const MAX_PENDING_FILES: usize = 64; +const MAX_CANCELLATION_TOMBSTONES: i64 = 4096; +const MAX_PENDING_BYTES: u64 = 300 * 1024 * 1024; +const STAGING_TTL_MS: i64 = 24 * 60 * 60 * 1000; +pub const ATTACHMENT_SOURCE_PREFIX: &str = "maple-attachment://"; +pub const UNAVAILABLE_IMAGE_SOURCE: &str = "maple-attachment-unavailable:retired"; + +#[derive(Default)] +pub(crate) struct AttachmentLocks { + held: Mutex>>, +} + +impl AttachmentLocks { + fn acquire(self: &Arc, key: String) -> StoreResult { + let mut held = self.held.lock().map_err(|_| StoreError::InvalidData)?; + if held.contains_key(&key) { + return Err(StoreError::Busy); + } + let cancelled = Arc::new(AtomicBool::new(false)); + held.insert(key.clone(), Arc::clone(&cancelled)); + Ok(UploadLease { + locks: Arc::clone(self), + key, + cancelled, + }) + } + fn cancel(&self, key: &str) -> StoreResult<()> { + if let Some(cancelled) = self + .held + .lock() + .map_err(|_| StoreError::InvalidData)? + .get(key) + { + cancelled.store(true, Ordering::Release); + } + Ok(()) + } +} + +struct UploadLease { + locks: Arc, + key: String, + cancelled: Arc, +} +impl UploadLease { + fn check(&self) -> StoreResult<()> { + if self.cancelled.load(Ordering::Acquire) { + Err(StoreError::Conflict) + } else { + Ok(()) + } + } +} +impl Drop for UploadLease { + fn drop(&mut self) { + self.locks + .held + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&self.key); + } +} + +#[derive(Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ImageUpload { + pub name: String, + pub data_url: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AttachmentUploadMetadata { + pub name: String, + pub mime_type: String, + pub size: u64, + /// Lower-case SHA-256 hex of the raw bytes, checked again while streaming. + pub sha256: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct StoredAttachment { + pub id: String, + pub name: String, + pub mime_type: String, + pub source: String, + pub size: u64, + pub execution_path: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "state", content = "attachment", rename_all = "camelCase")] +pub enum AttachmentUploadStatus { + Missing, + Uploading, + Complete(StoredAttachment), + Cancelled, +} + +pub enum AttachmentUploadStart { + Complete(StoredAttachment), + Writing(Box), +} + +pub struct AttachmentUploadWriter { + account: AccountId, + task_uuid: String, + row: AttachmentRow, + file: File, + partial: PathBuf, + published: PathBuf, + written: u64, + hash: Sha256, + header: Vec, + lease: UploadLease, +} + +/// A verified, synced partial file. Publication still requires current account, +/// task and (for remote uploads) grant authority; preparation does not publish. +pub struct PreparedAttachmentUpload { + account: AccountId, + task_uuid: String, + row: AttachmentRow, + partial: PathBuf, + published: PathBuf, + lease: UploadLease, +} + +impl AttachmentUploadWriter { + pub fn is_cancelled(&self) -> bool { + self.lease.cancelled.load(Ordering::Acquire) + } + + pub fn write_chunk(&mut self, bytes: &[u8]) -> StoreResult<()> { + self.lease.check()?; + if bytes.len() > MAX_ATTACHMENT_RANGE_BYTES + || self.written.saturating_add(bytes.len() as u64) > self.row.metadata.size + { + return Err(StoreError::InvalidInput); + } + let header_bytes = (16 - self.header.len()).min(bytes.len()); + self.header.extend_from_slice(&bytes[..header_bytes]); + self.file.write_all(bytes)?; + self.hash.update(bytes); + self.written += bytes.len() as u64; + Ok(()) + } + + pub fn finish(self, store: &Store) -> StoreResult { + self.prepare(store)?.publish(store) + } + + pub fn prepare(self, store: &Store) -> StoreResult { + self.lease.check()?; + if self.written != self.row.metadata.size + || hex(&self.hash.clone().finalize()) != self.row.metadata.sha256 + { + return Err(StoreError::InvalidInput); + } + if self.row.metadata.mime_type.starts_with("image/") + && detected_image_mime(&self.header) != Some(self.row.metadata.mime_type.as_str()) + { + return Err(StoreError::InvalidInput); + } + store.read(|tx| { + let task = active_task(tx, self.account, &self.task_uuid)?; + verify_upload_row(tx, task, &self.row, "uploading") + })?; + super::keys::sync_key_file(&self.file)?; + // Windows cannot rename while our writable handle remains open. + drop(self.file); + Ok(PreparedAttachmentUpload { + account: self.account, + task_uuid: self.task_uuid, + row: self.row, + partial: self.partial, + published: self.published, + lease: self.lease, + }) + } +} + +impl PreparedAttachmentUpload { + pub fn publish(self, store: &Store) -> StoreResult { + self.lease.check()?; + store.read(|tx| { + let task = active_task(tx, self.account, &self.task_uuid)?; + verify_upload_row(tx, task, &self.row, "uploading") + })?; + publish_file(&self.partial, &self.published)?; + store.finish_attachment_metadata(self.account, &self.task_uuid, &self.row)?; + Ok(self.row.descriptor(&self.published)) + } +} + +#[derive(Clone, PartialEq, Eq)] +pub struct AttachmentRange { + pub mime: String, + pub size: u64, + pub data: Vec, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct AttachmentCollectionProgress { + pub done: bool, + pub files: usize, + pub bytes: usize, + pub links: usize, +} + +#[derive(Clone)] +struct AttachmentRow { + id: String, + operation: String, + metadata: AttachmentUploadMetadata, + state: String, +} + +/// Deletion needs ownership and an operation lease, not readable display or +/// content metadata. In particular, an old invalid filename must be reclaimable. +struct AttachmentCleanupRow { + id: String, + operation: String, + size: u64, +} + +fn attachment_cleanup_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(AttachmentCleanupRow { + id: row.get(0)?, + operation: row.get(1)?, + size: { + let size = row.get::<_, i64>(2)?; + u64::try_from(size).map_err(|_| rusqlite::Error::IntegralValueOutOfRange(2, size))? + }, + }) +} +impl AttachmentRow { + fn descriptor(&self, path: &Path) -> StoredAttachment { + StoredAttachment { + id: self.id.clone(), + name: self.metadata.name.clone(), + mime_type: self.metadata.mime_type.clone(), + source: format!("{ATTACHMENT_SOURCE_PREFIX}{}", self.id), + size: self.metadata.size, + execution_path: path.to_string_lossy().into_owned(), + } + } +} + +pub(super) struct PreparedEntryImages { + pub(super) entry: SessionEntry, + account: AccountId, + task_uuid: String, + attachments: Vec, +} +pub(super) struct PreparedEntriesImages { + pub(super) entries: Vec, +} +impl PreparedEntryImages { + pub(super) fn has_images(&self) -> bool { + !self.attachments.is_empty() + } + pub(super) fn link(&self, tx: &Transaction<'_>, task_id: i64) -> StoreResult<()> { + if active_task(tx, self.account, &self.task_uuid)? != task_id { + return Err(StoreError::NotFound); + } + pin_tx(tx, task_id, &self.attachments) + } +} + +impl Store { + pub(super) fn cancel_attachment_writers(&self) -> StoreResult { + let held = self + .attachment_locks + .held + .lock() + .map_err(|_| StoreError::InvalidData)?; + for cancelled in held.values() { + cancelled.store(true, Ordering::Release); + } + Ok(held.is_empty()) + } + + pub fn begin_attachment_upload( + &self, + account: AccountId, + task_uuid: &str, + operation_id: &str, + metadata: &AttachmentUploadMetadata, + ) -> StoreResult { + self.begin_attachment_upload_inner(account, task_uuid, operation_id, metadata, false) + } + + fn begin_attachment_upload_inner( + &self, + account: AccountId, + task_uuid: &str, + operation_id: &str, + metadata: &AttachmentUploadMetadata, + source_image: bool, + ) -> StoreResult { + validate_operation(operation_id)?; + let metadata = normalized_metadata(metadata)?; + // Never durably insert metadata that our own readers would reject. + if normalized_name(&metadata.name) != metadata.name { + return Err(StoreError::InvalidInput); + } + if !source_image + && metadata.mime_type.starts_with("image/") + && metadata.size > MAX_UPLOAD_IMAGE_BYTES as u64 + { + return Err(StoreError::InvalidInput); + } + let lease = self + .attachment_locks + .acquire(operation_key(account, operation_id))?; + let row = self.write(Durability::Durable, |tx| { + let task = active_task(tx, account, task_uuid)?; + let existing = tx.query_row("SELECT public_id,operation_id,name,mime,size,sha256,state,task_id FROM file_attachments WHERE account_id=?1 AND operation_id=?2", params![account.0,operation_id], |row| Ok((attachment_row(row)?,row.get::<_,i64>(7)?))).optional()?; + if let Some((row, stored_task)) = existing { + if stored_task != task || row.metadata != metadata || row.state == "cancelled" { return Err(StoreError::Conflict); } + return Ok(row); + } + let (count,bytes,task_count):(i64,i64,i64) = tx.query_row("SELECT coalesce(sum(state!='cancelled'),0),coalesce(sum(size),0),coalesce(sum(task_id=?1 AND state!='cancelled'),0) FROM file_attachments WHERE bound=0 AND account_id=?2", params![task,account.0], |row| Ok((row.get(0)?,row.get(1)?,row.get(2)?)))?; + if !source_image && (count >= MAX_PENDING_FILES as i64 || task_count >= MAX_ATTACHMENTS as i64 || bytes as u64 + metadata.size > MAX_PENDING_BYTES) { return Err(StoreError::Busy); } + let tombstones:i64=tx.query_row("SELECT count(*) FROM file_attachments WHERE account_id=?1 AND state='cancelled'",[account.0],|row|row.get(0))?; + if !source_image && tombstones>=MAX_CANCELLATION_TOMBSTONES {return Err(StoreError::Busy);} + let row = AttachmentRow { id: format!("{:032x}",rand::random::()), operation: operation_id.to_owned(), metadata: metadata.clone(), state: "uploading".into() }; + tx.execute("INSERT INTO file_attachments(public_id,account_id,task_id,operation_id,name,mime,size,sha256,state,created_ms,updated_ms) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,'uploading',?9,?9)", params![row.id,account.0,task,row.operation,row.metadata.name,row.metadata.mime_type,i64::try_from(row.metadata.size).map_err(|_|StoreError::InvalidInput)?,row.metadata.sha256,pi_ai::now_ms()])?; + Ok(row) + })?; + let path = self.attachment_path(account, task_uuid, &row, true)?; + if row.state == "ready" { + checked_file(&path, row.metadata.size)?; + return Ok(AttachmentUploadStart::Complete(row.descriptor(&path))); + } + if path.try_exists()? { + verify_file_digest(&path, &row.metadata)?; + self.finish_attachment_metadata(account, task_uuid, &row)?; + return Ok(AttachmentUploadStart::Complete(row.descriptor(&path))); + } + let partial = path.parent().ok_or(StoreError::UnsafePath)?.join(".part"); + remove_private_file_if_present(&partial)?; + let file = private_new_file(&partial)?; + Ok(AttachmentUploadStart::Writing(Box::new( + AttachmentUploadWriter { + account, + task_uuid: task_uuid.to_owned(), + row, + file, + partial, + published: path, + written: 0, + hash: Sha256::new(), + header: Vec::with_capacity(16), + lease, + }, + ))) + } + + fn finish_attachment_metadata( + &self, + account: AccountId, + task_uuid: &str, + row: &AttachmentRow, + ) -> StoreResult<()> { + self.write(Durability::Durable, |tx| { + let task = active_task(tx, account, task_uuid)?; + verify_upload_row(tx, task, row, "uploading")?; + tx.execute( + "UPDATE file_attachments SET state='ready',updated_ms=?2 WHERE public_id=?1", + params![row.id, pi_ai::now_ms()], + )?; + Ok(()) + }) + } + + pub fn attachment_upload_status( + &self, + account: AccountId, + task_uuid: &str, + operation_id: &str, + ) -> StoreResult { + validate_operation(operation_id)?; + let row = self.read(|tx| { + let task = active_task(tx,account,task_uuid)?; + let found = tx.query_row("SELECT public_id,operation_id,name,mime,size,sha256,state,task_id FROM file_attachments WHERE account_id=?1 AND operation_id=?2",params![account.0,operation_id],|row| Ok((attachment_row(row)?,row.get::<_,i64>(7)?))).optional()?; + match found { Some((_,owner)) if owner != task => Err(StoreError::NotFound), other => Ok(other.map(|(row,_)|row)) } + })?; + match row { + None => Ok(AttachmentUploadStatus::Missing), + Some(row) if row.state == "cancelled" => Ok(AttachmentUploadStatus::Cancelled), + Some(row) if row.state == "ready" => { + let path = self.attachment_path(account, task_uuid, &row, false)?; + checked_file(&path, row.metadata.size)?; + Ok(AttachmentUploadStatus::Complete(row.descriptor(&path))) + } + Some(_) => Ok(AttachmentUploadStatus::Uploading), + } + } + + pub fn cancel_attachment_upload( + &self, + account: AccountId, + task_uuid: &str, + operation_id: &str, + ) -> StoreResult<()> { + validate_operation(operation_id)?; + self.write(Durability::Durable,|tx| { + let task=active_task(tx,account,task_uuid)?; + let existing:Option<(i64,bool)>=tx.query_row("SELECT task_id,bound FROM file_attachments WHERE account_id=?1 AND operation_id=?2",params![account.0,operation_id],|row|Ok((row.get(0)?,row.get(1)?))).optional()?; + match existing { + Some((owner,_)) if owner!=task => return Err(StoreError::NotFound), + Some((_,true)) => return Err(StoreError::Conflict), + Some(_) => {tx.execute("UPDATE file_attachments SET state='cancelled',updated_ms=?3 WHERE account_id=?1 AND operation_id=?2 AND bound=0",params![account.0,operation_id,pi_ai::now_ms()])?;}, + None => { + let count:i64=tx.query_row("SELECT count(*) FROM file_attachments WHERE account_id=?1 AND state='cancelled'",[account.0],|row|row.get(0))?; + if count>=MAX_CANCELLATION_TOMBSTONES {return Err(StoreError::Busy);} + tx.execute("INSERT INTO file_attachments(public_id,account_id,task_id,operation_id,name,mime,size,sha256,state,created_ms,updated_ms) VALUES(?1,?2,?3,?4,'attachment','application/octet-stream',0,?5,'cancelled',?6,?6)",params![format!("{:032x}",rand::random::()),account.0,task,operation_id,"0".repeat(64),pi_ai::now_ms()])?; + } + } + Ok(()) + })?; + self.attachment_locks + .cancel(&operation_key(account, operation_id))?; + // Cancellation is already durable. Reclaim an idle upload promptly in + // one bounded step; active writers and retryable filesystem failures + // remain for maintenance, with their reserved bytes still charged. + let candidate = self.read(|tx| { + let task = active_task(tx, account, task_uuid)?; + Ok(tx.query_row( + "SELECT public_id,operation_id,size FROM file_attachments WHERE task_id=?1 AND operation_id=?2 AND state='cancelled' AND size>0", + params![task, operation_id], attachment_cleanup_row, + ).optional()?) + })?; + if let Some(row) = candidate { + let _ = self.collect_attachment_file(account, task_uuid, &row, false); + } + Ok(()) + } + + /// Refresh persisted model references after a store-generation recovery. + /// This returns authorized metadata only; reads still use checked file APIs. + pub fn attachment_metadata( + &self, + account: AccountId, + task_uuid: &str, + id: &str, + ) -> StoreResult { + validate_id(id)?; + let row = self.read(|tx| lookup_ready(tx, active_task(tx, account, task_uuid)?, id))?; + let path = self.attachment_path_metadata(account, task_uuid, &row)?; + Ok(row.descriptor(&path)) + } + + pub fn resolve_attachments( + &self, + account: AccountId, + task_uuid: &str, + ids: &[String], + ) -> StoreResult> { + if ids.len() > MAX_ATTACHMENTS || ids.iter().collect::>().len() != ids.len() { + return Err(StoreError::InvalidInput); + } + ids.iter() + .map(|id| { + validate_id(id)?; + let row = self.read(|tx| { + let task = active_task(tx, account, task_uuid)?; + lookup_ready(tx, task, id) + })?; + let path = self.attachment_path(account, task_uuid, &row, false)?; + checked_file(&path, row.metadata.size)?; + Ok(row.descriptor(&path)) + }) + .collect() + } + + pub fn pin_attachments( + &self, + account: AccountId, + task_uuid: &str, + ids: &[String], + ) -> StoreResult<()> { + if ids.len() > MAX_ATTACHMENTS { + return Err(StoreError::InvalidInput); + } + let mut rows = self.read(|tx| { + let task = active_task(tx, account, task_uuid)?; + ids.iter() + .map(|id| { + validate_id(id)?; + lookup_ready(tx, task, id) + }) + .collect::>>() + })?; + rows.sort_by(|a, b| a.operation.cmp(&b.operation)); + rows.dedup_by(|a, b| a.operation == b.operation); + let _leases = rows + .iter() + .map(|row| { + self.attachment_locks + .acquire(operation_key(account, &row.operation)) + }) + .collect::>>()?; + for row in &rows { + checked_file( + &self.attachment_path(account, task_uuid, row, false)?, + row.metadata.size, + )?; + } + self.write(Durability::Durable, |tx| { + let task = active_task(tx, account, task_uuid)?; + pin_tx(tx, task, ids) + }) + } + + pub fn store_image_uploads( + &self, + account: AccountId, + task_uuid: &str, + uploads: &[ImageUpload], + ) -> StoreResult> { + if uploads.len() > MAX_ATTACHMENTS { + return Err(StoreError::InvalidInput); + } + let parsed = uploads + .iter() + .map(parse_upload) + .collect::>>()?; + parsed + .into_iter() + .map(|(metadata, bytes)| { + self.import_attachment_bytes( + account, + task_uuid, + &format!("{:032x}", rand::random::()), + metadata, + &bytes, + false, + ) + }) + .collect() + } + + fn import_attachment_bytes( + &self, + account: AccountId, + task_uuid: &str, + operation: &str, + metadata: AttachmentUploadMetadata, + bytes: &[u8], + source_image: bool, + ) -> StoreResult { + match self.begin_attachment_upload_inner( + account, + task_uuid, + operation, + &metadata, + source_image, + )? { + AttachmentUploadStart::Complete(value) => Ok(value), + AttachmentUploadStart::Writing(mut writer) => { + for chunk in bytes.chunks(MAX_ATTACHMENT_RANGE_BYTES) { + writer.write_chunk(chunk)?; + } + writer.finish(self) + } + } + } + + pub fn read_attachment( + &self, + account: AccountId, + task_uuid: &str, + public_id: &str, + offset: u64, + len: usize, + ) -> StoreResult { + validate_id(public_id)?; + if len > MAX_ATTACHMENT_RANGE_BYTES { + return Err(StoreError::InvalidInput); + } + let row = + self.read(|tx| lookup_ready(tx, active_task(tx, account, task_uuid)?, public_id))?; + if offset > row.metadata.size { + return Err(StoreError::InvalidInput); + } + let path = self.attachment_path(account, task_uuid, &row, false)?; + let mut file = checked_file(&path, row.metadata.size)?; + file.seek(SeekFrom::Start(offset))?; + let mut data = vec![0; len.min((row.metadata.size - offset) as usize)]; + file.read_exact(&mut data)?; + Ok(AttachmentRange { + mime: row.metadata.mime_type, + size: row.metadata.size, + data, + }) + } + + /// Call only on the effective context after compaction and image filtering. + /// Changed, inaccessible or unsafe historical files are unavailable media; + /// malformed persisted metadata/source stays an error. + pub fn materialize_message_images( + &self, + account: AccountId, + task_uuid: &str, + messages: &mut [Message], + ) -> StoreResult<()> { + self.materialize_images(account, task_uuid, messages, true) + } + + /// A new send must not silently discard an attachment the sender selected. + pub fn materialize_outgoing_images( + &self, + account: AccountId, + task_uuid: &str, + messages: &mut [Message], + ) -> StoreResult<()> { + self.materialize_images(account, task_uuid, messages, false) + } + + fn materialize_images( + &self, + account: AccountId, + task_uuid: &str, + messages: &mut [Message], + historical: bool, + ) -> StoreResult<()> { + let mut prepared = messages.to_vec(); + for message in &mut prepared { + let content = match message { + Message::User(m) => &mut m.content, + Message::ToolResult(m) => &mut m.content, + _ => continue, + }; + for block in content { + let Content::Image(image) = block else { + continue; + }; + if image.data == UNAVAILABLE_IMAGE_SOURCE { + if !historical { + return Err(StoreError::NotFound); + } + *block = Content::text("Image no longer available."); + continue; + } + let Some(id) = image.data.strip_prefix(ATTACHMENT_SOURCE_PREFIX) else { + continue; + }; + validate_id(id).map_err(|_| StoreError::InvalidData)?; + let row = + self.read(|tx| lookup_ready(tx, active_task(tx, account, task_uuid)?, id))?; + if image.mime_type != row.metadata.mime_type { + return Err(StoreError::InvalidData); + } + // Resolve database-owned identity outside the file-error path. + // Corrupt rows/store metadata must never become missing media. + let path = self.attachment_path_metadata(account, task_uuid, &row)?; + let bytes = (|| { + self.prepare_attachment_path(&path, false)?; + let mut file = checked_file(&path, row.metadata.size)?; + let mut bytes = vec![0; row.metadata.size as usize]; + file.read_exact(&mut bytes)?; + if hex(&Sha256::digest(&bytes)) != row.metadata.sha256 { + return Err(StoreError::InvalidData); + } + Ok(bytes) + })(); + match bytes { + Ok(bytes) => image.data = STANDARD.encode(bytes), + Err(error) + if historical + && matches!( + error, + StoreError::NotFound + | StoreError::UnsafePath + | StoreError::InvalidData + | StoreError::Io { .. } + ) => + { + *block = Content::text("Image no longer available.") + } + Err(error) => return Err(error), + } + } + } + messages.clone_from_slice(&prepared); + Ok(()) + } + + pub(super) fn prepare_entry_images( + &self, + account: AccountId, + task_uuid: &str, + entry: &SessionEntry, + ) -> StoreResult { + self.prepare_entries_images(account, task_uuid, std::slice::from_ref(entry))? + .entries + .pop() + .ok_or(StoreError::InvalidData) + } + pub(super) fn prepare_entries_images( + &self, + account: AccountId, + task_uuid: &str, + entries: &[SessionEntry], + ) -> StoreResult { + self.read(|tx| active_task(tx, account, task_uuid))?; + let mut prepared = Vec::with_capacity(entries.len()); + for source in entries { + let mut entry = source.clone(); + let mut ids = Vec::new(); + let mut ordinal = 0; + let entry_id = entry.id.clone(); + visit_images(&mut entry, |image| { + ordinal += 1; + if image.data == UNAVAILABLE_IMAGE_SOURCE { + return Ok(()); + } + if let Some(id) = image.data.strip_prefix(ATTACHMENT_SOURCE_PREFIX) { + validate_id(id).map_err(|_| StoreError::InvalidData)?; + let bound = self.read(|tx| { + let task = active_task(tx, account, task_uuid)?; + let stored = lookup_ready(tx, task, id)?; + if stored.metadata.mime_type != image.mime_type { + return Err(StoreError::InvalidData); + } + Ok(tx.query_row( + "SELECT bound FROM file_attachments WHERE public_id=?1", + [id], + |row| row.get::<_, bool>(0), + )?) + })?; + if !bound { + self.pin_attachments(account, task_uuid, &[id.to_owned()])?; + } + ids.push(id.to_owned()); + return Ok(()); + } + if image.data.len() > MAX_GENERATED_IMAGE_BYTES.div_ceil(3) * 4 + || !image.mime_type.starts_with("image/") + { + image.data = UNAVAILABLE_IMAGE_SOURCE.into(); + return Ok(()); + } + let encoded_data = if image.data.starts_with("data:") { + let expected = format!("data:{};base64,", image.mime_type); + match image.data.strip_prefix(&expected) { + Some(data) => data, + None => { + image.data = UNAVAILABLE_IMAGE_SOURCE.into(); + return Ok(()); + } + } + } else { + image.data.as_str() + }; + let encoded: String = encoded_data + .chars() + .filter(|character| !character.is_ascii_whitespace()) + .collect(); + let bytes = match STANDARD.decode(&encoded) { + Ok(bytes) if !bytes.is_empty() && bytes.len() <= MAX_GENERATED_IMAGE_BYTES => { + bytes + } + _ => { + image.data = UNAVAILABLE_IMAGE_SOURCE.into(); + return Ok(()); + } + }; + if detected_image_mime(&bytes) != Some(image.mime_type.as_str()) { + image.data = UNAVAILABLE_IMAGE_SOURCE.into(); + return Ok(()); + } + let sha256 = hex(&Sha256::digest(&bytes)); + let mut identity = Sha256::new(); + for part in [ + task_uuid, + entry_id.as_str(), + ordinal.to_string().as_str(), + sha256.as_str(), + ] { + identity.update((part.len() as u64).to_le_bytes()); + identity.update(part.as_bytes()); + } + let operation = format!("source-{}", hex(&identity.finalize())); + let metadata = AttachmentUploadMetadata { + name: format!("image-{ordinal}.{}", image_extension(&image.mime_type)), + mime_type: image.mime_type.clone(), + size: bytes.len() as u64, + sha256, + }; + if normalized_metadata(&metadata).is_err() { + image.data = UNAVAILABLE_IMAGE_SOURCE.into(); + return Ok(()); + } + let stored = self.import_attachment_bytes( + account, task_uuid, &operation, metadata, &bytes, true, + )?; + self.pin_attachments(account, task_uuid, std::slice::from_ref(&stored.id))?; + image.data = stored.source; + ids.push(stored.id); + Ok(()) + })?; + prepared.push(PreparedEntryImages { + entry, + account, + task_uuid: task_uuid.to_owned(), + attachments: ids, + }); + } + Ok(PreparedEntriesImages { entries: prepared }) + } + + fn attachment_path_metadata( + &self, + account: AccountId, + task_uuid: &str, + row: &AttachmentRow, + ) -> StoreResult { + if !safe_component(task_uuid) || normalized_name(&row.metadata.name) != row.metadata.name { + return Err(StoreError::InvalidData); + } + Ok(self + .attachment_directory_metadata(account, task_uuid, &row.id)? + .join(&row.metadata.name)) + } + + fn attachment_directory_metadata( + &self, + account: AccountId, + task_uuid: &str, + id: &str, + ) -> StoreResult { + validate_id(id).map_err(|_| StoreError::InvalidData)?; + if account.0 <= 0 || !safe_component(task_uuid) { + return Err(StoreError::InvalidData); + } + let store_id: Vec = self.read(|tx| { + Ok(tx.query_row( + "SELECT value FROM store_meta WHERE key='store_id'", + [], + |row| row.get(0), + )?) + })?; + if store_id.len() != 16 { + return Err(StoreError::InvalidData); + } + Ok(self + .paths + .attachments() + .join(hex(&store_id)) + .join(account.0.to_string()) + .join(task_uuid) + .join(id)) + } + + fn attachment_path( + &self, + account: AccountId, + task_uuid: &str, + row: &AttachmentRow, + create: bool, + ) -> StoreResult { + let path = self.attachment_path_metadata(account, task_uuid, row)?; + self.prepare_attachment_path(&path, create)?; + Ok(path) + } + + fn prepare_attachment_path(&self, path: &Path, create: bool) -> StoreResult<()> { + let relative = path + .strip_prefix(self.paths.root()) + .map_err(|_| StoreError::UnsafePath)?; + let mut checked = self.paths.root().to_owned(); + verify_private_path(&checked, true)?; + let components = relative.components().collect::>(); + for component in &components[..components.len() - 1] { + checked.push(component); + if create { + private_directory(&checked)?; + } else { + verify_private_path(&checked, true).map_err(missing_file)?; + } + } + Ok(()) + } + + pub(crate) fn collect_unbound_attachments_step( + &self, + ) -> StoreResult { + let after = self + .collection_cursors + .lock() + .map_err(|_| StoreError::Busy)? + .unbound + .clone(); + let candidate = self.read_bulk(|tx| { + let query = |after: Option<&(i64, String)>| { + tx.query_row("SELECT a.public_id,a.operation_id,a.size,a.account_id,t.uuid,a.updated_ms FROM file_attachments a JOIN tasks t ON t.id=a.task_id JOIN accounts c ON c.id=a.account_id WHERE a.bound=0 AND ((a.state='cancelled' AND a.size>0) OR a.updated_ms(?2,?3)) ORDER BY a.updated_ms,a.public_id LIMIT 1",params![pi_ai::now_ms().saturating_sub(STAGING_TTL_MS),after.map(|key|key.0),after.map(|key|key.1.as_str())],|row|Ok((AccountId(row.get(3)?),row.get::<_,String>(4)?,row.get::<_,i64>(5)?,attachment_cleanup_row(row)?))).optional() + }; + let candidate = query(after.as_ref())?; + Ok(if candidate.is_none() && after.is_some() { query(None)? } else { candidate }) + })?; + // Advance before touching the owned tree: an unsafe or inaccessible + // candidate remains durable, but cannot monopolize every future tick. + self.collection_cursors + .lock() + .map_err(|_| StoreError::Busy)? + .unbound = candidate + .as_ref() + .map(|(_, _, updated, row)| (*updated, row.id.clone())); + let Some((account, task, _, row)) = candidate else { + return Ok(AttachmentCollectionProgress { + done: true, + ..Default::default() + }); + }; + self.collect_attachment_file(account, &task, &row, false) + } + + pub(super) fn collect_task_attachments_step( + &self, + account: AccountId, + task_uuid: &str, + ) -> StoreResult { + let candidate=self.read_bulk(|tx| { + let task=deleting_task(tx,account,task_uuid)?; + Ok(tx.query_row("SELECT public_id,operation_id,size FROM file_attachments WHERE task_id=?1 LIMIT 1",[task],attachment_cleanup_row).optional()?) + })?; + let Some(row) = candidate else { + return Ok(AttachmentCollectionProgress { + done: true, + ..Default::default() + }); + }; + self.collect_attachment_file(account, task_uuid, &row, true) + } + + fn collect_attachment_file( + &self, + account: AccountId, + task_uuid: &str, + row: &AttachmentCleanupRow, + deleting: bool, + ) -> StoreResult { + let _lease = match self + .attachment_locks + .acquire(operation_key(account, &row.operation)) + { + Ok(value) => value, + Err(StoreError::Busy) => return Ok(AttachmentCollectionProgress::default()), + Err(error) => return Err(error), + }; + self.read(|tx| { + let task=if deleting {deleting_task(tx,account,task_uuid)?} else {active_task(tx,account,task_uuid)?}; + let collectable:bool=tx.query_row("SELECT EXISTS(SELECT 1 FROM file_attachments WHERE task_id=?1 AND public_id=?2 AND (?3 OR (bound=0 AND ((state='cancelled' AND size>0) OR updated_ms {} + Err(StoreError::NotFound) => { + missing = true; + break; + } + Err(error) => return Err(error), + } + } + let cleanup = if missing { + super::attachment_cleanup::AttachmentCleanupProgress { + done: true, + entries: 0, + } + } else { + super::attachment_cleanup::collect_owned_attachment_tree_step(&directory)? + }; + if !cleanup.done { + return Ok(AttachmentCollectionProgress { + files: cleanup.entries, + ..Default::default() + }); + } + self.write(Durability::Durable, |tx| { + let task = if deleting { + deleting_task(tx, account, task_uuid)? + } else { + active_task(tx, account, task_uuid)? + }; + let (state, updated): (String, i64) = tx.query_row( + "SELECT state,updated_ms FROM file_attachments WHERE public_id=?1", + [&row.id], + |row| Ok((row.get(0)?, row.get(1)?)), + )?; + let links = if !deleting && state == "cancelled" { + // Keep the cancellation identity after deleting its bytes, so + // a delayed sender cannot resurrect the cancelled operation. + let expired = updated < pi_ai::now_ms().saturating_sub(STAGING_TTL_MS); + if expired { + tx.execute( + "DELETE FROM file_attachments WHERE task_id=?1 AND public_id=?2", + params![task, row.id], + )? + } else { + tx.execute( + "UPDATE file_attachments SET size=0 WHERE task_id=?1 AND public_id=?2", + params![task, row.id], + )? + } + } else { + tx.execute( + "DELETE FROM file_attachments WHERE task_id=?1 AND public_id=?2", + params![task, row.id], + )? + }; + Ok(AttachmentCollectionProgress { + links, + files: cleanup.entries, + bytes: row.size as usize, + ..Default::default() + }) + }) + } +} + +fn operation_key(account: AccountId, operation: &str) -> String { + format!("{}:{operation}", account.0) +} +fn validate_operation(value: &str) -> StoreResult<()> { + if value.is_empty() + || value.len() > 128 + || !value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b)) + { + Err(StoreError::InvalidInput) + } else { + Ok(()) + } +} +fn validate_id(value: &str) -> StoreResult<()> { + if value.len() == 32 + && value + .bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) + { + Ok(()) + } else { + Err(StoreError::InvalidInput) + } +} +fn safe_component(value: &str) -> bool { + !value.is_empty() + && value.len() <= 128 + && value != "." + && value != ".." + && value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b)) +} +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} +fn normalized_name(value: &str) -> String { + let last = value.rsplit(['/', '\\']).next().unwrap_or_default(); + let mut name = String::new(); + for c in last + .chars() + .filter(|c| !c.is_control() && !"<>:\"|?*".contains(*c)) + .skip_while(|c| c.is_whitespace()) + { + if name.len() + c.len_utf8() > 180 { + break; + } + name.push(c); + } + let mut name = name + .trim_end_matches(|c: char| c == '.' || c.is_whitespace()) + .to_owned(); + let stem = name + .split('.') + .next() + .unwrap_or_default() + .to_ascii_uppercase(); + if name.is_empty() + || name.eq_ignore_ascii_case(".part") + || matches!( + stem.as_str(), + "CON" + | "PRN" + | "AUX" + | "NUL" + | "COM1" + | "COM2" + | "COM3" + | "COM4" + | "COM5" + | "COM6" + | "COM7" + | "COM8" + | "COM9" + | "LPT1" + | "LPT2" + | "LPT3" + | "LPT4" + | "LPT5" + | "LPT6" + | "LPT7" + | "LPT8" + | "LPT9" + ) + { + name = format!("attachment-{name}"); + } + // Prefixing reserved stems must share the same byte budget as normal names. + // Truncation can expose Unicode whitespace at the end, so trim once more. + while name.len() > 180 { + name.pop(); + } + name.trim_end_matches(|c: char| c == '.' || c.is_whitespace()) + .to_owned() +} +fn normalized_metadata(value: &AttachmentUploadMetadata) -> StoreResult { + if value.size > MAX_ATTACHMENT_BYTES as u64 + || value.sha256.len() != 64 + || !value + .sha256 + .bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) + || value.mime_type.len() > 128 + || !value.mime_type.contains('/') + || !value + .mime_type + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b"/!#$&^_.+-".contains(&b)) + { + return Err(StoreError::InvalidInput); + } + Ok(AttachmentUploadMetadata { + name: normalized_name(&value.name), + ..value.clone() + }) +} +fn attachment_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + attachment_row_offset(row, 0) +} +fn attachment_row_offset( + row: &rusqlite::Row<'_>, + offset: usize, +) -> rusqlite::Result { + let attachment = AttachmentRow { + id: row.get(offset)?, + operation: row.get(offset + 1)?, + metadata: AttachmentUploadMetadata { + name: row.get(offset + 2)?, + mime_type: row.get(offset + 3)?, + size: { + let size = row.get::<_, i64>(offset + 4)?; + u64::try_from(size) + .map_err(|_| rusqlite::Error::IntegralValueOutOfRange(offset + 4, size))? + }, + sha256: row.get(offset + 5)?, + }, + state: row.get(offset + 6)?, + }; + if validate_id(&attachment.id).is_err() + || validate_operation(&attachment.operation).is_err() + || normalized_metadata(&attachment.metadata).is_err() + || normalized_name(&attachment.metadata.name) != attachment.metadata.name + || !matches!( + attachment.state.as_str(), + "uploading" | "ready" | "cancelled" + ) + { + return Err(rusqlite::Error::FromSqlConversionFailure( + offset, + rusqlite::types::Type::Text, + Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "invalid attachment metadata", + )), + )); + } + Ok(attachment) +} +fn active_task(tx: &Transaction<'_>, account: AccountId, uuid: &str) -> StoreResult { + Ok(tx.query_row("SELECT t.id FROM tasks t JOIN accounts a ON a.id=t.account_id WHERE t.uuid=?1 AND t.account_id=?2 AND t.deleted_ms IS NULL AND a.deleting_ms IS NULL",params![uuid,account.0],|row|row.get(0))?) +} +fn deleting_task(tx: &Transaction<'_>, account: AccountId, uuid: &str) -> StoreResult { + Ok(tx.query_row("SELECT t.id FROM tasks t JOIN accounts a ON a.id=t.account_id WHERE t.uuid=?1 AND t.account_id=?2 AND t.deleted_ms IS NOT NULL AND (a.deletion_phase IS NULL OR a.deletion_phase='confirmed')",params![uuid,account.0],|row|row.get(0))?) +} +fn lookup_ready(tx: &Transaction<'_>, task: i64, id: &str) -> StoreResult { + Ok(tx.query_row("SELECT public_id,operation_id,name,mime,size,sha256,state FROM file_attachments WHERE task_id=?1 AND public_id=?2 AND state='ready'",params![task,id],attachment_row)?) +} +fn verify_upload_row( + tx: &Transaction<'_>, + task: i64, + row: &AttachmentRow, + state: &str, +) -> StoreResult<()> { + let actual=tx.query_row("SELECT public_id,operation_id,name,mime,size,sha256,state FROM file_attachments WHERE task_id=?1 AND public_id=?2",params![task,row.id],attachment_row)?; + if actual.operation != row.operation || actual.metadata != row.metadata || actual.state != state + { + Err(StoreError::Conflict) + } else { + Ok(()) + } +} +fn pin_tx(tx: &Transaction<'_>, task: i64, ids: &[String]) -> StoreResult<()> { + for id in ids { + validate_id(id)?; + lookup_ready(tx, task, id)?; + } + for id in ids { + tx.execute( + "UPDATE file_attachments SET bound=1 WHERE task_id=?1 AND public_id=?2", + params![task, id], + )?; + } + Ok(()) +} +fn missing_file(error: StoreError) -> StoreError { + match error { + StoreError::Io { + kind: std::io::ErrorKind::NotFound, + } => StoreError::NotFound, + other => other, + } +} +fn checked_file(path: &Path, size: u64) -> StoreResult { + checked_file_access(path, size, false) +} +fn checked_file_access(path: &Path, size: u64, writable: bool) -> StoreResult { + verify_private_path(path, false).map_err(missing_file)?; + let mut options = OpenOptions::new(); + options.read(true).write(writable); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC | libc::O_NONBLOCK); + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + options.custom_flags(windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT); + } + let file = options + .open(path) + .map_err(StoreError::from) + .map_err(missing_file)?; + let meta = file.metadata()?; + if !meta.is_file() || meta.len() != size { + return Err(StoreError::InvalidData); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + let named = fs::symlink_metadata(path)?; + if meta.dev() != named.dev() || meta.ino() != named.ino() || meta.nlink() != 1 { + return Err(StoreError::UnsafePath); + } + } + verify_private_path(path, false)?; + Ok(file) +} +fn verify_file_digest(path: &Path, metadata: &AttachmentUploadMetadata) -> StoreResult<()> { + let mut file = checked_file_access(path, metadata.size, true)?; + let mut hash = Sha256::new(); + let mut buffer = vec![0; MAX_ATTACHMENT_RANGE_BYTES]; + loop { + let read = file.read(&mut buffer)?; + if read == 0 { + break; + } + hash.update(&buffer[..read]); + } + if hex(&hash.finalize()) != metadata.sha256 { + return Err(StoreError::InvalidData); + } + super::keys::sync_key_file(&file)?; + sync_directory(path.parent().ok_or(StoreError::UnsafePath)?)?; + Ok(()) +} +fn publish_file(partial: &Path, published: &Path) -> StoreResult<()> { + if published.try_exists()? { + return Err(StoreError::Conflict); + } + super::keys::replace_key_file(partial, published)?; + sync_directory(published.parent().ok_or(StoreError::UnsafePath)?)?; + Ok(()) +} +fn remove_private_file_if_present(path: &Path) -> StoreResult<()> { + match verify_private_path(path, false) { + Ok(()) => fs::remove_file(path).map_err(StoreError::from), + Err(StoreError::Io { + kind: std::io::ErrorKind::NotFound, + }) => Ok(()), + Err(error) => Err(error), + } +} +fn image_extension(mime: &str) -> &'static str { + match mime { + "image/png" => "png", + "image/jpeg" | "image/jpg" => "jpg", + "image/webp" => "webp", + "image/gif" => "gif", + _ => "image", + } +} +fn detected_image_mime(bytes: &[u8]) -> Option<&'static str> { + if bytes.starts_with(b"\x89PNG\r\n\x1a\n") { + Some("image/png") + } else if bytes.starts_with(&[0xff, 0xd8, 0xff]) { + Some("image/jpeg") + } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" { + Some("image/webp") + } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") { + Some("image/gif") + } else { + None + } +} + +fn parse_upload(upload: &ImageUpload) -> StoreResult<(AttachmentUploadMetadata, Vec)> { + let (header, data) = upload + .data_url + .split_once(',') + .ok_or(StoreError::InvalidInput)?; + let mime = header + .strip_prefix("data:") + .and_then(|h| h.strip_suffix(";base64")) + .ok_or(StoreError::InvalidInput)?; + if !mime.starts_with("image/") + || data.is_empty() + || data.len() > MAX_UPLOAD_IMAGE_BYTES.div_ceil(3) * 4 + { + return Err(StoreError::InvalidInput); + } + let bytes = STANDARD + .decode(data) + .map_err(|_| StoreError::InvalidInput)?; + if bytes.len() > MAX_UPLOAD_IMAGE_BYTES { + return Err(StoreError::InvalidInput); + } + let detected = detected_image_mime(&bytes).ok_or(StoreError::InvalidInput)?; + if mime != detected && !(mime == "image/jpg" && detected == "image/jpeg") { + return Err(StoreError::InvalidInput); + } + let metadata = AttachmentUploadMetadata { + name: upload.name.clone(), + mime_type: detected.into(), + size: bytes.len() as u64, + sha256: hex(&Sha256::digest(&bytes)), + }; + Ok((normalized_metadata(&metadata)?, bytes)) +} + +/// Stored source is reference-only. Unlike incoming provider content, a raw +/// image in a v4 source entry is a corrupt storage invariant, not an upload. +pub(super) fn validate_stored_entry_images(entry: &SessionEntry) -> StoreResult<()> { + for image in entry_images(entry) { + if image.data == UNAVAILABLE_IMAGE_SOURCE { + continue; + } + let id = image + .data + .strip_prefix(ATTACHMENT_SOURCE_PREFIX) + .ok_or(StoreError::InvalidData)?; + validate_id(id).map_err(|_| StoreError::InvalidData)?; + } + Ok(()) +} + +pub(super) fn entry_image_sizes_tx( + tx: &Transaction<'_>, + task: super::tasks::TaskKey, + entry: &SessionEntry, +) -> StoreResult> { + entry_images(entry) + .map(|image| { + let size = if image.data == UNAVAILABLE_IMAGE_SOURCE { + 0 + } else { + let id = image + .data + .strip_prefix(ATTACHMENT_SOURCE_PREFIX) + .ok_or(StoreError::InvalidData)?; + validate_id(id).map_err(|_| StoreError::InvalidData)?; + lookup_ready(tx, task.0, id)?.metadata.size as usize + }; + size.div_ceil(3) + .checked_mul(4) + .ok_or(StoreError::InvalidData) + }) + .collect() +} + +fn entry_images(entry: &SessionEntry) -> impl Iterator { + let content = match &entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::User(message)), + } => Some(&message.content), + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(message)), + } => Some(&message.content), + EntryKind::Message { + message: SessionMessage::Custom(message), + } => Some(&message.content), + EntryKind::CustomMessage { content, .. } => Some(content), + EntryKind::ContextEdit { replacement, .. } => replacement.as_ref(), + _ => None, + }; + content + .into_iter() + .flatten() + .filter_map(|block| match block { + Content::Image(image) => Some(image), + _ => None, + }) +} + +fn visit_images( + entry: &mut SessionEntry, + mut visit: impl FnMut(&mut pi_ai::ImageContent) -> StoreResult<()>, +) -> StoreResult<()> { + let content = match &mut entry.kind { + EntryKind::Message { + message: SessionMessage::Llm(Message::User(message)), + } => Some(&mut message.content), + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(message)), + } => Some(&mut message.content), + EntryKind::Message { + message: SessionMessage::Custom(message), + } => Some(&mut message.content), + EntryKind::CustomMessage { content, .. } => Some(content), + EntryKind::ContextEdit { replacement, .. } => replacement.as_mut(), + _ => None, + }; + if let Some(content) = content { + for block in content { + if let Content::Image(image) = block { + visit(image)?; + } + } + } + Ok(()) +} + +#[cfg(test)] +#[path = "attachment_tests.rs"] +mod tests; diff --git a/apps/maple-agent/crates/maple-agent/src/store/bench.rs b/apps/maple-agent/crates/maple-agent/src/store/bench.rs index eaa2d8244..895dc105c 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/bench.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/bench.rs @@ -360,7 +360,7 @@ impl Store { } /// Deterministic planner guards with and without statistics. Fixed queries - /// exercise source, paging, projector merge/hide/plan and blob lookups. + /// exercise source, paging, projector merge/hide/plan and attachment metadata. pub fn bench_query_plans( &self, account: AccountId, @@ -396,8 +396,8 @@ impl Store { ("call", "timeline_rows_call", "SELECT row_key FROM timeline_rows WHERE task_id=?1 AND epoch=0 AND call_id='call' AND hidden=0 ORDER BY row_seq DESC LIMIT 1",task.0), ("hide", "timeline_rows_entry", "SELECT id FROM timeline_rows WHERE task_id=?1 AND epoch=0 AND entry_seq=2",task.0), ("plan", "timeline_rows_plan", "SELECT item FROM timeline_rows WHERE task_id=?1 AND epoch=0 AND is_plan=1 AND hidden=0 ORDER BY row_seq DESC LIMIT 1",task.0), - ("blob-links", "task_blobs", "SELECT task_id FROM task_blobs WHERE blob_id=?1",-1), - ("attachment-collection", "attachments_task_blob", "SELECT public_id FROM attachments WHERE task_id=?1 AND blob_id=1 LIMIT 1000",task.0), + ("attachment-unbound", "file_attachments_unbound", "SELECT public_id FROM file_attachments WHERE bound=0 AND updated_ms bool { }) { return false; } - // WITHOUT ROWID secondary indexes include the primary key (task_id, - // public_id). Either index can cover the selected public_id while seeking - // both equality constraints; a blob-only search would not be bounded by task. + // The payload is on disk. Collection must seek only this task's metadata; + // the index need not cover the selected public identifier. details.iter().any(|detail| { - ["attachments_task_blob", "attachments_blob"] - .iter() - .any(|index| { - let prefix = format!("SEARCH attachments USING COVERING INDEX {index} ("); - detail.strip_prefix(&prefix).is_some_and(|constraints| { - matches!( - constraints, - "task_id=? AND blob_id=?)" | "blob_id=? AND task_id=?)" - ) - }) - }) + detail == "SEARCH file_attachments USING INDEX file_attachments_task (task_id=?)" }) } @@ -635,28 +624,25 @@ mod tests { } #[test] - fn attachment_plan_accepts_both_covering_indexes_only_with_both_equalities() { + fn attachment_plan_requires_the_task_metadata_index() { let (_temp, store, account, _writers) = fixture(); - for index in ["attachments_task_blob", "attachments_blob"] { - for (columns, predicate, accepted) in [ - ("public_id", "task_id=?1 AND blob_id=1", true), - ("public_id", "task_id=?1", false), - ("public_id", "blob_id=?1", false), - ("created_ms", "task_id=?1 AND blob_id=1", false), - ] { - let details = store.read(|tx| { + for (predicate, accepted) in [ + ("task_id=?1", true), + ("task_id>?1", false), + ("account_id=?1", false), + ] { + let details = store.read(|tx| { let task = super::super::tasks::task_key_tx(tx, account, FIXTURE_A)?; let mut statement = tx.prepare(&format!( - "EXPLAIN QUERY PLAN SELECT {columns} FROM attachments INDEXED BY {index} WHERE {predicate} LIMIT 1000"))?; + "EXPLAIN QUERY PLAN SELECT public_id FROM file_attachments INDEXED BY file_attachments_task WHERE {predicate} LIMIT 1000"))?; Ok(statement.query_map([task.0], |row| row.get::<_, String>(3))? .collect::>>()?) }).unwrap(); - assert_eq!( - plan_matches("attachment-collection", "attachments_task_blob", &details), - accepted, - "{details:?}" - ); - } + assert_eq!( + plan_matches("attachment-collection", "file_attachments_task", &details), + accepted, + "{details:?}" + ); } } diff --git a/apps/maple-agent/crates/maple-agent/src/store/blobs.rs b/apps/maple-agent/crates/maple-agent/src/store/blobs.rs deleted file mode 100644 index 61907fc64..000000000 --- a/apps/maple-agent/crates/maple-agent/src/store/blobs.rs +++ /dev/null @@ -1,865 +0,0 @@ -//! Chunked new-store images. Digests and database ids never cross the host API. -//! -//! A digest lease spans both chunk preparation and the transaction that links -//! the bytes. Collectors take the same lease before deciding whether a blob is -//! shared. No filesystem attachment cache or previous application data is read. - -use std::collections::{BTreeMap, HashSet}; -use std::sync::{Arc, Condvar, Mutex}; - -use base64::{Engine as _, engine::general_purpose::STANDARD}; -use pi_ai::{Content, Message}; -use pi_coding_agent::{ - SessionMessage, - session::{EntryKind, SessionEntry}, -}; -use rusqlite::{OptionalExtension, Transaction, params}; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; - -use super::{AccountId, Durability, Store, StoreError, StoreResult}; - -pub const BLOB_CHUNK_BYTES: usize = 256 * 1024; -pub const MAX_ATTACHMENT_RANGE_BYTES: usize = 1024 * 1024; -pub const MAX_UPLOAD_IMAGE_BYTES: usize = 10 * 1024 * 1024; -pub const MAX_UPLOAD_IMAGES: usize = 10; -const CHUNKS_PER_WRITE: usize = MAX_ATTACHMENT_RANGE_BYTES / BLOB_CHUNK_BYTES; -const METADATA_ROWS_PER_COLLECTION: usize = 1000; -const BLOB_MARKER: &str = "maple-blob:"; -// Pi deliberately retains third-party images it cannot decode. Escape those -// strings rather than requiring them to meet the attachment-upload contract. -// Every such string is wrapped, including literal values of either marker. -const INLINE_IMAGE_MARKER: &str = "maple-inline-image:"; - -#[derive(Default)] -pub(crate) struct BlobLocks { - held: Mutex>, - changed: Condvar, - orphan_cursor: Mutex, -} - -impl BlobLocks { - fn acquire(self: &Arc, digest: [u8; 32]) -> StoreResult { - let mut held = self.held.lock().map_err(|_| StoreError::InvalidData)?; - while held.contains(&digest) { - held = self - .changed - .wait(held) - .map_err(|_| StoreError::InvalidData)?; - } - held.insert(digest); - Ok(BlobLease { - locks: Arc::clone(self), - digest, - }) - } -} - -struct BlobLease { - locks: Arc, - digest: [u8; 32], -} - -impl Drop for BlobLease { - fn drop(&mut self) { - // Poisoning cannot strand other waiters during shutdown/error unwinding. - let mut held = self - .locks - .held - .lock() - .unwrap_or_else(|error| error.into_inner()); - held.remove(&self.digest); - self.locks.changed.notify_all(); - } -} - -#[derive(Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct ImageUpload { - pub name: String, - pub data_url: String, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct StoredAttachment { - pub id: String, - pub name: String, - pub mime_type: String, - pub source: String, -} - -/// Raw bounded bytes; transports choose their own base64 encoding. -#[derive(Clone, PartialEq, Eq)] -pub struct AttachmentRange { - pub mime: String, - pub size: u64, - pub data: Vec, -} - -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub struct BlobCollectionProgress { - pub done: bool, - pub chunks: usize, - pub bytes: usize, - pub links: usize, -} - -/// Keep this value alive until the source-entry transaction has committed. Its -/// typed entry has only actual Pi image blocks encoded, never opaque JSON. -pub(super) struct PreparedEntryImages { - pub(super) entry: SessionEntry, - account: AccountId, - task_uuid: String, - blobs: Vec, - has_encoded_inline_images: bool, - _leases: Vec, -} - -/// A retry/write-all batch owns the union of its digest leases exactly once. -/// Preparing individual entries independently would deadlock when they share -/// an image, because the first prepared entry deliberately retains its lease. -pub(super) struct PreparedEntriesImages { - pub(super) entries: Vec, - _leases: Vec, -} - -impl PreparedEntryImages { - /// The entry flag means image data needs decoding, including escaped inline - /// strings which intentionally have no blob or task-blob link. - pub(super) fn has_images(&self) -> bool { - !self.blobs.is_empty() || self.has_encoded_inline_images - } - - pub(super) fn link(&self, transaction: &Transaction<'_>, task_id: i64) -> StoreResult<()> { - let actual = active_task(transaction, self.account, &self.task_uuid)?; - if actual != task_id { - return Err(StoreError::NotFound); - } - for id in &self.blobs { - if !transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM blobs WHERE id=?1 AND complete=1)", - [id], - |row| row.get::<_, bool>(0), - )? { - return Err(StoreError::InvalidData); - } - transaction.execute( - "INSERT INTO task_blobs(task_id,blob_id) VALUES(?1,?2) ON CONFLICT DO NOTHING", - params![task_id, id], - )?; - } - Ok(()) - } -} - -impl Store { - /// Validate the complete batch before writing. Public ids are random and - /// unrelated to digest, account or task identifiers. All links commit together. - pub fn store_image_uploads( - &self, - account: AccountId, - task_uuid: &str, - uploads: &[ImageUpload], - ) -> StoreResult> { - if uploads.len() > MAX_UPLOAD_IMAGES { - return Err(StoreError::InvalidInput); - } - self.read(|transaction| active_task(transaction, account, task_uuid))?; - let parsed = uploads - .iter() - .map(parse_upload) - .collect::>>()?; - let mut unique = BTreeMap::new(); - for image in &parsed { - unique.entry(image.digest).or_insert(image); - } - // A total ordering prevents deadlocks when two batches overlap. - let _leases = unique - .keys() - .map(|digest| self.blob_locks.acquire(*digest)) - .collect::>>()?; - let mut ids = BTreeMap::new(); - for (digest, image) in unique { - ids.insert(digest, self.prepare_blob(&image.bytes, image.mime, digest)?); - } - let attachments = parsed - .iter() - .map(|image| { - let id = format!("{:032x}", rand::random::()); - StoredAttachment { - source: format!("maple-attachment://{id}"), - id, - name: image.name.clone(), - mime_type: image.mime.to_string(), - } - }) - .collect::>(); - self.write(Durability::Durable, |transaction| { - let task = active_task(transaction, account, task_uuid)?; - for (image, attachment) in parsed.iter().zip(&attachments) { - let id = ids[&image.digest]; - transaction.execute("UPDATE blobs SET complete=1,mime=?2 WHERE id=?1", params![id,image.mime])?; - transaction.execute( - "INSERT INTO task_blobs(task_id,blob_id) VALUES(?1,?2) ON CONFLICT DO NOTHING", - params![task, id], - )?; - transaction.execute( - "INSERT INTO attachments(task_id,public_id,blob_id,created_ms) VALUES(?1,?2,?3,?4)", - params![task, attachment.id, id, pi_ai::now_ms()], - )?; - } - Ok(()) - })?; - Ok(attachments) - } - - /// Metadata authorization and chunk reads share one short snapshot. No caller - /// may read a digest/id directly, even if the same bytes belong to its account. - pub fn read_attachment( - &self, - account: AccountId, - task_uuid: &str, - public_id: &str, - offset: u64, - len: usize, - ) -> StoreResult { - if public_id.len() != 32 - || !public_id.bytes().all(|byte| byte.is_ascii_hexdigit()) - || len > MAX_ATTACHMENT_RANGE_BYTES - { - return Err(StoreError::InvalidInput); - } - self.read(|transaction| { - let task = active_task(transaction, account, task_uuid)?; - let (id, mime, size): (i64, String, i64) = transaction.query_row( - "SELECT b.id,b.mime,b.size FROM attachments a JOIN blobs b ON b.id=a.blob_id \ - JOIN task_blobs tb ON tb.task_id=a.task_id AND tb.blob_id=b.id \ - WHERE a.task_id=?1 AND a.public_id=?2 AND b.complete=1", - params![task, public_id], - |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), - )?; - let size = u64::try_from(size).map_err(|_| StoreError::InvalidData)?; - let data = read_range(transaction, id, size, offset, len)?; - Ok(AttachmentRange { mime, size, data }) - }) - } - - /// Prepare chunks without changing source data. Source insertion and `link` - /// must be the same Durable transaction; dropping on error leaves only an - /// unlinked blob, safely collected on next startup. - pub(super) fn prepare_entry_images( - &self, - account: AccountId, - task_uuid: &str, - entry: &SessionEntry, - ) -> StoreResult { - let mut prepared = - self.prepare_entries_images(account, task_uuid, std::slice::from_ref(entry))?; - let mut entry = prepared.entries.pop().ok_or(StoreError::InvalidData)?; - entry._leases = prepared._leases; - Ok(entry) - } - - pub(super) fn prepare_entries_images( - &self, - account: AccountId, - task_uuid: &str, - entries: &[SessionEntry], - ) -> StoreResult { - self.read(|transaction| active_task(transaction, account, task_uuid))?; - let mut images = BTreeMap::new(); - let mut entries = entries.to_vec(); - for entry in &mut entries { - visit_images(entry, |image| { - let Some(bytes) = blob_image_bytes(image) else { - return Ok(()); - }; - let digest: [u8; 32] = Sha256::digest(&bytes).into(); - images - .entry(digest) - .or_insert((bytes, image.mime_type.clone())); - Ok(()) - })?; - } - let leases = images - .keys() - .map(|digest| self.blob_locks.acquire(*digest)) - .collect::>>()?; - let mut ids = BTreeMap::new(); - for (digest, (bytes, mime)) in &images { - let id = self.prepare_blob(bytes, mime, *digest)?; - // Completion is not externally observable until link; the source - // commit's FULL sync covers these preceding chunk transactions. - self.write(Durability::Normal, |transaction| { - transaction.execute("UPDATE blobs SET complete=1 WHERE id=?1", [id])?; - Ok(()) - })?; - ids.insert(*digest, id); - } - let mut prepared = Vec::with_capacity(entries.len()); - for mut entry in entries { - let mut entry_ids = HashSet::new(); - let mut has_encoded_inline_images = false; - visit_images(&mut entry, |image| { - let Some(bytes) = blob_image_bytes(image) else { - has_encoded_inline_images = true; - image.data = format!( - "{INLINE_IMAGE_MARKER}{}", - STANDARD.encode(image.data.as_bytes()) - ); - return Ok(()); - }; - let digest: [u8; 32] = Sha256::digest(&bytes).into(); - let id = ids[&digest]; - entry_ids.insert(id); - image.data = format!("{BLOB_MARKER}{id}"); - Ok(()) - })?; - prepared.push(PreparedEntryImages { - entry, - account, - task_uuid: task_uuid.to_string(), - blobs: entry_ids.into_iter().collect(), - has_encoded_inline_images, - _leases: Vec::new(), - }); - } - Ok(PreparedEntriesImages { - entries: prepared, - _leases: leases, - }) - } - - /// Model-context loading only. Timeline pages never call this method. - pub(super) fn rehydrate_entry_images( - &self, - account: AccountId, - task_uuid: &str, - entry: &mut SessionEntry, - ) -> StoreResult<()> { - // Work on a clone so a missing/corrupt later image cannot leave a partly - // decoded entry in the caller's session. - let mut restored = entry.clone(); - let mut cache = BTreeMap::::new(); - visit_images(&mut restored, |image| { - if let Some(data) = decode_inline_image(&image.data)? { - // Decode exactly one layer: the original is allowed to look - // like either internal marker and must remain opaque. - image.data = data; - return Ok(()); - } - let Some(marker) = image.data.strip_prefix(BLOB_MARKER) else { - return Ok(()); - }; - let id = marker - .parse::() - .ok() - .filter(|id| *id > 0) - .ok_or(StoreError::InvalidData)?; - if let Some(data) = cache.get(&id) { - image.data.clone_from(data); - return Ok(()); - } - let size = self.read_bulk(|transaction| { - let task = active_task(transaction, account, task_uuid)?; - linked_blob_size(transaction, task, id) - })?; - let mut bytes = Vec::new(); - let mut offset = 0; - while offset < size { - // Each range releases its snapshot before base64 encoding - // and the next read, allowing checkpoints/foreground writes. - let data = self.read_bulk(|transaction| { - let task = active_task(transaction, account, task_uuid)?; - if linked_blob_size(transaction, task, id)? != size { - return Err(StoreError::InvalidData); - } - read_range(transaction, id, size, offset, MAX_ATTACHMENT_RANGE_BYTES) - })?; - offset += data.len() as u64; - bytes.extend(data); - } - image.data = STANDARD.encode(bytes); - cache.insert(id, image.data.clone()); - Ok(()) - })?; - *entry = restored; - Ok(()) - } - - /// One startup/idle orphan collection step. Upload and source preparation - /// hold the same digest lease through their link commit, so the final link - /// recheck also protects an upload which was in progress during discovery. - /// Metadata is removed after its chunks, keeping the final cascade bounded. - pub(crate) fn collect_orphan_blobs_step(&self) -> StoreResult { - // Scan each linked prefix once, at most 1,000 metadata rows per read - // snapshot. A restart simply starts from zero and safely rediscovers any - // partially collected row. The host runs one resumable pass per open. - let mut cursor = self - .blob_locks - .orphan_cursor - .lock() - .map_err(|_| StoreError::InvalidData)?; - let (candidate, last, scanned) = self.read_bulk(|transaction| { - let mut statement = transaction.prepare_cached( - "SELECT b.id,b.sha256, \ - NOT EXISTS(SELECT 1 FROM task_blobs tb WHERE tb.blob_id=b.id) \ - AND NOT EXISTS(SELECT 1 FROM attachments a WHERE a.blob_id=b.id) \ - FROM blobs b WHERE b.id>=?1 ORDER BY b.id LIMIT ?2", - )?; - let mut rows = - statement.query(params![*cursor, METADATA_ROWS_PER_COLLECTION as i64])?; - let mut last = *cursor; - let mut scanned = 0; - while let Some(row) = rows.next()? { - last = row.get(0)?; - scanned += 1; - if row.get::<_, bool>(2)? { - return Ok((Some((last, row.get::<_, Vec>(1)?)), last, scanned)); - } - } - Ok((None, last, scanned)) - })?; - let Some((id, digest)) = candidate else { - let done = scanned < METADATA_ROWS_PER_COLLECTION; - *cursor = if done { 0 } else { last.saturating_add(1) }; - return Ok(BlobCollectionProgress { - done, - ..Default::default() - }); - }; - *cursor = id; - let digest = digest.try_into().map_err(|_| StoreError::InvalidData)?; - let _lease = self.blob_locks.acquire(digest)?; - let mut removed = false; - let progress = self.write_idle(|transaction| { - // INTEGER PRIMARY KEY ids can be reused after another collector - // removes the candidate while we wait for its digest lease. - if !transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM blobs WHERE id=?1 AND sha256=?2)", - params![id,digest.as_slice()],|row| row.get::<_,bool>(0), - )? { return Ok(BlobCollectionProgress::default()); } - let linked: bool = transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM task_blobs WHERE blob_id=?1) OR EXISTS(SELECT 1 FROM attachments WHERE blob_id=?1)", - [id], |row| row.get(0), - )?; - if linked { return Ok(BlobCollectionProgress::default()); } - transaction.execute("UPDATE blobs SET complete=0 WHERE id=?1", [id])?; - let mut progress = delete_chunks(transaction,id)?; - if progress.chunks == 0 { - transaction.execute("DELETE FROM blobs WHERE id=?1", [id])?; - removed = true; - } - progress.done = false; - Ok(progress) - })?; - if removed { - *cursor = id.saturating_add(1); - } - Ok(progress) - } - - /// Task collector hook; the task must already be durably hidden. It also - /// precollects attachments, preventing a large final task cascade. - pub(super) fn collect_task_blobs_step( - &self, - account: AccountId, - task_uuid: &str, - ) -> StoreResult { - let (task, candidate) = self.read_bulk(|transaction| { - let task = deleting_task(transaction, account, task_uuid)?; - let candidate = transaction - .query_row( - "SELECT b.id,b.sha256 FROM blobs b WHERE b.id=coalesce( \ - (SELECT blob_id FROM task_blobs WHERE task_id=?1 LIMIT 1), \ - (SELECT blob_id FROM attachments WHERE task_id=?1 LIMIT 1))", - [task], - |row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec>(1)?)), - ) - .optional()?; - Ok((task, candidate)) - })?; - let Some((id, digest)) = candidate else { - return Ok(BlobCollectionProgress { - done: true, - ..Default::default() - }); - }; - let digest = digest.try_into().map_err(|_| StoreError::InvalidData)?; - let _lease = self.blob_locks.acquire(digest)?; - self.write_idle(|transaction| { - if deleting_task(transaction, account, task_uuid)? != task { - return Err(StoreError::NotFound); - } - if !transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM blobs b WHERE b.id=?1 AND b.sha256=?2 AND \ - (EXISTS(SELECT 1 FROM task_blobs tb WHERE tb.task_id=?3 AND tb.blob_id=b.id) \ - OR EXISTS(SELECT 1 FROM attachments a WHERE a.task_id=?3 AND a.blob_id=b.id)))", - params![id, digest.as_slice(), task], - |row| row.get::<_, bool>(0), - )? { - return Ok(BlobCollectionProgress::default()); - } - let links = transaction.execute( - "DELETE FROM attachments WHERE task_id=?1 AND public_id IN \ - (SELECT public_id FROM attachments WHERE task_id=?1 AND blob_id=?2 LIMIT ?3)", - params![task, id, METADATA_ROWS_PER_COLLECTION as i64], - )?; - // Metadata batches and chunk batches are separate transactions: a - // 1 MiB chunk delete must not also carry 1,000 attachment records. - if links > 0 { - return Ok(BlobCollectionProgress { - links, - ..Default::default() - }); - } - if transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM attachments WHERE task_id=?1 AND blob_id=?2)", - params![task, id], - |row| row.get::<_, bool>(0), - )? { - return Ok(BlobCollectionProgress { - links, - ..Default::default() - }); - } - let shared: bool = transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM task_blobs WHERE blob_id=?1 AND task_id<>?2) \ - OR EXISTS(SELECT 1 FROM attachments WHERE blob_id=?1 AND task_id<>?2)", - params![id, task], - |row| row.get(0), - )?; - if shared { - let links = links - + transaction.execute( - "DELETE FROM task_blobs WHERE task_id=?1 AND blob_id=?2", - params![task, id], - )?; - return Ok(BlobCollectionProgress { - links, - ..Default::default() - }); - } - transaction.execute("UPDATE blobs SET complete=0 WHERE id=?1", [id])?; - let mut progress = delete_chunks(transaction, id)?; - progress.links = links; - if progress.chunks == 0 { - progress.links += transaction.execute( - "DELETE FROM task_blobs WHERE task_id=?1 AND blob_id=?2", - params![task, id], - )?; - transaction.execute("DELETE FROM blobs WHERE id=?1", [id])?; - } - Ok(progress) - }) - } - - /// Caller owns the digest lease. Incomplete rows left by an interrupted - /// preparation are rewritten in capped transactions, never trusted blindly. - fn prepare_blob(&self, bytes: &[u8], mime: &str, digest: [u8; 32]) -> StoreResult { - let size = i64::try_from(bytes.len()).map_err(|_| StoreError::InvalidInput)?; - let (id,complete) = self.write(Durability::Normal,|transaction| { - let existing = transaction.query_row("SELECT id,complete,size FROM blobs WHERE sha256=?1",[digest.as_slice()], - |row| Ok((row.get::<_,i64>(0)?,row.get::<_,bool>(1)?,row.get::<_,i64>(2)?))).optional()?; - if let Some((id,complete,stored_size)) = existing { - if stored_size != size { return Err(StoreError::InvalidData); } - if !complete && transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM task_blobs WHERE blob_id=?1) OR EXISTS(SELECT 1 FROM attachments WHERE blob_id=?1)", - [id],|row| row.get::<_,bool>(0), - )? { return Err(StoreError::Busy); } - return Ok((id,complete)); - } - transaction.execute("INSERT INTO blobs(sha256,mime,size,complete,created_ms) VALUES(?1,?2,?3,0,?4)", - params![digest.as_slice(),mime,size,pi_ai::now_ms()])?; - Ok((transaction.last_insert_rowid(),false)) - })?; - if complete { - return Ok(id); - } - loop { - if self - .write(Durability::Normal, |transaction| { - delete_chunks(transaction, id) - })? - .chunks - == 0 - { - break; - } - } - for (batch, chunks) in bytes - .chunks(BLOB_CHUNK_BYTES) - .collect::>() - .chunks(CHUNKS_PER_WRITE) - .enumerate() - { - self.write(Durability::Normal, |transaction| { - for (index, chunk) in chunks.iter().enumerate() { - let index = i64::try_from(batch * CHUNKS_PER_WRITE + index) - .map_err(|_| StoreError::InvalidInput)?; - transaction.execute( - "INSERT INTO blob_chunks(blob_id,idx,data) VALUES(?1,?2,?3)", - params![id, index, *chunk], - )?; - } - Ok(()) - })?; - } - Ok(id) - } -} - -fn active_task(transaction: &Transaction<'_>, account: AccountId, uuid: &str) -> StoreResult { - Ok(transaction.query_row( - "SELECT t.id FROM tasks t JOIN accounts a ON a.id=t.account_id \ - WHERE t.uuid=?1 AND t.account_id=?2 AND t.deleted_ms IS NULL AND a.deleting_ms IS NULL", - params![uuid, account.0], - |row| row.get(0), - )?) -} - -fn deleting_task( - transaction: &Transaction<'_>, - account: AccountId, - uuid: &str, -) -> StoreResult { - Ok(transaction.query_row( - "SELECT t.id FROM tasks t JOIN accounts a ON a.id=t.account_id \ - WHERE t.uuid=?1 AND t.account_id=?2 AND t.deleted_ms IS NOT NULL \ - AND (a.deletion_phase IS NULL OR a.deletion_phase='confirmed')", - params![uuid, account.0], - |row| row.get(0), - )?) -} - -fn linked_blob_size(transaction: &Transaction<'_>, task: i64, id: i64) -> StoreResult { - let size: i64 = transaction.query_row( - "SELECT b.size FROM blobs b JOIN task_blobs tb ON tb.blob_id=b.id \ - WHERE b.id=?1 AND tb.task_id=?2 AND b.complete=1", - params![id, task], - |row| row.get(0), - )?; - u64::try_from(size).map_err(|_| StoreError::InvalidData) -} - -/// Metadata-only projector support, in typed image-block order. The caller -/// already authorized the TaskKey in this same transaction. No image bytes are -/// read and no base64 string is fabricated for history rendering/rebuilds. -pub(super) fn entry_image_sizes_tx( - transaction: &Transaction<'_>, - task: super::tasks::TaskKey, - entry: &SessionEntry, -) -> StoreResult> { - let mut entry = entry.clone(); - let mut sizes = Vec::new(); - visit_images(&mut entry, |image| { - let chars = if let Some(data) = decode_inline_image(&image.data)? { - data.len() - } else if let Some(marker) = image.data.strip_prefix(BLOB_MARKER) { - let id = marker - .parse::() - .ok() - .filter(|id| *id > 0) - .ok_or(StoreError::InvalidData)?; - let size = usize::try_from(linked_blob_size(transaction, task.0, id)?) - .map_err(|_| StoreError::InvalidData)?; - size.div_ceil(3) - .checked_mul(4) - .ok_or(StoreError::InvalidData)? - } else { - image.data.len() - }; - sizes.push(chars); - Ok(()) - })?; - Ok(sizes) -} - -fn delete_chunks(transaction: &Transaction<'_>, id: i64) -> StoreResult { - let chunks = { - let mut statement = transaction.prepare_cached( - "SELECT idx,length(data) FROM blob_chunks WHERE blob_id=?1 ORDER BY idx LIMIT ?2", - )?; - statement - .query_map(params![id, CHUNKS_PER_WRITE as i64], |row| { - Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)) - })? - .collect::, _>>()? - }; - let mut progress = BlobCollectionProgress::default(); - for (index, size) in chunks { - let size = usize::try_from(size).map_err(|_| StoreError::InvalidData)?; - if size > BLOB_CHUNK_BYTES { - return Err(StoreError::InvalidData); - } - transaction.execute( - "DELETE FROM blob_chunks WHERE blob_id=?1 AND idx=?2", - params![id, index], - )?; - progress.chunks += 1; - progress.bytes += size; - } - Ok(progress) -} - -fn read_range( - transaction: &Transaction<'_>, - id: i64, - size: u64, - offset: u64, - len: usize, -) -> StoreResult> { - if offset > size || len > MAX_ATTACHMENT_RANGE_BYTES { - return Err(StoreError::InvalidInput); - } - let end = offset.saturating_add(len as u64).min(size); - if offset == end { - return Ok(Vec::new()); - } - let first = offset / BLOB_CHUNK_BYTES as u64; - let last = (end - 1) / BLOB_CHUNK_BYTES as u64; - let mut statement = transaction.prepare_cached( - "SELECT idx,data FROM blob_chunks WHERE blob_id=?1 AND idx BETWEEN ?2 AND ?3 ORDER BY idx", - )?; - let mut rows = statement.query(params![id, first as i64, last as i64])?; - let mut next = first; - let mut data = Vec::with_capacity((end - offset) as usize); - while let Some(row) = rows.next()? { - let index: i64 = row.get(0)?; - let bytes: Vec = row.get(1)?; - if index != next as i64 { - return Err(StoreError::InvalidData); - } - let start = next * BLOB_CHUNK_BYTES as u64; - let expected = (size - start).min(BLOB_CHUNK_BYTES as u64) as usize; - if bytes.len() != expected { - return Err(StoreError::InvalidData); - } - let from = offset.saturating_sub(start) as usize; - let to = (end - start).min(bytes.len() as u64) as usize; - data.extend_from_slice(&bytes[from..to]); - next += 1; - } - if next != last + 1 || data.len() != (end - offset) as usize { - return Err(StoreError::InvalidData); - } - Ok(data) -} - -fn blob_image_bytes(image: &pi_ai::ImageContent) -> Option> { - if image.mime_type.is_empty() { - return None; - } - let bytes = STANDARD.decode(&image.data).ok()?; - (!bytes.is_empty()).then_some(bytes) -} - -fn decode_inline_image(data: &str) -> StoreResult> { - data.strip_prefix(INLINE_IMAGE_MARKER) - .map(|encoded| { - let bytes = STANDARD - .decode(encoded) - .map_err(|_| StoreError::InvalidData)?; - String::from_utf8(bytes).map_err(|_| StoreError::InvalidData) - }) - .transpose() -} - -fn visit_images( - entry: &mut SessionEntry, - mut visit: impl FnMut(&mut pi_ai::ImageContent) -> StoreResult<()>, -) -> StoreResult<()> { - let content = match &mut entry.kind { - EntryKind::Message { - message: SessionMessage::Llm(Message::User(message)), - } => Some(&mut message.content), - EntryKind::Message { - message: SessionMessage::Llm(Message::ToolResult(message)), - } => Some(&mut message.content), - EntryKind::Message { - message: SessionMessage::Custom(message), - } => Some(&mut message.content), - EntryKind::CustomMessage { content, .. } => Some(content), - EntryKind::ContextEdit { replacement, .. } => replacement.as_mut(), - _ => None, - }; - if let Some(content) = content { - for block in content { - if let Content::Image(image) = block { - visit(image)?; - } - } - } - Ok(()) -} - -struct ParsedUpload { - name: String, - mime: &'static str, - bytes: Vec, - digest: [u8; 32], -} - -fn parse_upload(upload: &ImageUpload) -> StoreResult { - let (header, data) = upload - .data_url - .split_once(',') - .ok_or(StoreError::InvalidInput)?; - let declared = header - .strip_prefix("data:") - .and_then(|header| header.strip_suffix(";base64")) - .ok_or(StoreError::InvalidInput)?; - if data.is_empty() || data.len() > MAX_UPLOAD_IMAGE_BYTES.div_ceil(3) * 4 { - return Err(StoreError::InvalidInput); - } - let bytes = STANDARD - .decode(data) - .map_err(|_| StoreError::InvalidInput)?; - if bytes.len() > MAX_UPLOAD_IMAGE_BYTES { - return Err(StoreError::InvalidInput); - } - let mime = if bytes.starts_with(b"\x89PNG\r\n\x1a\n") { - "image/png" - } else if bytes.starts_with(&[0xff, 0xd8, 0xff]) { - "image/jpeg" - } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" { - "image/webp" - } else { - return Err(StoreError::InvalidInput); - }; - let declared = if declared.eq_ignore_ascii_case("image/jpg") { - "image/jpeg" - } else { - declared - }; - if declared != mime { - return Err(StoreError::InvalidInput); - } - let name = upload - .name - .rsplit(['/', '\\']) - .next() - .unwrap_or_default() - .trim() - .chars() - .filter(|character| !character.is_control()) - .take(255) - .collect::(); - let name = if name.is_empty() { - "image".to_string() - } else { - name - }; - let digest = Sha256::digest(&bytes).into(); - Ok(ParsedUpload { - name, - mime, - bytes, - digest, - }) -} - -#[cfg(test)] -#[path = "blobtests.rs"] -mod tests; diff --git a/apps/maple-agent/crates/maple-agent/src/store/blobtests.rs b/apps/maple-agent/crates/maple-agent/src/store/blobtests.rs deleted file mode 100644 index f4d851359..000000000 --- a/apps/maple-agent/crates/maple-agent/src/store/blobtests.rs +++ /dev/null @@ -1,633 +0,0 @@ -use std::sync::{Arc, Barrier, mpsc}; -use std::time::Duration; - -use super::*; -use crate::store::StorePaths; - -fn fixture() -> (tempfile::TempDir, Arc, AccountId, AccountId) { - let directory = tempfile::tempdir().unwrap(); - let store = Arc::new(Store::open(StorePaths::new(directory.path())).unwrap()); - let alice = store.upsert_account("alice", 1).unwrap(); - let bob = store.upsert_account("bob", 1).unwrap(); - create_task(&store, alice, "alice-task"); - create_task(&store, bob, "bob-task"); - (directory, store, alice, bob) -} - -fn create_task(store: &Store, account: AccountId, uuid: &str) -> i64 { - store.write(Durability::Durable, |transaction| { - transaction.execute( - "INSERT INTO projects(account_id,root,last_used_ms) VALUES(?1,?2,1) ON CONFLICT DO NOTHING", - params![account.0, format!("/fixture/{}",account.0)], - )?; - let project: i64 = transaction.query_row("SELECT id FROM projects WHERE account_id=?1",[account.0],|row| row.get(0))?; - transaction.execute( - "INSERT INTO tasks(uuid,account_id,project_id,kind,state,title,title_fold,created_ms,updated_ms) \ - VALUES(?1,?2,?3,'desktop','active','fixture','fixture',1,1)", - params![uuid,account.0,project], - )?; - Ok(transaction.last_insert_rowid()) - }).unwrap() -} - -fn png(size: usize) -> Vec { - let mut bytes = vec![17; size.max(8)]; - bytes[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); - bytes -} - -fn upload(bytes: &[u8]) -> ImageUpload { - ImageUpload { - name: "../photo.png".to_string(), - data_url: format!("data:image/png;base64,{}", STANDARD.encode(bytes)), - } -} - -fn count(store: &Store, table: &str) -> i64 { - assert!(["blobs", "blob_chunks", "attachments", "task_blobs"].contains(&table)); - store - .read(|transaction| { - Ok( - transaction.query_row(&format!("SELECT count(*) FROM {table}"), [], |row| { - row.get(0) - })?, - ) - }) - .unwrap() -} - -fn mark_deleted(store: &Store, uuid: &str) { - store - .write(Durability::Durable, |transaction| { - transaction.execute("UPDATE tasks SET deleted_ms=2 WHERE uuid=?1", [uuid])?; - Ok(()) - }) - .unwrap(); -} - -fn collect_task(store: &Store, account: AccountId, uuid: &str) -> usize { - for step in 0..100 { - let progress = store.collect_task_blobs_step(account, uuid).unwrap(); - assert!(progress.chunks <= CHUNKS_PER_WRITE); - assert!(progress.bytes <= MAX_ATTACHMENT_RANGE_BYTES); - assert!(progress.links <= METADATA_ROWS_PER_COLLECTION + 1); - if progress.done { - return step; - } - } - panic!("blob collector failed to finish"); -} - -#[test] -fn uploads_deduplicate_without_exposing_digest_or_crossing_accounts() { - let (_directory, store, alice, bob) = fixture(); - let bytes = png(2 * BLOB_CHUNK_BYTES + 73); - let first = store - .store_image_uploads(alice, "alice-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - let second = store - .store_image_uploads(bob, "bob-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - assert_ne!(first.id, second.id); - assert_eq!(first.id.len(), 32); - assert_eq!(first.name, "photo.png"); - assert_eq!(first.source, format!("maple-attachment://{}", first.id)); - assert!( - !first - .source - .contains(&format!("{:x}", Sha256::digest(&bytes))) - ); - assert_eq!(count(&store, "blobs"), 1); - assert_eq!(count(&store, "blob_chunks"), 3); - assert_eq!(count(&store, "task_blobs"), 2); - assert!(matches!( - store.read_attachment(bob, "alice-task", &first.id, 0, 1), - Err(StoreError::NotFound) - )); - assert!(matches!( - store.read_attachment(bob, "bob-task", &first.id, 0, 1), - Err(StoreError::NotFound) - )); - assert_eq!( - store - .read_attachment(alice, "alice-task", &first.id, 0, bytes.len()) - .unwrap() - .data, - bytes - ); -} - -#[test] -fn ranges_cross_chunk_boundaries_and_reassemble_exact_bytes() { - let (_directory, store, alice, _) = fixture(); - let mut bytes = png(3 * MAX_ATTACHMENT_RANGE_BYTES + 19); - for (index, byte) in bytes.iter_mut().enumerate().skip(8) { - *byte = (index % 251) as u8; - } - let image = store - .store_image_uploads(alice, "alice-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - let mut assembled = Vec::new(); - while assembled.len() < bytes.len() { - let read = store - .read_attachment( - alice, - "alice-task", - &image.id, - assembled.len() as u64, - MAX_ATTACHMENT_RANGE_BYTES - 7, - ) - .unwrap(); - assert_eq!(read.size, bytes.len() as u64); - assert_eq!(read.mime, "image/png"); - assembled.extend(read.data); - } - assert_eq!(assembled, bytes); - let crossing = store - .read_attachment( - alice, - "alice-task", - &image.id, - (BLOB_CHUNK_BYTES - 3) as u64, - 11, - ) - .unwrap(); - assert_eq!( - crossing.data, - bytes[BLOB_CHUNK_BYTES - 3..BLOB_CHUNK_BYTES + 8] - ); - assert!( - store - .read_attachment( - alice, - "alice-task", - &image.id, - 0, - MAX_ATTACHMENT_RANGE_BYTES + 1 - ) - .is_err() - ); - assert!( - store - .read_attachment(alice, "alice-task", &image.id, bytes.len() as u64 + 1, 1) - .is_err() - ); - assert!( - store - .read_attachment(alice, "alice-task", &image.id, bytes.len() as u64, 1) - .unwrap() - .data - .is_empty() - ); -} - -#[test] -fn upload_limits_validate_whole_batch_before_creating_any_blob() { - let (_directory, store, alice, _) = fixture(); - let okay = upload(&png(100)); - let bad = ImageUpload { - name: "bad.jpg".into(), - data_url: okay.data_url.replace("image/png", "image/jpeg"), - }; - assert!( - store - .store_image_uploads(alice, "alice-task", &[okay.clone(), bad]) - .is_err() - ); - assert_eq!(count(&store, "blobs"), 0); - assert!( - store - .store_image_uploads(alice, "alice-task", &vec![okay; MAX_UPLOAD_IMAGES + 1]) - .is_err() - ); - assert!( - store - .store_image_uploads( - alice, - "alice-task", - &[upload(&png(MAX_UPLOAD_IMAGE_BYTES + 1))] - ) - .is_err() - ); - assert_eq!(count(&store, "blobs"), 0); -} - -#[test] -fn all_existing_upload_formats_are_preserved() { - let (_directory, store, alice, _) = fixture(); - let jpeg = ImageUpload { - name: "a.jpg".into(), - data_url: format!( - "data:image/jpg;base64,{}", - STANDARD.encode([0xff, 0xd8, 0xff, 7]) - ), - }; - let webp = ImageUpload { - name: "a.webp".into(), - data_url: format!( - "data:image/webp;base64,{}", - STANDARD.encode(b"RIFF\x04\0\0\0WEBP") - ), - }; - let result = store - .store_image_uploads(alice, "alice-task", &[jpeg, webp]) - .unwrap(); - assert_eq!(result[0].mime_type, "image/jpeg"); - assert_eq!(result[1].mime_type, "image/webp"); -} - -fn tool_entry(bytes: &[u8]) -> SessionEntry { - serde_json::from_value(serde_json::json!({ - "id":"0123abcd","parentId":null,"timestamp":1,"type":"message", - "message":{"role":"toolResult","toolCallId":"call","toolName":"read_image", - "content":[{"type":"image","data":STANDARD.encode(bytes),"mimeType":"image/png"}], - "details":{"type":"image","data":"opaque-provider-payload","mimeType":"kept"}, - "isError":false,"timestamp":1} - })) - .unwrap() -} - -#[test] -fn inline_images_over_upload_limit_round_trip_exactly_without_touching_opaque_details() { - let (_directory, store, alice, bob) = fixture(); - let original = tool_entry(&png(20 * 1024 * 1024)); - let prepared = store - .prepare_entry_images(alice, "alice-task", &original) - .unwrap(); - assert!(prepared.has_images()); - assert_ne!(prepared.entry, original); - let encoded = serde_json::to_value(&prepared.entry).unwrap(); - assert_eq!( - encoded["message"]["details"]["data"], - "opaque-provider-payload" - ); - assert!( - encoded["message"]["content"][0]["data"] - .as_str() - .unwrap() - .starts_with(BLOB_MARKER) - ); - store - .write(Durability::Durable, |transaction| { - let task = active_task(transaction, alice, "alice-task")?; - prepared.link(transaction, task) - }) - .unwrap(); - let mut stored = prepared.entry.clone(); - drop(prepared); - assert!( - store - .rehydrate_entry_images(bob, "bob-task", &mut stored) - .is_err() - ); - store - .rehydrate_entry_images(alice, "alice-task", &mut stored) - .unwrap(); - assert_eq!(stored, original); -} - -#[test] -fn third_party_images_escape_losslessly_without_aliasing_internal_markers() { - let (_directory, store, alice, _) = fixture(); - let bytes = png(103); - let canonical = STANDARD.encode(&bytes); - // Existing unencoded inline data is still opaque to both decoding and - // metadata projection; a flag-zero source row does not need a rewrite. - let original_inline = tool_entry(&bytes); - let mut unchanged_inline = original_inline.clone(); - store - .rehydrate_entry_images(alice, "alice-task", &mut unchanged_inline) - .unwrap(); - assert_eq!(unchanged_inline, original_inline); - store - .read(|transaction| { - let task = active_task(transaction, alice, "alice-task")?; - assert_eq!( - entry_image_sizes_tx( - transaction, - super::super::tasks::TaskKey(task), - &original_inline, - )?, - vec![canonical.len()] - ); - Ok(()) - }) - .unwrap(); - let non_blob_images = [ - ("not base64!".to_string(), "image/png"), - (String::new(), "image/png"), - (format!("data:image/png;base64,{canonical}"), "image/png"), - ( - format!("{}\n{}\n", &canonical[..76], &canonical[76..]), - "image/png", - ), - (canonical.trim_end_matches('=').to_string(), "image/png"), - ( - STANDARD - .encode([0xfb, 0xff]) - .replace('+', "-") - .replace('/', "_"), - "image/png", - ), - (canonical.clone(), ""), - ("maple-blob:1".to_string(), "image/png"), - ("maple-blob:not-an-id".to_string(), "image/png"), - ("maple-inline-image:".to_string(), "image/png"), - ( - "maple-inline-image:bWFwbGUtYmxvYjox".to_string(), - "image/png", - ), - ("opaque image 🦊".to_string(), "image/png"), - ]; - for (data, mime) in non_blob_images { - for mixed in [false, true] { - let mut original = tool_entry(&bytes); - let EntryKind::Message { - message: SessionMessage::Llm(Message::ToolResult(result)), - } = &mut original.kind - else { - unreachable!() - }; - result.content = vec![Content::image(data.clone(), mime)]; - if mixed { - result.content.extend([ - Content::image(canonical.clone(), "image/png"), - Content::image(canonical.clone(), "image/png"), - ]); - } - let prepared = store - .prepare_entry_images(alice, "alice-task", &original) - .unwrap(); - assert!(prepared.has_images(), "escaped-only entries need decoding"); - assert!(prepared.has_encoded_inline_images); - assert_eq!(prepared.blobs.len(), usize::from(mixed)); - let mut expected_sizes = vec![data.len()]; - if mixed { - expected_sizes.extend([canonical.len(); 2]); - } - store - .write(Durability::Durable, |transaction| { - let task = active_task(transaction, alice, "alice-task")?; - prepared.link(transaction, task)?; - assert_eq!( - entry_image_sizes_tx( - transaction, - super::super::tasks::TaskKey(task), - &prepared.entry - )?, - expected_sizes - ); - Ok(()) - }) - .unwrap(); - let mut restored = prepared.entry.clone(); - drop(prepared); - store - .rehydrate_entry_images(alice, "alice-task", &mut restored) - .unwrap(); - assert_eq!(restored, original); - } - } - assert_eq!(count(&store, "blobs"), 1); - assert_eq!(count(&store, "task_blobs"), 1); -} - -#[test] -fn retry_batch_holds_shared_digest_once_and_projector_reads_metadata_only() { - let (_directory, store, alice, _) = fixture(); - let bytes = png(103); - let first = tool_entry(&bytes); - let mut second = first.clone(); - second.id = "0123abce".to_string(); - let originals = vec![first, second]; - let prepared = store - .prepare_entries_images(alice, "alice-task", &originals) - .unwrap(); - assert_eq!(prepared.entries.len(), 2); - store - .write(Durability::Durable, |transaction| { - let task = active_task(transaction, alice, "alice-task")?; - for entry in &prepared.entries { - entry.link(transaction, task)?; - } - Ok(()) - }) - .unwrap(); - assert_eq!(count(&store, "blobs"), 1); - assert_eq!(count(&store, "task_blobs"), 1); - let stored = prepared - .entries - .iter() - .map(|entry| entry.entry.clone()) - .collect::>(); - drop(prepared); - for (mut entry, original) in stored.iter().cloned().zip(originals) { - store - .rehydrate_entry_images(alice, "alice-task", &mut entry) - .unwrap(); - assert_eq!(entry, original); - } - // A projector requires metadata even if bytes have been damaged; it must - // not silently load/decode them just to compute the image summary length. - store - .write(Durability::Normal, |transaction| { - transaction.execute("DELETE FROM blob_chunks", [])?; - let task = super::super::tasks::TaskKey(active_task(transaction, alice, "alice-task")?); - let sizes = entry_image_sizes_tx(transaction, task, &stored[0])?; - assert_eq!(sizes, vec![STANDARD.encode(&bytes).len()]); - Ok(()) - }) - .unwrap(); - let mut damaged = stored[0].clone(); - assert!( - store - .rehydrate_entry_images(alice, "alice-task", &mut damaged) - .is_err() - ); - assert_eq!(damaged, stored[0]); -} - -#[test] -fn shared_blob_survives_task_collection_and_final_exclusive_collection_is_bounded() { - let (_directory, store, alice, bob) = fixture(); - let bytes = png(3 * MAX_ATTACHMENT_RANGE_BYTES + 1); - let first = store - .store_image_uploads(alice, "alice-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - let second = store - .store_image_uploads(bob, "bob-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - mark_deleted(&store, "alice-task"); - assert!( - store - .read_attachment(alice, "alice-task", &first.id, 0, 1) - .is_err() - ); - collect_task(&store, alice, "alice-task"); - assert_eq!(count(&store, "blobs"), 1); - assert_eq!( - store - .read_attachment(bob, "bob-task", &second.id, 0, 1) - .unwrap() - .data, - bytes[..1] - ); - mark_deleted(&store, "bob-task"); - assert!(collect_task(&store, bob, "bob-task") >= 5); - assert_eq!(count(&store, "blobs"), 0); - assert_eq!(count(&store, "blob_chunks"), 0); - assert_eq!(count(&store, "attachments"), 0); - assert_eq!(count(&store, "task_blobs"), 0); -} - -#[test] -fn concurrent_same_digest_uploads_write_one_complete_blob() { - let (_directory, store, alice, bob) = fixture(); - let barrier = Arc::new(Barrier::new(3)); - let mut handles = Vec::new(); - for (account, task) in [(alice, "alice-task"), (bob, "bob-task")] { - let store = Arc::clone(&store); - let barrier = Arc::clone(&barrier); - handles.push(std::thread::spawn(move || { - let image = upload(&png(2 * MAX_ATTACHMENT_RANGE_BYTES + 9)); - barrier.wait(); - store.store_image_uploads(account, task, &[image]).unwrap() - })); - } - barrier.wait(); - for handle in handles { - handle.join().unwrap(); - } - assert_eq!(count(&store, "blobs"), 1); - assert_eq!(count(&store, "attachments"), 2); - assert_eq!(count(&store, "blob_chunks"), 9); -} - -#[test] -fn prepare_to_link_lease_prevents_deletion_race() { - let (_directory, store, alice, bob) = fixture(); - let bytes = png(MAX_ATTACHMENT_RANGE_BYTES + 9); - store - .store_image_uploads(alice, "alice-task", &[upload(&bytes)]) - .unwrap(); - let prepared = store - .prepare_entry_images(bob, "bob-task", &tool_entry(&bytes)) - .unwrap(); - mark_deleted(&store, "alice-task"); - let (started_tx, started_rx) = mpsc::channel(); - let (done_tx, done_rx) = mpsc::channel(); - let collector = Arc::clone(&store); - let handle = std::thread::spawn(move || { - started_tx.send(()).unwrap(); - done_tx - .send(collector.collect_task_blobs_step(alice, "alice-task")) - .unwrap(); - }); - started_rx.recv().unwrap(); - assert!(matches!( - done_rx.recv_timeout(Duration::from_millis(30)), - Err(mpsc::RecvTimeoutError::Timeout) - )); - store - .write(Durability::Durable, |transaction| { - prepared.link(transaction, active_task(transaction, bob, "bob-task")?) - }) - .unwrap(); - drop(prepared); - done_rx - .recv_timeout(Duration::from_secs(5)) - .unwrap() - .unwrap(); - handle.join().unwrap(); - assert_eq!(count(&store, "blobs"), 1); - let mut image = tool_entry(&bytes); - let prepared = store.prepare_entry_images(bob, "bob-task", &image).unwrap(); - image = prepared.entry.clone(); - drop(prepared); - store - .rehydrate_entry_images(bob, "bob-task", &mut image) - .unwrap(); - assert_eq!(image, tool_entry(&bytes)); -} - -#[test] -fn orphan_collection_resumes_after_reopen_without_touching_previous_files() { - let (directory, store, alice, _) = fixture(); - let sentinel = directory.path().join("previous-attachments"); - std::fs::write(&sentinel, b"leave previous data untouched").unwrap(); - let bytes = png(3 * MAX_ATTACHMENT_RANGE_BYTES + 7); - let prepared = store - .prepare_entry_images(alice, "alice-task", &tool_entry(&bytes)) - .unwrap(); - drop(prepared); // Source transaction never happened. - let first = store.collect_orphan_blobs_step().unwrap(); - assert_eq!(first.chunks, 4); - assert_eq!(first.bytes, MAX_ATTACHMENT_RANGE_BYTES); - drop(store); - let reopened = Store::open(StorePaths::new(directory.path())).unwrap(); - for _ in 0..100 { - let progress = reopened.collect_orphan_blobs_step().unwrap(); - assert!(progress.bytes <= MAX_ATTACHMENT_RANGE_BYTES); - if progress.done { - break; - } - } - assert_eq!(count(&reopened, "blobs"), 0); - assert_eq!(count(&reopened, "blob_chunks"), 0); - assert_eq!( - std::fs::read(sentinel).unwrap(), - b"leave previous data untouched" - ); -} - -#[test] -fn interrupted_preparation_is_rewritten_and_pending_account_deletion_never_collects() { - let (_directory, store, alice, _) = fixture(); - let bytes = png(MAX_ATTACHMENT_RANGE_BYTES + 3); - let digest: [u8; 32] = Sha256::digest(&bytes).into(); - let id = { - let _lease = store.blob_locks.acquire(digest).unwrap(); - store.prepare_blob(&bytes, "image/png", digest).unwrap() - }; - // Simulate chunks torn between transactions: the retry cannot treat them - // as a completed deduplication hit. - store - .write(Durability::Normal, |transaction| { - transaction.execute("DELETE FROM blob_chunks WHERE blob_id=?1 AND idx>0", [id])?; - transaction.execute("UPDATE blob_chunks SET data=x'00' WHERE blob_id=?1", [id])?; - Ok(()) - }) - .unwrap(); - let attachment = store - .store_image_uploads(alice, "alice-task", &[upload(&bytes)]) - .unwrap() - .remove(0); - let read = store - .read_attachment( - alice, - "alice-task", - &attachment.id, - 0, - MAX_ATTACHMENT_RANGE_BYTES, - ) - .unwrap(); - assert_eq!(read.data, bytes[..MAX_ATTACHMENT_RANGE_BYTES]); - mark_deleted(&store, "alice-task"); - store - .write(Durability::Durable, |transaction| { - transaction.execute( - "UPDATE accounts SET deleting_ms=2,deletion_phase='pending_server' WHERE id=?1", - [alice.0], - )?; - Ok(()) - }) - .unwrap(); - assert!(store.collect_task_blobs_step(alice, "alice-task").is_err()); - assert_eq!(count(&store, "blobs"), 1); - assert_eq!(count(&store, "attachments"), 1); -} diff --git a/apps/maple-agent/crates/maple-agent/src/store/gc.rs b/apps/maple-agent/crates/maple-agent/src/store/gc.rs index a563194e6..93f234935 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/gc.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/gc.rs @@ -12,12 +12,20 @@ use super::{AccountId, Store, StoreError, StoreResult}; const ROW_CAP: usize = 1_000; const BYTE_CAP: u64 = 1024 * 1024; +/// Cursors affect scheduling only. Durable tombstones remain the authority, and +/// restarting a host merely starts another bounded pass from the beginning. +#[derive(Default)] +pub(super) struct CollectionCursors { + pub(super) unbound: Option<(i64, String)>, + deleted_task: Option<(i64, i64)>, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CollectionStage { Timeline, Entries, Summaries, - Blobs, + Attachments, TaskMetadata, AccountMetadata, Done, @@ -60,6 +68,43 @@ pub struct DeletedTask { } impl Store { + /// Rotate past failed or busy tasks so a few retained tombstones cannot + /// indefinitely starve later tasks. Keep `deleted_tasks` a stable snapshot. + pub(crate) fn deleted_tasks_for_collection(&self, limit: u16) -> StoreResult> { + if limit == 0 || limit > 200 { + return Err(StoreError::InvalidInput); + } + let after = self + .collection_cursors + .lock() + .map_err(|_| StoreError::Busy)? + .deleted_task; + let candidates = self.read(|tx| { + let query = |after: Option<(i64, i64)>| -> StoreResult> { + let mut statement = tx.prepare_cached( + "SELECT t.deleted_ms,t.id,t.account_id,t.uuid FROM tasks t INDEXED BY tasks_deleting + JOIN accounts a ON a.id=t.account_id WHERE t.deleted_ms IS NOT NULL + AND (a.deletion_phase IS NULL OR a.deletion_phase='confirmed') + AND (?1 IS NULL OR (t.deleted_ms,t.id)>(?1,?2)) + ORDER BY t.deleted_ms,t.id LIMIT ?3", + )?; + Ok(statement.query_map(params![after.map(|key| key.0), after.map(|key| key.1), limit], |row| { + Ok(((row.get(0)?, row.get(1)?), DeletedTask { + account: AccountId(row.get(2)?), + uuid: row.get(3)?, + })) + })?.collect::>>()?) + }; + let candidates = query(after)?; + if candidates.is_empty() && after.is_some() { query(None) } else { Ok(candidates) } + })?; + self.collection_cursors + .lock() + .map_err(|_| StoreError::Busy)? + .deleted_task = candidates.last().map(|(key, _)| *key); + Ok(candidates.into_iter().map(|(_, task)| task).collect()) + } + /// Startup and idle collection use only authorized tombstones. Pending /// account attempts are excluded even if a task had an earlier tombstone. pub fn deleted_tasks(&self, limit: u16) -> StoreResult> { @@ -84,8 +129,8 @@ impl Store { }) } - /// Perform at most one byte-capped payload transaction, or one digest-locked - /// blob step. The caller yields to foreground work between calls. + /// Perform at most one byte-capped payload transaction, or one attachment + /// filesystem cleanup step. The caller yields to foreground work between calls. pub fn collect_task_step( &self, account: AccountId, @@ -113,14 +158,14 @@ impl Store { return Ok(progress.measured(started)); } - // Blob collection takes the per-digest lock before the writer. Never - // call it from inside a write closure, which would reverse lock order. - let blobs = self.collect_task_blobs_step(account, uuid)?; - if !blobs.done || blobs.chunks != 0 || blobs.links != 0 { + // Attachment cleanup removes the exact owned files before their metadata. + // Run outside the writer so filesystem I/O never holds a SQL transaction. + let attachments = self.collect_task_attachments_step(account, uuid)?; + if !attachments.done || attachments.files != 0 || attachments.links != 0 { return Ok(CollectionProgress::new( - CollectionStage::Blobs, - blobs.chunks + blobs.links, - blobs.bytes as u64, + CollectionStage::Attachments, + attachments.files + attachments.links, + attachments.bytes as u64, ) .measured(started)); } @@ -132,8 +177,7 @@ impl Store { EXISTS(SELECT 1 FROM entries WHERE task_id=?1) OR EXISTS(SELECT 1 FROM timeline_rows WHERE task_id=?1) OR EXISTS(SELECT 1 FROM tool_summaries WHERE task_id=?1) OR - EXISTS(SELECT 1 FROM attachments WHERE task_id=?1) OR - EXISTS(SELECT 1 FROM task_blobs WHERE task_id=?1)", + EXISTS(SELECT 1 FROM file_attachments WHERE task_id=?1)", [task], |row| row.get(0), )?; diff --git a/apps/maple-agent/crates/maple-agent/src/store/idle_maintenance.rs b/apps/maple-agent/crates/maple-agent/src/store/idle_maintenance.rs index 43e5a8bbe..647417111 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/idle_maintenance.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/idle_maintenance.rs @@ -58,70 +58,80 @@ mod tests { .unwrap() } - fn seed_deleted_blob(store: &Store) -> Vec { - let live = vec![0xa6; 64 * 1024]; + fn seed_deleted_source(store: &Store) -> Vec { + let payload = |id, size| { + serde_json::to_vec(&pi_coding_agent::session::SessionEntry { + id: format!("entry{id}"), + parent_id: None, + timestamp: 1, + kind: pi_coding_agent::session::EntryKind::CustomMessage { + custom_type: "vacuum-fixture".into(), + content: vec![pi_ai::Content::text("x".repeat(size))], + display: false, + details: None, + }, + }) + .unwrap() + }; + let retired = payload(1, 1024 * 1024); + let live = payload(2, 64 * 1024); store .write(Durability::Durable, |transaction| { - transaction.execute( - "INSERT INTO blobs(id,sha256,mime,size,complete,created_ms) - VALUES(1,?1,'application/octet-stream',1048576,1,1), - (2,?2,'application/octet-stream',65536,1,1)", - params![[1_u8; 32].as_slice(), [2_u8; 32].as_slice()], - )?; - for index in 0..4 { + transaction.execute("INSERT INTO accounts(id,user_key,user_id,created_ms,last_seen_ms) VALUES (1,'a','A',1,1)", [])?; + transaction.execute("INSERT INTO projects(id,account_id,root,last_used_ms) VALUES (1,1,'/vacuum-fixture',1)", [])?; + transaction.execute("INSERT INTO tasks(id,uuid,account_id,project_id,kind,state,title,title_fold,created_ms,updated_ms) VALUES (1,'vacuum-task',1,1,'desktop','settled','vacuum','vacuum',1,1)", [])?; + for (id, bytes) in [(1, &retired), (2, &live)] { + // Leave the fixture uncompressed to create real free pages. transaction.execute( - "INSERT INTO blob_chunks(blob_id,idx,data) VALUES(1,?1,zeroblob(262144))", - [index], + "INSERT INTO entries(id,task_id,seq,entry_id,kind,created_ms,bytes,payload) \ + VALUES (?1,1,?1,?2,'custom_message',1,?3,?4)", + params![id, format!("entry{id}"), bytes.len() as i64, bytes], )?; } - transaction.execute( - "INSERT INTO blob_chunks(blob_id,idx,data) VALUES(2,0,?1)", - [&live], - )?; Ok(()) }) .unwrap(); store .write(Durability::Durable, |transaction| { - transaction.execute("DELETE FROM blobs WHERE id=1", [])?; + transaction.execute("DELETE FROM entries WHERE id=1", [])?; Ok(()) }) .unwrap(); live } - fn assert_live_blob(store: &Store, expected: &[u8]) { + fn assert_live_source(store: &Store, expected: &[u8]) { let actual: Vec = store .read(|transaction| { - Ok(transaction.query_row( - "SELECT data FROM blob_chunks WHERE blob_id=2 AND idx=0", - [], - |row| row.get(0), - )?) + Ok( + transaction.query_row("SELECT payload FROM entries WHERE id=2", [], |row| { + row.get(0) + })?, + ) }) .unwrap(); assert_eq!(actual, expected); } #[test] - fn incremental_chunks_reclaim_multiple_pages_with_a_fixed_bound_and_preserve_live_blobs() { + fn incremental_chunks_reclaim_multiple_pages_with_a_fixed_bound_and_preserve_live_source() { for keyed in [false, true] { let (_directory, store) = fixture(keyed); - let live = seed_deleted_blob(&store); + let live = seed_deleted_source(&store); let mut before = free_pages(&store); assert!(before > i64::from(VACUUM_PAGE_CAP) * 3); store.idle_storage_maintenance(false).unwrap(); let after = free_pages(&store); assert_eq!(before - after, i64::from(VACUUM_PAGE_CAP)); - assert_live_blob(&store, &live); + assert_live_source(&store, &live); before = after; for _ in 0..16 { store.idle_storage_maintenance(false).unwrap(); let after = free_pages(&store); assert!((0..=i64::from(VACUUM_PAGE_CAP)).contains(&(before - after))); - assert_live_blob(&store, &live); + assert_live_source(&store, &live); if after == 0 { break; } @@ -131,7 +141,7 @@ mod tests { assert_eq!(free_pages(&store), 0); // Repeated idle scheduling remains harmless once reclamation ends. store.idle_storage_maintenance(true).unwrap(); - assert_live_blob(&store, &live); + assert_live_source(&store, &live); store.close().unwrap(); } } diff --git a/apps/maple-agent/crates/maple-agent/src/store/integrity.rs b/apps/maple-agent/crates/maple-agent/src/store/integrity.rs index 41541e0fb..9ffc0ad17 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/integrity.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/integrity.rs @@ -108,10 +108,9 @@ pub enum IntegrityTable { Entries, TimelineRows, ToolSummaries, + /// Read-only diagnostics for a pre-file-store schema. Blobs, - BlobChunks, - TaskBlobs, - Attachments, + FileAttachments, SendReceipts, RemoteIdentities, } @@ -934,10 +933,7 @@ fn classify(message: &str, cipher: bool) -> IntegrityFinding { ("entries", IntegrityTable::Entries), ("timeline_rows", IntegrityTable::TimelineRows), ("tool_summaries", IntegrityTable::ToolSummaries), - ("blob_chunks", IntegrityTable::BlobChunks), - ("task_blobs", IntegrityTable::TaskBlobs), - ("blobs", IntegrityTable::Blobs), - ("attachments", IntegrityTable::Attachments), + ("file_attachments", IntegrityTable::FileAttachments), ("send_receipts", IntegrityTable::SendReceipts), ("remote_identities", IntegrityTable::RemoteIdentities), ] diff --git a/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints.rs b/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints.rs index 2a5a9fbc1..0d587a52e 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints.rs @@ -13,6 +13,7 @@ use crate::store::schema; struct Constraints { table: &'static str, finding: IntegrityTable, + removed: Option, // The first schema version containing each predicate. checks: &'static [(i64, &'static str)], } @@ -21,26 +22,31 @@ const CONSTRAINTS: &[Constraints] = &[ Constraints { table: "accounts", finding: IntegrityTable::Accounts, + removed: None, checks: &[(1, "deletion_phase IN ('pending_server', 'confirmed')")], }, Constraints { table: "credentials", finding: IntegrityTable::Credentials, + removed: None, checks: &[(1, "slot = 1")], }, Constraints { table: "app_settings", finding: IntegrityTable::AppSettings, + removed: None, checks: &[(1, "json_valid(value)")], }, Constraints { table: "account_settings", finding: IntegrityTable::AccountSettings, + removed: None, checks: &[(1, "json_valid(value)")], }, Constraints { table: "projects", finding: IntegrityTable::Projects, + removed: None, checks: &[ (1, "trusted IN (0, 1)"), (1, "json_valid(mcp_defaults)"), @@ -50,16 +56,19 @@ const CONSTRAINTS: &[Constraints] = &[ Constraints { table: "mcp_servers", finding: IntegrityTable::McpServers, + removed: None, checks: &[(1, "json_valid(definition)")], }, Constraints { table: "mcp_server_secrets", finding: IntegrityTable::McpServerSecrets, + removed: None, checks: &[(1, "kind IN ('env', 'header')"), (1, "length(target) = 32")], }, Constraints { table: "tasks", finding: IntegrityTable::Tasks, + removed: None, checks: &[ (1, "kind IN ('desktop', 'acp')"), (1, "state IN ('active', 'settled', 'archived')"), @@ -72,11 +81,25 @@ const CONSTRAINTS: &[Constraints] = &[ Constraints { table: "blobs", finding: IntegrityTable::Blobs, + removed: Some(4), checks: &[(1, "length(sha256) = 32")], }, + Constraints { + table: "file_attachments", + finding: IntegrityTable::FileAttachments, + removed: None, + checks: &[ + (3, "length(public_id) = 32"), + (3, "size >= 0 AND size <= 31457280"), + (3, "length(sha256) = 64"), + (3, "state IN ('uploading', 'ready', 'cancelled')"), + (3, "bound IN (0, 1)"), + ], + }, Constraints { table: "send_receipts", finding: IntegrityTable::SendReceipts, + removed: None, checks: &[ (1, "result IS NULL OR json_valid(result)"), (1, "interrupted IN (0, 1)"), @@ -85,6 +108,7 @@ const CONSTRAINTS: &[Constraints] = &[ Constraints { table: "remote_identities", finding: IntegrityTable::RemoteIdentities, + removed: None, checks: &[(1, "length(secret) = 32"), (1, "json_valid(record)")], }, ]; @@ -103,6 +127,9 @@ pub(super) fn collect( return Err(rusqlite::Error::InvalidQuery); } for table in CONSTRAINTS { + if table.removed.is_some_and(|removed| version >= removed) { + continue; + } if findings.len() >= MAX_FINDINGS { break; } @@ -122,6 +149,9 @@ pub(super) fn collect( .map(|(_, expression)| format!("NOT ({expression})")) .collect::>() .join(" OR "); + if predicate.is_empty() { + continue; + } let sql = format!( "SELECT 1 FROM main.\"{}\" NOT INDEXED WHERE {predicate} LIMIT 1", table.table, diff --git a/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints_tests.rs b/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints_tests.rs index cce884f99..45e1a711f 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints_tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/integrity_constraints_tests.rs @@ -5,7 +5,7 @@ use std::sync::{ }; use std::time::Duration; -use rusqlite::{Connection, params}; +use rusqlite::Connection; use super::super::{ CheckIncomplete, IntegrityCancellation, IntegrityKind, IntegrityOutcome, check_connection, @@ -106,14 +106,7 @@ fn expressions(sql: &str) -> Vec { fn catalog_matches_every_check_in_each_accepted_owned_schema() { for version in 1..=schema::VERSION { let mut reference = Connection::open_in_memory().unwrap(); - if version == 1 { - reference.execute_batch(schema::REFERENCE_DDL).unwrap(); - } else { - reference - .pragma_update(None, "application_id", schema::APPLICATION_ID) - .unwrap(); - schema::migrate(&mut reference, 1).unwrap(); - } + schema::migrate_to_version(&mut reference, 1, version).unwrap(); let mut statement = reference .prepare("SELECT name,sql FROM sqlite_schema WHERE type='table'") .unwrap(); @@ -131,6 +124,8 @@ fn catalog_matches_every_check_in_each_accepted_owned_schema() { .collect::>(); let actual = CONSTRAINTS .iter() + .filter(|table| table.removed.is_none_or(|removed| version < removed)) + .filter(|table| table.checks.iter().any(|(since, _)| *since <= version)) .map(|table| { let mut checks = table .checks @@ -194,20 +189,18 @@ fn readonly_plain_and_keyed_checks_detect_json_and_integer_constraints() { #[test] fn readonly_checks_accept_nullable_checks_and_version_one_without_new_column() { for version in [1, 2] { - let (_directory, store) = fixture(false); - let account = store.upsert_account("nullable-fixture", 1).unwrap(); - { - let guard = store.writer.lock().unwrap(); - let writer = guard.as_ref().unwrap(); - if version == 1 { - writer - .execute_batch( - "ALTER TABLE projects DROP COLUMN name_revision; PRAGMA user_version=1", - ) - .unwrap(); - } - writer.execute("INSERT INTO projects(account_id,root,last_used_ms,trusted,mcp_defaults) VALUES (?1,'nullable-fixture',1,NULL,'{}')", params![account.0]).unwrap(); - } + let directory = tempfile::tempdir().unwrap(); + let paths = StorePaths::new(directory.path()).prepare().unwrap(); + paths.precreate_database().unwrap(); + let mut writer = Connection::open(paths.database()).unwrap(); + writer + .pragma_update(None, "application_id", schema::APPLICATION_ID) + .unwrap(); + schema::migrate_to_version(&mut writer, 1, version).unwrap(); + writer.execute("INSERT INTO accounts(id,user_key,user_id,created_ms,last_seen_ms) VALUES (1,'nullable-fixture','nullable-fixture',1,1)",[]).unwrap(); + writer.execute("INSERT INTO projects(account_id,root,last_used_ms,trusted,mcp_defaults) VALUES (1,'nullable-fixture',1,NULL,'{}')", []).unwrap(); + drop(writer); + let store = Store::unavailable(paths, crate::store::StoreState::Maintenance).unwrap(); let connection = open_readonly(&store.paths, None).unwrap(); let report = check_connection( &connection, diff --git a/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage.rs b/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage.rs index e510be8fd..227fc7f24 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage.rs @@ -44,10 +44,7 @@ pub enum SalvageTable { Entries, TimelineRows, ToolSummaries, - Blobs, - BlobChunks, - TaskBlobs, - Attachments, + FileAttachments, SendReceipts, RemoteIdentities, } @@ -68,10 +65,7 @@ impl SalvageTable { Self::Entries => "entries", Self::TimelineRows => "timeline_rows", Self::ToolSummaries => "tool_summaries", - Self::Blobs => "blobs", - Self::BlobChunks => "blob_chunks", - Self::TaskBlobs => "task_blobs", - Self::Attachments => "attachments", + Self::FileAttachments => "file_attachments", Self::SendReceipts => "send_receipts", Self::RemoteIdentities => "remote_identities", } @@ -90,8 +84,7 @@ impl SalvageTable { "mcp_server_secrets INDEXED BY sqlite_autoindex_mcp_server_secrets_1" } Self::ToolSummaries => "tool_summaries INDEXED BY sqlite_autoindex_tool_summaries_1", - Self::TaskBlobs => "task_blobs INDEXED BY sqlite_autoindex_task_blobs_1", - Self::Attachments => "attachments INDEXED BY sqlite_autoindex_attachments_1", + Self::FileAttachments => "file_attachments NOT INDEXED", Self::Accounts => "accounts NOT INDEXED", Self::Credentials => "credentials NOT INDEXED", Self::AppSettings => "app_settings NOT INDEXED", @@ -101,17 +94,14 @@ impl SalvageTable { Self::Tasks => "tasks NOT INDEXED", Self::Entries => "entries NOT INDEXED", Self::TimelineRows => "timeline_rows NOT INDEXED", - Self::Blobs => "blobs NOT INDEXED", - Self::BlobChunks => "blob_chunks NOT INDEXED", + Self::SendReceipts => "send_receipts NOT INDEXED", Self::RemoteIdentities => "remote_identities NOT INDEXED", } } fn key_columns(self) -> &'static [usize] { match self { - Self::AccountSettings | Self::BlobChunks | Self::TaskBlobs | Self::Attachments => { - &[0, 1] - } + Self::AccountSettings => &[0, 1], Self::McpServerSecrets | Self::ToolSummaries => &[0, 1, 2], _ => &[0], } @@ -119,7 +109,7 @@ impl SalvageTable { } // Derived timeline rows never determine whether authoritative source is readable. -const TABLES: [SalvageTable; 18] = [ +const TABLES: [SalvageTable; 15] = [ SalvageTable::StoreMeta, SalvageTable::Accounts, SalvageTable::Credentials, @@ -132,10 +122,7 @@ const TABLES: [SalvageTable; 18] = [ SalvageTable::Tasks, SalvageTable::Entries, SalvageTable::ToolSummaries, - SalvageTable::Blobs, - SalvageTable::BlobChunks, - SalvageTable::TaskBlobs, - SalvageTable::Attachments, + SalvageTable::FileAttachments, SalvageTable::SendReceipts, SalvageTable::RemoteIdentities, ]; @@ -161,7 +148,6 @@ pub enum SalvageRowKey { #[serde(rename_all = "snake_case")] pub enum SalvageReason { UnreadableTask, - InvalidBlob, OrphanedRow, } @@ -225,7 +211,6 @@ impl From for SalvageError { struct Scan { preview: SalvagePreview, bad_tasks: BTreeSet, - bad_blobs: BTreeSet, } impl Store { @@ -376,27 +361,6 @@ fn scan(source: &Connection) -> StoreResult { })?; } validate_metadata(source)?; - let has_deleted_tasks: bool = source.query_row( - "SELECT EXISTS(SELECT 1 FROM tasks NOT INDEXED WHERE deleted_ms IS NOT NULL)", - [], - |row| row.get(0), - )?; - let mut bad_blobs = BTreeSet::new(); - walk(source, SalvageTable::Blobs, |row, _| { - let id = row.get(0)?; - // A task collector can leave an incomplete blob while deleting its - // chunks. Preserve that hidden, resumable state rather than reporting - // already requested deletion as new corruption loss. Any live or - // orphan ownership still requires full payload validation. - // Do the ownership scan only for invalid payloads, avoiding extra - // whole-link scans for the normal complete-blob case. - if !valid_blob(source, row)? - && !(has_deleted_tasks && blob_has_only_deleted_owners(source, id)?) - { - insert_bad(&mut bad_blobs, id)?; - } - Ok(()) - })?; let mut bad_tasks = BTreeSet::new(); walk(source, SalvageTable::Tasks, |row, _| { let id = row.get(0)?; @@ -405,37 +369,18 @@ fn scan(source: &Connection) -> StoreResult { } Ok(()) })?; - for table in [SalvageTable::TaskBlobs, SalvageTable::Attachments] { - walk(source, table, |row, _| { - let task = row.get(0)?; - let blob = row.get(if table == SalvageTable::TaskBlobs { - 1 - } else { - 2 - })?; - let deleted = task_is_deleted(source, task)?; - let blob_exists = exists(source, "blobs", blob)?; - if deleted && !blob_exists { - // Normal collection removes a final link and its blob in one - // transaction. A broken FK is additional source corruption, - // not proof that this hidden ownership was fully collected. - return Err(StoreError::InvalidData); - } - if !deleted && (bad_blobs.contains(&blob) || !blob_exists) { - insert_bad(&mut bad_tasks, task)?; - } - Ok(()) - })?; - } // An orphan source row is still enumerated and explicitly reported. for table in [ SalvageTable::Entries, SalvageTable::ToolSummaries, - SalvageTable::TaskBlobs, - SalvageTable::Attachments, + SalvageTable::FileAttachments, ] { walk(source, table, |row, _| { - let task = row.get(if table == SalvageTable::Entries { 1 } else { 0 })?; + let task = row.get(match table { + SalvageTable::Entries => 1, + SalvageTable::FileAttachments => 2, + _ => 0, + })?; if !exists(source, "tasks", task)? { insert_bad(&mut bad_tasks, task)?; } @@ -456,10 +401,7 @@ fn scan(source: &Connection) -> StoreResult { if table == SalvageTable::StoreMeta { return Ok(()); } - if discard_deleted_blob_link(source, table, row, &bad_blobs)? { - return Ok(()); - } - if let Some(reason) = drop_reason(source, table, row, &bad_tasks, &bad_blobs)? { + if let Some(reason) = drop_reason(table, row, &bad_tasks)? { if report.dropped_rows.len() >= MAX_LOSS_ROWS { return Err(StoreError::InvalidInput); } @@ -494,7 +436,6 @@ fn scan(source: &Connection) -> StoreResult { fingerprint: fingerprint.finalize().into(), }, bad_tasks, - bad_blobs, }) } @@ -514,6 +455,10 @@ fn validate_metadata(source: &Connection) -> StoreResult<()> { return Err(StoreError::InvalidData); } } + if source.query_row( + "SELECT EXISTS(SELECT 1 FROM file_attachments a LEFT JOIN tasks t ON t.id=a.task_id LEFT JOIN accounts owner ON owner.id=a.account_id WHERE owner.id IS NULL OR (t.id IS NOT NULL AND t.account_id<>a.account_id))", + [], |row| row.get::<_, bool>(0), + )? { return Err(StoreError::InvalidData); } // The all-table pragma would also read the disposable derived timeline. // Scope each check to owned authoritative tables, keeping every other // metadata FK refusal unchanged. @@ -527,9 +472,7 @@ fn validate_metadata(source: &Connection) -> StoreResult<()> { SalvageTable::Tasks | SalvageTable::Entries | SalvageTable::ToolSummaries - | SalvageTable::BlobChunks - | SalvageTable::TaskBlobs - | SalvageTable::Attachments + | SalvageTable::FileAttachments ) { return Err(StoreError::InvalidData); } @@ -697,14 +640,20 @@ fn valid_images( }; let mut moved = false; for block in content.into_iter().flatten() { - if let Content::Image(image) = block - && let Some(marker) = image.data.strip_prefix("maple-blob:") - { - moved = true; - let Some(blob) = marker.parse::().ok().filter(|id| *id > 0) else { + if let Content::Image(image) = block { + if image.data == "maple-attachment-unavailable:retired" { + continue; + } + let Some(id) = image.data.strip_prefix("maple-attachment://") else { return Ok(false); }; - if !task_has_blob(source, task, blob)? { + moved = true; + if id.len() != 32 + || !id.bytes().all(|byte| byte.is_ascii_hexdigit()) + || !source.query_row(TASK_HAS_ATTACHMENT_SQL, params![task, id], |row| { + row.get::<_, bool>(0) + })? + { return Ok(false); } } @@ -712,97 +661,18 @@ fn valid_images( Ok(moved == (flags & 1 != 0)) } -const TASK_HAS_BLOB_SQL: &str = "SELECT EXISTS(SELECT 1 FROM task_blobs INDEXED BY sqlite_autoindex_task_blobs_1 WHERE task_id=?1 AND blob_id=?2)"; - -fn task_has_blob(source: &Connection, task: i64, blob: i64) -> StoreResult { - Ok(source.query_row(TASK_HAS_BLOB_SQL, params![task, blob], |row| row.get(0))?) -} - -fn valid_blob(source: &Connection, row: &Row<'_>) -> StoreResult { - let id: i64 = row.get(0)?; - let expected: Vec = row.get(1)?; - let size: i64 = row.get(3)?; - let complete: i64 = row.get(4)?; - if size < 0 || expected.len() != 32 || complete != 1 { - return Ok(false); - } - let mut hash = Sha256::new(); - let mut length = 0_i64; - let mut next = 0_i64; - let mut statement = source - .prepare("SELECT idx,data FROM blob_chunks NOT INDEXED WHERE blob_id=?1 ORDER BY idx")?; - let mut chunks = statement.query([id])?; - while let Some(row) = chunks.next()? { - let index: i64 = row.get(0)?; - let data: Vec = row.get(1)?; - let remaining = size.saturating_sub(length); - let expected_length = remaining.min(crate::store::blobs::BLOB_CHUNK_BYTES as i64); - if index != next || data.len() as i64 != expected_length || expected_length <= 0 { - return Ok(false); - } - hash.update(&data); - length += data.len() as i64; - next += 1; - } - let actual: [u8; 32] = hash.finalize().into(); - Ok(length == size && actual.as_slice() == expected) -} - -fn task_is_deleted(source: &Connection, task: i64) -> StoreResult { - Ok(source.query_row( - "SELECT EXISTS(SELECT 1 FROM tasks NOT INDEXED WHERE id=?1 AND deleted_ms IS NOT NULL)", - [task], - |row| row.get(0), - )?) -} - -fn blob_has_only_deleted_owners(source: &Connection, blob: i64) -> StoreResult { - // Both link scans name the WITHOUT ROWID primary tree. Missing task rows - // are not proof of deletion, and shared live ownership must stay validated. - Ok(source.query_row( - "SELECT (EXISTS(SELECT 1 FROM task_blobs INDEXED BY sqlite_autoindex_task_blobs_1 WHERE blob_id=?1) - OR EXISTS(SELECT 1 FROM attachments INDEXED BY sqlite_autoindex_attachments_1 WHERE blob_id=?1)) - AND NOT EXISTS(SELECT 1 FROM task_blobs AS link INDEXED BY sqlite_autoindex_task_blobs_1 - LEFT JOIN tasks AS task NOT INDEXED ON task.id=link.task_id - WHERE link.blob_id=?1 AND (task.id IS NULL OR task.deleted_ms IS NULL)) - AND NOT EXISTS(SELECT 1 FROM attachments AS link INDEXED BY sqlite_autoindex_attachments_1 - LEFT JOIN tasks AS task NOT INDEXED ON task.id=link.task_id - WHERE link.blob_id=?1 AND (task.id IS NULL OR task.deleted_ms IS NULL))", - [blob], - |row| row.get(0), - )?) -} - -fn discard_deleted_blob_link( - source: &Connection, - table: SalvageTable, - row: &Row<'_>, - bad_blobs: &BTreeSet, -) -> StoreResult { - let blob_column = match table { - SalvageTable::TaskBlobs => 1, - SalvageTable::Attachments => 2, - _ => return Ok(false), - }; - let blob = row.get(blob_column)?; - // A shared corrupt blob is omitted for its live owners. Its hidden owner - // keeps its tombstone/history, but cannot retain an FK to the omitted blob. - // That link removal is already authorized deletion, not a new loss row. - Ok(bad_blobs.contains(&blob) && task_is_deleted(source, row.get(0)?)?) -} +const TASK_HAS_ATTACHMENT_SQL: &str = "SELECT EXISTS(SELECT 1 FROM file_attachments NOT INDEXED WHERE task_id=?1 AND public_id=?2 AND state='ready')"; fn drop_reason( - source: &Connection, table: SalvageTable, row: &Row<'_>, tasks: &BTreeSet, - blobs: &BTreeSet, ) -> StoreResult> { - use SalvageTable::*; let task_column = match table { - Tasks => Some(0), - Entries => Some(1), - ToolSummaries | TaskBlobs | Attachments => Some(0), + SalvageTable::Tasks => Some(0), + SalvageTable::Entries => Some(1), + SalvageTable::ToolSummaries => Some(0), + SalvageTable::FileAttachments => Some(2), _ => None, }; if let Some(column) = task_column { @@ -811,20 +681,11 @@ fn drop_reason( return Ok(Some(SalvageReason::UnreadableTask)); } } - if matches!(table, Blobs | BlobChunks) { - let id = row.get(0)?; - if blobs.contains(&id) { - return Ok(Some(SalvageReason::InvalidBlob)); - } - if table == BlobChunks && !exists(source, "blobs", id)? { - return Ok(Some(SalvageReason::OrphanedRow)); - } - } Ok(None) } fn copy_rows(source: &Connection, destination: &mut Connection, scan: &Scan) -> StoreResult<()> { - // FK order differs from DDL order because task blob links follow blobs. + // Copy authoritative metadata after its account/task owners. let order = [ SalvageTable::Accounts, SalvageTable::Credentials, @@ -837,10 +698,7 @@ fn copy_rows(source: &Connection, destination: &mut Connection, scan: &Scan) -> SalvageTable::Tasks, SalvageTable::Entries, SalvageTable::ToolSummaries, - SalvageTable::Blobs, - SalvageTable::BlobChunks, - SalvageTable::TaskBlobs, - SalvageTable::Attachments, + SalvageTable::FileAttachments, SalvageTable::SendReceipts, SalvageTable::RemoteIdentities, ]; @@ -848,9 +706,7 @@ fn copy_rows(source: &Connection, destination: &mut Connection, scan: &Scan) -> for table in order { let mut insert = None; walk(source, table, |row, columns| { - if discard_deleted_blob_link(source, table, row, &scan.bad_blobs)? - || drop_reason(source, table, row, &scan.bad_tasks, &scan.bad_blobs)?.is_some() - { + if drop_reason(table, row, &scan.bad_tasks)?.is_some() { return Ok(()); } let mut values = (0..columns) diff --git a/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage_tests.rs b/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage_tests.rs index 49bcae0d7..4bc504bfd 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage_tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/integrity_salvage_tests.rs @@ -94,7 +94,7 @@ fn upload(store: &Store, account: AccountId, task: &TaskRecord) -> String { } #[test] -fn invalid_source_drops_whole_task_and_keeps_readable_history_blobs_and_settings() { +fn invalid_source_drops_whole_task_and_keeps_readable_history_files_and_settings() { let (_directory, store, account) = fixture(); let good = task(&store, account, 1); let bad = task(&store, account, 2); @@ -206,56 +206,36 @@ fn invalid_source_drops_whole_task_and_keeps_readable_history_blobs_and_settings } #[test] -fn corrupt_shared_blob_drops_each_linked_task_with_exact_chunk_locator() { +fn missing_attachment_file_does_not_discard_readable_session_history() { let (_directory, store, account) = fixture(); let first = task(&store, account, 1); - let second = task(&store, account, 2); - upload(&store, account, &first); - upload(&store, account, &second); - store - .write(Durability::Durable, |transaction| { - transaction.execute("UPDATE blob_chunks SET data=x'00'", [])?; - Ok(()) - }) - .unwrap(); + let attachment = upload(&store, account, &first); + let path = store + .paths() + .root() + .join("attachments") + .join(hex(&store.store_id().unwrap())) + .join(account.0.to_string()) + .join(&first.uuid) + .join(&attachment) + .join("image.png"); + fs::remove_file(path).unwrap(); damage(&store); let preview = store .preview_salvage(Duration::from_secs(5), &IntegrityCancellation::default()) .unwrap(); - assert_eq!(preview.report().dropped_tasks, 2); - assert!( - preview - .report() - .dropped_rows - .iter() - .any(|row| row.table == SalvageTable::BlobChunks - && matches!(row.key, SalvageRowKey::IntegerPair { second: 0, .. })) - ); + assert_eq!(preview.report().retained_tasks, 1); + assert_eq!(preview.report().dropped_tasks, 0); + assert!(preview.report().dropped_rows.is_empty()); store .save_readable_data(&preview, SALVAGE_CONFIRMATION) .unwrap(); - store - .read(|transaction| { - for table in [ - "tasks", - "entries", - "blobs", - "blob_chunks", - "task_blobs", - "attachments", - ] { - assert_eq!( - transaction.query_row::( - &format!("SELECT count(*) FROM {table}"), - [], - |row| row.get(0) - )?, - 0 - ); - } - Ok(()) - }) - .unwrap(); + assert!(store.task(account, &first.uuid).unwrap().is_some()); + assert!( + store + .read_attachment(account, &first.uuid, &attachment, 0, 1024) + .is_err() + ); } #[test] @@ -500,15 +480,16 @@ fn authoritative_walk_and_image_lookup_open_only_the_primary_tree_plain_and_keye assert_eq!(roots, [root], "{} authoritative root", table.name()); } let root: i64 = source.query_row( - "SELECT rootpage FROM sqlite_schema WHERE name='task_blobs'", + "SELECT rootpage FROM sqlite_schema WHERE name='file_attachments'", [], |row| row.get(0), )?; - let mut explain = source.prepare(&format!("EXPLAIN {TASK_HAS_BLOB_SQL}"))?; + let mut explain = source.prepare(&format!("EXPLAIN {TASK_HAS_ATTACHMENT_SQL}"))?; let roots = explain - .query_map(params![task.key.0, 1_i64], |row| { - Ok((row.get::<_, String>(1)?, row.get::<_, i64>(3)?)) - })? + .query_map( + params![task.key.0, "00000000000000000000000000000000"], + |row| Ok((row.get::<_, String>(1)?, row.get::<_, i64>(3)?)), + )? .collect::>>()? .into_iter() .filter_map(|(opcode, root)| (opcode == "OpenRead").then_some(root)) @@ -521,14 +502,14 @@ fn authoritative_walk_and_image_lookup_open_only_the_primary_tree_plain_and_keye } #[test] -fn corrupt_covering_blob_index_does_not_lose_primary_links_or_image_markers() { +fn corrupt_attachment_index_does_not_lose_primary_metadata_or_image_markers() { let (_directory, store, account) = fixture(); let task = task(&store, account, 1); upload(&store, account, &task); store .write(Durability::Durable, |transaction| { - let blob: i64 = transaction.query_row( - "SELECT blob_id FROM task_blobs WHERE task_id=?1", + let attachment: String = transaction.query_row( + "SELECT public_id FROM file_attachments WHERE task_id=?1", [task.key.0], |row| row.get(0), )?; @@ -538,7 +519,10 @@ fn corrupt_covering_blob_index_does_not_lose_primary_links_or_image_markers() { timestamp: 2, kind: EntryKind::Message { message: SessionMessage::Llm(Message::User(pi_ai::UserMessage { - content: vec![Content::image(format!("maple-blob:{blob}"), "image/png")], + content: vec![Content::image( + format!("maple-attachment://{attachment}"), + "image/png", + )], timestamp: 2, })), }, @@ -551,12 +535,12 @@ fn corrupt_covering_blob_index_does_not_lose_primary_links_or_image_markers() { Ok(()) }) .unwrap(); - let bytes = corrupt_fixture_root(&store, "task_blobs_blob"); + let bytes = corrupt_fixture_root(&store, "file_attachments_task"); let source = super::super::open_readonly(store.paths(), None).unwrap(); assert!( source .query_row::( - "SELECT task_id FROM task_blobs INDEXED BY task_blobs_blob LIMIT 1", + "SELECT task_id FROM file_attachments INDEXED BY file_attachments_task LIMIT 1", [], |row| row.get(0), ) @@ -566,13 +550,7 @@ fn corrupt_covering_blob_index_does_not_lose_primary_links_or_image_markers() { assert_eq!(scanned.preview.report.retained_tasks, 1); assert_eq!(scanned.preview.report.tasks_to_rebuild, 1); assert!(scanned.preview.report.dropped_rows.is_empty()); - for table in [ - "task_blobs", - "attachments", - "blobs", - "blob_chunks", - "entries", - ] { + for table in ["file_attachments", "entries"] { assert_eq!( destination .query_row::(&format!("SELECT count(*) FROM {table}"), [], |row| row @@ -638,11 +616,9 @@ fn tombstoned_partial_collection_is_preserved_without_exhausting_the_loss_review SELECT ?1,seq,printf('%08x',seq),'session_info',2,0,2,x'7b7d' FROM sequence", [deleted.key.0], )?; - // A collector can also leave an incomplete, partially removed blob - // still linked to this tombstone. Account deletion is only pending: + // An unfinished upload remains owned by this tombstone. Account deletion is only pending: // salvage must preserve its phase, credentials and remaining source. - transaction.execute("UPDATE blobs SET complete=0", [])?; - transaction.execute("DELETE FROM blob_chunks", [])?; + transaction.execute("UPDATE file_attachments SET state='uploading'", [])?; transaction.execute("UPDATE accounts SET deleting_ms=3,deletion_phase='pending_server',deletion_generation='pending-fixture' WHERE id=?1", [account.0])?; transaction.execute("INSERT INTO credentials(slot,account_id,api_url,session_id,revision,access_token,refresh_token,updated_ms) VALUES (1,?1,'https://example.invalid','fixture',1,?2,?2,1)",params![account.0,PRIVATE_MARKER.as_bytes()])?; transaction.execute("INSERT INTO store_meta(key,value) VALUES ('integrity_incomplete_attempts',1),('integrity_retry_after_ms',5000)", [])?; @@ -719,12 +695,12 @@ fn tombstoned_partial_collection_is_preserved_without_exhausting_the_loss_review PRIVATE_MARKER.as_bytes() ); assert_eq!( - destination.query_row::( - "SELECT complete FROM blobs", + destination.query_row::( + "SELECT state FROM file_attachments", [], |row| row.get(0) )?, - 0 + "uploading" ); Ok(()) }) @@ -733,66 +709,19 @@ fn tombstoned_partial_collection_is_preserved_without_exhausting_the_loss_review } #[test] -fn shared_bad_blob_still_validates_live_ownership_without_reviving_confirmed_deletion() { - let (_directory, store, live_account) = fixture(); - let deleted_account = store - .upsert_account("confirmed-deleted-fixture", 1) - .unwrap(); - let live = task(&store, live_account, 1); - let deleted = task(&store, deleted_account, 2); - upload(&store, live_account, &live); - upload(&store, deleted_account, &deleted); - store - .begin_account_deletion(deleted_account, "deletion-fixture", 3) - .unwrap(); - store - .confirm_account_deletion(deleted_account, "deletion-fixture") - .unwrap(); - store - .write(Durability::Durable, |transaction| { - transaction.execute("UPDATE blob_chunks SET data=x'00'", [])?; - Ok(()) - }) - .unwrap(); +fn attachment_account_mismatch_cannot_be_salvaged_into_a_different_owner() { + let (_directory, store, account) = fixture(); + let other = store.upsert_account("other-file-owner", 1).unwrap(); + let session = task(&store, account, 1); + upload(&store, account, &session); store - .with_authenticated_store_key(|key| { - let source = super::super::open_readonly(store.paths(), key)?; - let (scanned, destination) = candidate(&source); - assert_eq!(scanned.preview.report.retained_tasks, 0); - assert_eq!(scanned.preview.report.dropped_tasks, 1); - assert!(scanned.preview.report.dropped_rows.iter().all(|loss| { - loss.table != SalvageTable::Tasks - || loss.key != SalvageRowKey::Integer { id: deleted.key.0 } - })); - assert_eq!( - destination.query_row::( - "SELECT deleted_ms FROM tasks WHERE id=?1", - [deleted.key.0], - |row| row.get(0), - )?, - 3 - ); - assert_eq!( - destination.query_row::( - "SELECT deletion_phase FROM accounts WHERE id=?1", - [deleted_account.0], - |row| row.get(0), - )?, - "confirmed" - ); - for table in ["blobs", "blob_chunks", "task_blobs", "attachments"] { - assert_eq!( - destination.query_row::( - &format!("SELECT count(*) FROM {table}"), - [], - |row| row.get(0), - )?, - 0 - ); - } + .write(Durability::Durable, |tx| { + tx.execute("UPDATE file_attachments SET account_id=?1", [other.0])?; Ok(()) }) .unwrap(); + let source = super::super::open_readonly(store.paths(), None).unwrap(); + assert!(matches!(scan(&source), Err(StoreError::InvalidData))); } #[test] @@ -814,7 +743,7 @@ fn tombstone_does_not_turn_broken_authoritative_ownership_into_completed_deletio } else { source .execute( - "UPDATE task_blobs SET blob_id=999 WHERE task_id=?1", + "UPDATE file_attachments SET account_id=999 WHERE task_id=?1", [deleted.key.0], ) .unwrap(); @@ -825,3 +754,40 @@ fn tombstone_does_not_turn_broken_authoritative_ownership_into_completed_deletio assert!(!store.paths().root().join("db.tmp").exists()); } } + +#[test] +fn keyed_salvage_copies_retained_files_to_the_replacement_generation() { + use crate::store::keys::{KdfParams, Password}; + let (_directory, store, account) = fixture(); + let session = task(&store, account, 1); + let attachment = upload(&store, account, &session); + let old = hex(&store.store_id().unwrap()); + let password = Password::normalize("Z7jv!9qP#3nMx@8Ks^2dH").unwrap(); + store + .set_store_password(&password, &[], KdfParams::floor()) + .unwrap(); + assert_eq!( + store + .read_attachment(account, &session.uuid, &attachment, 0, 1024) + .unwrap() + .data, + b"\x89PNG\r\n\x1a\nSALVAGE_IMAGE_BYTES" + ); + damage(&store); + store.authenticate_repair_password(&password).unwrap(); + let preview = store + .preview_salvage(Duration::from_secs(5), &IntegrityCancellation::default()) + .unwrap(); + store + .save_readable_data(&preview, SALVAGE_CONFIRMATION) + .unwrap(); + assert_ne!(hex(&store.store_id().unwrap()), old); + assert_eq!( + store + .read_attachment(account, &session.uuid, &attachment, 0, 1024) + .unwrap() + .data, + b"\x89PNG\r\n\x1a\nSALVAGE_IMAGE_BYTES" + ); + assert!(!store.paths().root().join("attachments").join(old).exists()); +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle.rs index 7bb24aa90..a08787923 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/lifecycle.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle.rs @@ -1000,7 +1000,7 @@ impl Store { source.pragma_update(Some("dst"), "application_id", schema::APPLICATION_ID)?; let value = notice::new_value(notice)?; source.execute("INSERT INTO dst.app_settings(key,value,updated_ms) VALUES ('store.mode_notice',?1,?2) ON CONFLICT(key) DO UPDATE SET value=excluded.value,updated_ms=excluded.updated_ms", params![value,super::now_ms()])?; - for table in ["tasks", "entries", "blobs"] { + for table in ["tasks", "entries", "file_attachments"] { let count: i64 = source.query_row( &format!("SELECT count(*) FROM main.{table}"), [], @@ -1168,7 +1168,7 @@ impl Store { self.verify_or_create_generation()?; } let mut journal = if pending.is_some() { - set_state(self, StoreState::Maintenance)?; + freeze_attachment_writers(self, true)?; self.close_connections_without_checkpoint()?; match recover_filesystem(&self.paths)? { Some(journal) => journal, @@ -1203,7 +1203,7 @@ impl Store { // Publishing the reset intent can succeed before parent syncing // fails. Freeze admission first so no newer acknowledged writes can // be lost when startup finishes that confirmed reset later. - set_state(self, StoreState::Maintenance)?; + freeze_attachment_writers(self, true)?; // An in-flight checkpoint can publish its own quarantine journal. // Join it before publishing reset intent so it cannot overwrite the // operation identity used to retire the confirmed source below. @@ -1224,6 +1224,9 @@ impl Store { write_json(&self.paths.setup(), &journal)?; } set_state(self, StoreState::Maintenance)?; + if !self.cancel_attachment_writers()? { + return Err(StoreError::Busy.into()); + } self.close_connections_without_checkpoint()?; if journal.after == ResetAfter::Erase { self.lifecycle_control @@ -1373,6 +1376,37 @@ fn set_state(store: &Store, state: StoreState) -> StoreResult<()> { Ok(()) } +/// Freeze new upload admission before cancelling existing writers. A refusal +/// before durable replacement intent leaves the original store usable while +/// those already-cancelled handles drain; it must not strand the UI in upkeep. +fn freeze_attachment_writers(store: &Store, restore_on_busy: bool) -> StoreResult<()> { + let previous = { + let mut state = store.state.write().map_err(|_| StoreError::Busy)?; + let previous = state.clone(); + *state = StoreState::Maintenance; + previous + }; + if store.cancel_attachment_writers()? { + return Ok(()); + } + if restore_on_busy { + let mut state = store.state.write().map_err(|_| StoreError::Busy)?; + if *state == StoreState::Maintenance { + // A checkpoint failure can race the request without taking the + // lifecycle gate. Never revive Open after damage was observed. + *state = if previous == StoreState::Open && store.runtime_damage.load(Ordering::Acquire) + { + StoreState::Damaged { + reason: DamageReason::Database, + } + } else { + previous + }; + } + } + Err(StoreError::Busy) +} + fn disk_mode(paths: &StorePaths) -> LifecycleResult { if path_present(&paths.keys())? || path_present(&paths.db().join("keys.backup.json"))? { return Ok(StoreMode::Password); @@ -1563,6 +1597,7 @@ fn recover_filesystem(paths: &StorePaths) -> LifecycleResult LifecycleResult LifecycleResult<()> { journal.validate()?; + attachments::verify_generation(paths, &journal.old_store_id)?; let replacement = journal.replacement(); let current = read_generation(&paths.db())?; if journal.phase == SetupPhase::BuildReplacement @@ -1728,6 +1764,7 @@ fn retire_old(paths: &StorePaths, journal: &mut SetupJournal) -> LifecycleResult // A confirmed reset discards its current source before asking for new // setup choices. A committed salvage keeps its temporary source until // the verified replacement has been installed. + attachments::remove_generation(paths, &journal.old_store_id)?; remove_retired_source(paths, &journal.old_store_id)?; } if journal.phase == SetupPhase::RetireOld { @@ -1790,6 +1827,7 @@ fn finish_setup(paths: &StorePaths, journal: &mut SetupJournal) -> LifecycleResu journal.validate()?; // A crash may have installed the replacement before source cleanup. Its // journal stays authoritative until the exact old source is gone. + attachments::remove_generation(paths, &journal.old_store_id)?; remove_retired_source(paths, &journal.old_store_id)?; journal.phase = SetupPhase::Installed; write_json(&paths.setup(), journal)?; @@ -2109,6 +2147,7 @@ fn decode_id(value: &str) -> Option<[u8; 16]> { Some(id) } +mod attachments; #[cfg(test)] mod generated_tests; mod generation_repair; diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/attachments.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/attachments.rs new file mode 100644 index 000000000..d016b48f5 --- /dev/null +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/attachments.rs @@ -0,0 +1,300 @@ +//! Attachment filesystem ownership during generation replacement. +//! +//! A reset removes only the named old generation. Salvage copies retained ready +//! files before publishing its replacement database, keeping the two generations +//! independent and making both crash cleanup and later reset unambiguous. + +use sha2::{Digest, Sha256}; + +use super::{ + unclassified_identity::{FileIdentity, capture_path}, + *, +}; + +const COPY_BYTES: usize = 1024 * 1024; +const ROW_BATCH: i64 = 100; + +pub(super) fn root_identity(paths: &StorePaths) -> StoreResult> { + let root = paths.root().join("attachments"); + match capture_path(&root, true) { + Ok((identity, _)) => Ok(Some(identity)), + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => Ok(None), + Err(error) => Err(error), + } +} + +pub(super) fn verify_root(paths: &StorePaths, expected: &Option) -> StoreResult<()> { + let root = paths.root().join("attachments"); + match (capture_path(&root, true), expected) { + (Ok((actual, _)), Some(expected)) if &actual == expected => Ok(()), + ( + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }), + _, + ) => Ok(()), + (Err(error), _) => Err(error), + _ => Err(StoreError::UnsafePath), + } +} + +pub(super) fn remove_root(paths: &StorePaths, expected: &Option) -> StoreResult<()> { + verify_root(paths, expected)?; + let root = paths.root().join("attachments"); + if path_present(&root).map_err(startup_error)? { + collect_tree(&root)?; + } + sync_directory(paths.root())?; + Ok(()) +} + +pub(super) fn verify_generation(paths: &StorePaths, id: &str) -> StoreResult<()> { + if decode_id(id).is_none() { + return Err(StoreError::InvalidData); + } + let root = paths.root().join("attachments"); + if !path_present(&root).map_err(startup_error)? { + return Ok(()); + } + verify_private_path(&root, true)?; + let generation = root.join(id); + if path_present(&generation).map_err(startup_error)? { + capture_path(&generation, true)?; + } + Ok(()) +} + +pub(super) fn remove_generation(paths: &StorePaths, id: &str) -> StoreResult<()> { + verify_generation(paths, id)?; + let root = paths.root().join("attachments"); + let generation = root.join(id); + if path_present(&generation).map_err(startup_error)? { + collect_tree(&generation)?; + } + if path_present(&root).map_err(startup_error)? { + sync_directory(&root)?; + } + Ok(()) +} + +fn collect_tree(path: &Path) -> StoreResult<()> { + while !crate::store::attachment_cleanup::collect_owned_attachment_tree_step(path)?.done {} + Ok(()) +} + +struct RetainedFile { + id: String, + account: i64, + task: String, + name: String, + size: u64, + hash: String, +} + +pub(super) fn copy_retained( + paths: &StorePaths, + destination: &Connection, + old_id: &str, + new_id: &str, + control: &LifecycleControl, +) -> LifecycleResult<()> { + if old_id == new_id || decode_id(old_id).is_none() || decode_id(new_id).is_none() { + return Err(LifecycleError::Generation); + } + // This is the unpublished replacement namespace from this exact journal. + // A failed previous build may have left a strict subset of its files. + remove_generation(paths, new_id)?; + let mut cursor = String::new(); + loop { + control.check_cancelled()?; + let batch = { + let mut statement = destination.prepare( + "SELECT a.public_id,a.account_id,t.uuid,a.name,a.size,a.sha256 FROM file_attachments a \ + JOIN tasks t ON t.id=a.task_id WHERE a.state='ready' AND a.public_id>?1 ORDER BY a.public_id LIMIT ?2")?; + statement + .query_map(params![cursor, ROW_BATCH], |row| { + Ok(RetainedFile { + id: row.get(0)?, + account: row.get(1)?, + task: row.get(2)?, + name: row.get(3)?, + size: u64::try_from(row.get::<_, i64>(4)?) + .map_err(|_| rusqlite::Error::IntegralValueOutOfRange(4, -1))?, + hash: row.get(5)?, + }) + })? + .collect::, _>>()? + }; + if batch.is_empty() { + return Ok(()); + } + for file in batch { + control.check_cancelled()?; + cursor.clone_from(&file.id); + if file.account <= 0 + || decode_id(&file.id).is_none() + || !component(&file.task) + || !component(&file.name) + || file.size > 30 * 1024 * 1024 + || file.hash.len() != 64 + || !file.hash.bytes().all(|byte| byte.is_ascii_hexdigit()) + { + return Err(StoreError::InvalidData.into()); + } + let source = match file_path(paths, old_id, &file, false) { + Ok(path) => path, + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => continue, + Err(error) => return Err(error.into()), + }; + let (identity, size) = match capture_path(&source, false) { + Ok(value) => value, + Err(StoreError::Io { + kind: io::ErrorKind::NotFound, + }) => continue, + Err(error) => return Err(error.into()), + }; + if size != file.size { + return Err(StoreError::InvalidData.into()); + } + let mut input = open_input(&source)?; + if capture_path(&source, false)?.0 != identity { + return Err(StoreError::UnsafePath.into()); + } + let target = file_path(paths, new_id, &file, true)?; + let partial = target.parent().ok_or(StoreError::UnsafePath)?.join(".part"); + let mut output = private_new_file(&partial)?; + let mut hash = Sha256::new(); + let mut copied = 0_u64; + let mut bytes = vec![0_u8; COPY_BYTES]; + loop { + control.check_cancelled()?; + let count = input.read(&mut bytes)?; + if count == 0 { + break; + } + copied = copied + .checked_add(count as u64) + .ok_or(StoreError::InvalidData)?; + if copied > file.size { + return Err(StoreError::InvalidData.into()); + } + hash.update(&bytes[..count]); + output.write_all(&bytes[..count])?; + } + if copied != file.size + || format!("{:x}", hash.finalize()) != file.hash + || capture_path(&source, false)?.0 != identity + { + return Err(StoreError::InvalidData.into()); + } + keys::sync_key_file(&output)?; + drop(output); + if path_present(&target)? { + return Err(StoreError::Conflict.into()); + } + keys::replace_key_file(&partial, &target)?; + sync_directory(target.parent().ok_or(StoreError::UnsafePath)?)?; + } + } +} + +fn component(value: &str) -> bool { + let mut parts = Path::new(value).components(); + matches!(parts.next(), Some(std::path::Component::Normal(_))) + && parts.next().is_none() + && !value.contains(['/', '\\']) +} + +fn file_path( + paths: &StorePaths, + generation: &str, + file: &RetainedFile, + create: bool, +) -> StoreResult { + let mut path = paths.root().to_owned(); + verify_private_path(&path, true)?; + for part in [ + "attachments".to_owned(), + generation.to_owned(), + file.account.to_string(), + file.task.clone(), + file.id.clone(), + ] { + path.push(part); + if create { + private_directory(&path)?; + } else { + verify_private_path(&path, true)?; + } + } + Ok(path.join(&file.name)) +} + +fn open_input(path: &Path) -> io::Result { + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC | libc::O_NONBLOCK); + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + options.custom_flags(windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT); + } + options.open(path) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(unix)] + #[test] + fn generation_cleanup_unlinks_model_links_without_touching_their_targets() { + use std::os::unix::fs::symlink; + let directory = tempfile::tempdir().unwrap(); + let paths = StorePaths::new(directory.path()).prepare().unwrap(); + let root = paths.root().join("attachments"); + private_directory(&root).unwrap(); + let generation = root.join(encode_id(&[7; 16])); + private_directory(&generation).unwrap(); + let mut owned = private_new_file(&generation.join("owned-file")).unwrap(); + owned + .write_all(b"keep until the entire tree is qualified") + .unwrap(); + drop(owned); + let outside = directory.path().join("outside"); + fs::write(&outside, b"never follow this link").unwrap(); + symlink(&outside, generation.join("unsafe-file")).unwrap(); + remove_generation(&paths, &encode_id(&[7; 16])).unwrap(); + assert!(!generation.exists()); + assert_eq!(fs::read(outside).unwrap(), b"never follow this link"); + } + + #[test] + fn exact_generation_cleanup_is_repeatable_after_a_partial_delete() { + let directory = tempfile::tempdir().unwrap(); + let paths = StorePaths::new(directory.path()).prepare().unwrap(); + let root = paths.root().join("attachments"); + private_directory(&root).unwrap(); + let generation = root.join(encode_id(&[7; 16])); + private_directory(&generation).unwrap(); + let nested = generation.join("1"); + private_directory(&nested).unwrap(); + let mut owned = private_new_file(&nested.join("remaining-file")).unwrap(); + owned.write_all(b"remove after retry").unwrap(); + drop(owned); + let other = root.join(encode_id(&[8; 16])); + private_directory(&other).unwrap(); + remove_generation(&paths, &encode_id(&[7; 16])).unwrap(); + remove_generation(&paths, &encode_id(&[7; 16])).unwrap(); + assert!(!generation.exists()); + assert!(other.exists()); + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/salvage.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/salvage.rs index 49143a045..1c285e275 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/salvage.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/salvage.rs @@ -169,7 +169,7 @@ impl Store { }; self.lifecycle_control.begin(); let _operation = ActiveOperation(&self.lifecycle_control); - set_state(self, StoreState::Maintenance)?; + freeze_attachment_writers(self, true)?; // Do not checkpoint or truncate the damaged source, including on close. if let Err(error) = self.close_connections_without_checkpoint() { set_state( @@ -240,6 +240,13 @@ impl Store { return Err(StoreError::InvalidData.into()); } validate_replacement(&destination, &journal)?; + attachments::copy_retained( + &self.paths, + &destination, + &journal.old_store_id, + &journal.new_store_id, + &self.lifecycle_control, + )?; destination .close() .map_err(|(_, error)| StoreError::from(error))?; @@ -310,6 +317,7 @@ impl Store { if path_present(&temporary)? { discard_temporary(&temporary, &generation)?; } + attachments::remove_generation(&self.paths, &journal.new_store_id)?; *self.dek.lock().map_err(|_| StoreError::Busy)? = source_key; set_state( self, diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/tests.rs index ab2355740..e94fa5797 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/tests.rs @@ -572,6 +572,19 @@ fn confirmed_reset_discards_only_current_source_and_survives_cold_reopen() { .unwrap(); } let old = store.store_id().unwrap(); + let attachment_root = store.paths.root().join("attachments"); + private_directory(&attachment_root).unwrap(); + let old_attachments = attachment_root.join(encode_id(&old)); + private_directory(&old_attachments).unwrap(); + let mut original = private_new_file(&old_attachments.join("current-file")).unwrap(); + original.write_all(b"current session attachment").unwrap(); + drop(original); + let historical_attachments = attachment_root.join(encode_id(&[89; 16])); + private_directory(&historical_attachments).unwrap(); + let mut previous = + private_new_file(&historical_attachments.join("historical-file")).unwrap(); + previous.write_all(b"keep unrelated generation").unwrap(); + drop(previous); store .upsert_account("current-history-to-discard", 1) .unwrap(); @@ -604,6 +617,11 @@ fn confirmed_reset_discards_only_current_source_and_survives_cold_reopen() { } let replacement = store.store_id().unwrap(); assert_ne!(replacement, old); + assert!(!old_attachments.exists()); + assert_eq!( + fs::read(historical_attachments.join("historical-file")).unwrap(), + b"keep unrelated generation" + ); assert!( store .account_id("current-history-to-discard") @@ -1408,3 +1426,160 @@ fn status_read_failure_does_not_reclassify_an_open_generation() { .upsert_account("status-retry-remains-usable", 1) .unwrap(); } + +#[test] +fn reset_busy_on_an_upload_keeps_the_store_usable_until_retry() { + use crate::store::{AttachmentUploadMetadata, AttachmentUploadStart, NewTask, TaskKind}; + use sha2::{Digest, Sha256}; + let directory = tempfile::tempdir().unwrap(); + let store = opened(directory.path()); + let account = store.upsert_account("reset-upload", 1).unwrap(); + let task = store + .create_task( + account, + &NewTask { + uuid: "00000000-0000-7000-8000-000000000001".into(), + project_root: if cfg!(windows) { + "C:/reset-upload" + } else { + "/reset-upload" + } + .into(), + kind: TaskKind::Desktop, + title: "Upload".into(), + provisional: true, + model: None, + created_ms: 1, + }, + ) + .unwrap(); + let AttachmentUploadStart::Writing(mut upload) = store + .begin_attachment_upload( + account, + &task.uuid, + "reset-active-upload", + &AttachmentUploadMetadata { + name: "file.txt".into(), + mime_type: "text/plain".into(), + size: 3, + sha256: format!("{:x}", Sha256::digest(b"abc")), + }, + ) + .unwrap() + else { + panic!("fresh upload unexpectedly complete"); + }; + upload.write_chunk(b"a").unwrap(); + assert_eq!( + store.erase_local_data(ERASE_CONFIRMATION), + Err(LifecycleError::Store(StoreError::Busy)) + ); + assert_eq!(store.state(), StoreState::Open); + assert_eq!(store.account_id("reset-upload").unwrap(), Some(account)); + assert!(!store.paths.setup().exists()); + assert_eq!(upload.write_chunk(b"bc"), Err(StoreError::Conflict)); + drop(upload); + store.erase_local_data(ERASE_CONFIRMATION).unwrap(); + assert_eq!(store.state(), StoreState::Open); + assert_eq!(store.account_id("reset-upload").unwrap(), None); +} + +#[cfg(unix)] +fn readonly_attachment_generation(store: &Store) -> PathBuf { + use std::os::unix::fs::PermissionsExt; + let root = store.paths.attachments(); + private_directory(&root).unwrap(); + let generation = root.join(encode_id(&store.store_id().unwrap())); + let mut current = generation.clone(); + private_directory(¤t).unwrap(); + for name in [ + "1", + "00000000-0000-7000-8000-000000000001", + "0123456789abcdef0123456789abcdef", + ] { + current.push(name); + private_directory(¤t).unwrap(); + } + fs::write(current.join("archive.zip"), b"archive fixture").unwrap(); + let nested = current.join("unpacked").join("nested"); + fs::create_dir_all(&nested).unwrap(); + fs::write(nested.join("plaintext.txt"), b"owned extracted plaintext").unwrap(); + fs::set_permissions(&nested, fs::Permissions::from_mode(0o555)).unwrap(); + fs::set_permissions(nested.parent().unwrap(), fs::Permissions::from_mode(0o500)).unwrap(); + generation +} + +#[cfg(unix)] +#[test] +fn reset_removes_readonly_attachment_trees_and_cold_reopens_cleanly() { + for after in [ResetAfter::Erase, ResetAfter::StartFresh] { + let directory = tempfile::tempdir().unwrap(); + let store = opened(directory.path()); + store.upsert_account("old-readonly-reset", 1).unwrap(); + let generation = readonly_attachment_generation(&store); + match after { + ResetAfter::Erase => { + store.erase_local_data(ERASE_CONFIRMATION).unwrap(); + } + ResetAfter::StartFresh => { + set_state( + &store, + StoreState::Damaged { + reason: DamageReason::Database, + }, + ) + .unwrap(); + store.start_fresh(START_FRESH_CONFIRMATION).unwrap(); + } + } + assert_eq!(store.state(), StoreState::Open); + assert!(!generation.exists()); + assert!(!store.paths.setup().exists()); + drop(store); + let reopened = opened(directory.path()); + assert_eq!(reopened.account_id("old-readonly-reset").unwrap(), None); + assert!(!generation.exists()); + assert!(!reopened.paths.setup().exists()); + } +} + +#[cfg(unix)] +#[test] +fn interrupted_reset_reclaims_readonly_attachment_trees_on_cold_recovery() { + fn fail_sync(_: &Path) -> io::Result<()> { + Err(io::Error::from(io::ErrorKind::Other)) + } + for renamed in [false, true] { + let directory = tempfile::tempdir().unwrap(); + let store = opened(directory.path()); + store.upsert_account("old-interrupted-reset", 1).unwrap(); + let generation = readonly_attachment_generation(&store); + { + let _gate = store.lifecycle.lock().unwrap(); + assert_eq!( + store.begin_reset_with_intent_sync_locked(ResetAfter::Erase, fail_sync), + Err(LifecycleError::RecoveryRequired { + mode: StoreMode::Plain + }), + ); + } + assert!(store.paths.setup().exists()); + assert!(generation.exists()); + if renamed { + durable_rename( + &store.paths.db(), + &store.paths.root().join("db.old"), + store.paths.root(), + ) + .unwrap(); + } + drop(store); + let reopened = opened(directory.path()); + assert_eq!(reopened.account_id("old-interrupted-reset").unwrap(), None); + assert!(!generation.exists()); + assert!(!reopened.paths.setup().exists()); + assert!(!reopened.paths.root().join("db.old").exists()); + drop(reopened); + assert_eq!(opened(directory.path()).state(), StoreState::Open); + } +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified.rs index 10579e2a2..492eed56f 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified.rs @@ -6,7 +6,7 @@ use super::{unclassified_identity::DirectoryManifest, *}; -const FORMAT: &str = "maple-store-unclassified-reset/2"; +const FORMAT: &str = "maple-store-unclassified-reset/3"; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -26,6 +26,7 @@ pub(super) struct Journal { previous_mode: StoreMode, phase: Phase, source: DirectoryManifest, + attachments: Option, } impl Journal { @@ -123,6 +124,7 @@ pub(super) fn recover(paths: &StorePaths) -> LifecycleResult> { fn preflight(paths: &StorePaths, journal: &Journal) -> LifecycleResult<()> { journal.validate()?; verify_private_path(paths.root(), true)?; + attachments::verify_root(paths, &journal.attachments)?; if path_present(&paths.setup())? { return Err(LifecycleError::Generation); } @@ -209,6 +211,7 @@ fn retire_original(paths: &StorePaths, journal: &mut Journal) -> LifecycleResult // Both confirmed reset variants permanently discard their exact current // source. The manifest permits partial deletion after a crash, never a // replacement file or a directory belonging to another operation. + attachments::remove_root(paths, &journal.attachments)?; journal.source.delete_remaining(&archive)?; if journal.phase == Phase::RetireOriginal { journal.phase = Phase::AwaitChoice; @@ -221,6 +224,7 @@ fn finish(paths: &StorePaths, journal: &Journal) -> LifecycleResult<()> { if read_generation(&paths.db())?.as_ref() != Some(&journal.replacement()) { return Err(LifecycleError::Generation); } + attachments::remove_root(paths, &journal.attachments)?; journal .source .delete_remaining(&journal.retired_source(paths))?; @@ -233,6 +237,7 @@ fn finish(paths: &StorePaths, journal: &Journal) -> LifecycleResult<()> { fn finish_installed(paths: &StorePaths, journal: &Journal) -> LifecycleResult<()> { // The installed pair wins on retry. Cleanup can touch only this journal's // exact temporary source, not a later active generation or historical copy. + attachments::remove_root(paths, &journal.attachments)?; journal .source .delete_remaining(&journal.retired_source(paths))?; @@ -252,7 +257,7 @@ impl Store { after: ResetAfter, intent_sync: fn(&Path) -> io::Result<()>, ) -> LifecycleResult<()> { - set_state(self, StoreState::Maintenance)?; + freeze_attachment_writers(self, true)?; self.close_connections_without_checkpoint()?; if let Some(mut journal) = pending(&self.paths)? { if matches!(read_generation(&self.paths.db()),Ok(Some(actual)) if actual==journal.replacement()) @@ -292,6 +297,7 @@ impl Store { previous_mode: mode, phase: Phase::RetireOriginal, source: DirectoryManifest::capture(&self.paths.db())?, + attachments: attachments::root_identity(&self.paths)?, }; preflight(&self.paths, &journal)?; if write_json_with_sync(&intent_path(&self.paths), &journal, intent_sync).is_err() { diff --git a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified_tests.rs b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified_tests.rs index b01c9a00e..12cc98ebe 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified_tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/lifecycle/unclassified_tests.rs @@ -44,6 +44,7 @@ fn fixture_journal(paths: &StorePaths, after: ResetAfter) -> Journal { previous_mode: StoreMode::Password, phase: Phase::RetireOriginal, source: DirectoryManifest::capture(&paths.db()).unwrap(), + attachments: attachments::root_identity(paths).unwrap(), } } @@ -476,3 +477,45 @@ fn installed_generation_wins_over_reused_source_directory_identity() { assert!(!intent_path(&store.paths).exists()); assert!(!store.paths.root().join("db.old").exists()); } + +#[test] +fn unclassified_reset_removes_only_its_captured_attachment_root() { + let directory = tempfile::tempdir().unwrap(); + let store = unqualified(directory.path()); + let root = store.paths.root().join("attachments"); + private_directory(&root).unwrap(); + let nested = root.join("captured"); + private_directory(&nested).unwrap(); + let mut file = private_new_file(&nested.join("file.txt")).unwrap(); + file.write_all(b"owned upload").unwrap(); + drop(file); + store.start_fresh(START_FRESH_CONFIRMATION).unwrap(); + assert_eq!(store.state(), StoreState::Open); + assert!(!root.exists()); +} + +#[test] +fn unclassified_reset_refuses_a_replaced_attachment_root() { + let directory = tempfile::tempdir().unwrap(); + let store = unqualified(directory.path()); + let root = store.paths.root().join("attachments"); + private_directory(&root).unwrap(); + let journal = fixture_journal(&store.paths, ResetAfter::Erase); + write_json(&intent_path(&store.paths), &journal).unwrap(); + let original = directory.path().join("preserved-original-attachments"); + fs::rename(&root, &original).unwrap(); + private_directory(&root).unwrap(); + let mut file = private_new_file(&root.join("other-file")).unwrap(); + file.write_all(b"do not delete replacement").unwrap(); + drop(file); + assert!(matches!( + recover(&store.paths), + Err(LifecycleError::Store(StoreError::UnsafePath)) + )); + assert_eq!( + fs::read(root.join("other-file")).unwrap(), + b"do not delete replacement" + ); + assert!(original.exists()); + assert!(store.paths.database().exists()); +} diff --git a/apps/maple-agent/crates/maple-agent/src/store/mod.rs b/apps/maple-agent/crates/maple-agent/src/store/mod.rs index e13d48d02..88101a738 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/mod.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/mod.rs @@ -3,9 +3,11 @@ //! Only typed APIs leave this module. Raw connections and transactions are //! restricted to the store subtree, and errors never retain SQL or row values. +mod attachment_cleanup; +mod attachment_migration; +mod attachments; #[cfg(feature = "store-bench")] pub mod bench; -mod blobs; mod codec; mod connection; #[cfg(test)] @@ -45,9 +47,12 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use rusqlite::{Connection, OptionalExtension, Transaction, TransactionBehavior, params}; use serde::{Deserialize, Serialize}; -pub use blobs::{ - AttachmentRange, BlobCollectionProgress, ImageUpload, MAX_ATTACHMENT_RANGE_BYTES, - StoredAttachment, +pub use attachments::{ + ATTACHMENT_SOURCE_PREFIX, AttachmentCollectionProgress, AttachmentRange, + AttachmentUploadMetadata, AttachmentUploadStart, AttachmentUploadStatus, + AttachmentUploadWriter, ImageUpload, MAX_ATTACHMENT_BYTES, MAX_ATTACHMENT_RANGE_BYTES, + MAX_ATTACHMENTS, MAX_UPLOAD_IMAGE_BYTES, PreparedAttachmentUpload, StoredAttachment, + UNAVAILABLE_IMAGE_SOURCE, }; pub use error::{DamageReason, RefusalReason, StoreError, StoreResult}; pub use gc::{CollectionProgress, CollectionStage, DeletedTask}; @@ -149,7 +154,8 @@ pub struct Store { lifecycle_control: lifecycle::LifecycleControl, next_reader: AtomicUsize, foreground_waiters: AtomicUsize, - pub(super) blob_locks: Arc, + pub(super) attachment_locks: Arc, + collection_cursors: Mutex, changes: tokio::sync::broadcast::Sender, persistence: persistence::PersistenceState, replays: replay::ReplayRegistry, @@ -187,7 +193,8 @@ impl Store { lifecycle_control: lifecycle::LifecycleControl::default(), next_reader: AtomicUsize::new(0), foreground_waiters: AtomicUsize::new(0), - blob_locks: Arc::new(blobs::BlobLocks::default()), + attachment_locks: Arc::new(attachments::AttachmentLocks::default()), + collection_cursors: Mutex::new(gc::CollectionCursors::default()), changes: tokio::sync::broadcast::channel(256).0, persistence: persistence::PersistenceState::default(), replays: replay::ReplayRegistry::default(), diff --git a/apps/maple-agent/crates/maple-agent/src/store/paging.rs b/apps/maple-agent/crates/maple-agent/src/store/paging.rs index 8d58ee68a..56f75ceb3 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/paging.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/paging.rs @@ -225,6 +225,22 @@ struct TaskSnapshot { } impl Store { + /// A format migration can intentionally remove the old derived rows. The + /// host must reconstruct that task before presenting it as an empty chat. + /// Ordinary background rebuilds retain usable rows and need no foreground + /// wait. This read never instantiates a Pi session or loads source payloads. + pub(crate) fn history_needs_rebuild( + &self, + account: AccountId, + task: &str, + ) -> StoreResult { + validate_task(task)?; + self.read(|tx| { + let snapshot = task_snapshot(tx, account, task)?; + history_needs_rebuild(tx, &snapshot) + }) + } + pub fn task_history_stamp( &self, account: AccountId, @@ -280,6 +296,9 @@ impl Store { validate_page_limits(rows, bytes, 500, caps.max_reply_bytes)?; self.read(|tx| { let task = task_snapshot(tx, account, &request.task)?; + if history_needs_rebuild(tx, &task)? { + return Err(StoreError::Busy); + } let mut summary = navigation::task_summary_tx(tx, account, &request.task)? .ok_or(StoreError::NotFound)?; let inline = caps.inline_row_bytes.min((bytes / 3).max(256)); @@ -466,6 +485,17 @@ impl Store { } } +fn history_needs_rebuild(tx: &Transaction<'_>, task: &TaskSnapshot) -> StoreResult { + if task.format == projector::TIMELINE_FORMAT { + return Ok(false); + } + Ok(tx.query_row( + "SELECT last_seq>0 AND NOT EXISTS(SELECT 1 FROM timeline_rows WHERE task_id=?1 AND epoch=?2) FROM tasks WHERE id=?1", + params![task.key.0, task.epoch], + |row| row.get(0), + )?) +} + fn task_snapshot( tx: &Transaction<'_>, account: AccountId, diff --git a/apps/maple-agent/crates/maple-agent/src/store/paths.rs b/apps/maple-agent/crates/maple-agent/src/store/paths.rs index a50b9532f..52e2a8008 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/paths.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/paths.rs @@ -82,6 +82,10 @@ impl StorePaths { Ok(file) } + pub fn attachments(&self) -> PathBuf { + self.root.join("attachments") + } + pub fn logs(&self) -> PathBuf { self.root.join("logs") } diff --git a/apps/maple-agent/crates/maple-agent/src/store/projector.rs b/apps/maple-agent/crates/maple-agent/src/store/projector.rs index df136b231..8362fd7d0 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/projector.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/projector.rs @@ -6,6 +6,7 @@ //! prepared on the bulk reader and replayed in bounded idle transactions. use std::collections::{HashMap, HashSet}; +use std::time::Instant; use pi_ai::{Message, content_text}; use pi_coding_agent::{ @@ -25,7 +26,7 @@ mod stored; use stored::StoredItem; /// Bump whenever the shared timeline builders or their merge rules change. -pub(crate) const TIMELINE_FORMAT: i64 = 1; +pub(crate) const TIMELINE_FORMAT: i64 = 2; const CHUNK_ENTRIES: usize = 1_000; const CHUNK_ROWS: usize = 1_000; const CHUNK_BYTES: usize = 1024 * 1024; @@ -78,8 +79,29 @@ struct PreparedSourceOperation { result_call: Option, } -/// Complete source-sized formatting before entering the writer transaction. +/// Raw-entry fixture formatting. Production projects the persisted image form. +#[cfg(test)] pub(super) fn prepare_entry(entry: &SessionEntry, seq: i64) -> StoreResult { + prepare_operations(entry, seq, operations(entry, |_| Ok(None))?) +} + +/// Complete source-sized formatting before entering the writer transaction. +/// Image lengths come from metadata collected separately, never file bytes. +pub(super) fn prepare_stored_entry( + entry: &SessionEntry, + seq: i64, + image_sizes: &[usize], +) -> StoreResult { + let mut rows = operations(entry, |_| Ok(None))?; + apply_image_sizes(entry, &mut rows, image_sizes)?; + prepare_operations(entry, seq, rows) +} + +fn prepare_operations( + entry: &SessionEntry, + seq: i64, + rows: Vec, +) -> StoreResult { if seq < 1 { return Err(StoreError::InvalidInput); } @@ -89,7 +111,7 @@ pub(super) fn prepare_entry(entry: &SessionEntry, seq: i64) -> StoreResult Some(result.tool_call_id.clone()), _ => None, }; - let operations = operations(entry, |_| Ok(None))? + let operations = rows .into_iter() .enumerate() .map(|(index, operation)| { @@ -221,7 +243,7 @@ fn operations( } // Image payloads in stored entries are references. Timeline rows expose only -// their length, so rebuilds read blob metadata rather than loading image bytes. +// their length, so append and rebuild use metadata without loading image bytes. fn operations_tx( tx: &Transaction<'_>, task: TaskKey, @@ -235,9 +257,21 @@ fn operations_tx( message: SessionMessage::Llm(Message::ToolResult(_)) } ) { - let sizes = super::blobs::entry_image_sizes_tx(tx, task, entry)?; - let mut sizes = sizes.into_iter(); - for row in &mut rows { + let sizes = super::attachments::entry_image_sizes_tx(tx, task, entry)?; + apply_image_sizes(entry, &mut rows, &sizes)?; + } + Ok(rows) +} + +fn apply_image_sizes(entry: &SessionEntry, rows: &mut [RowOp], sizes: &[usize]) -> StoreResult<()> { + if matches!( + &entry.kind, + EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(_)) + } + ) { + let mut sizes = sizes.iter(); + for row in rows { if let Some(content) = row .item .output @@ -257,7 +291,7 @@ fn operations_tx( return Err(StoreError::InvalidData); } } - Ok(rows) + Ok(()) } /// The UI's todo predicate, shared with persisted plan-page selection. @@ -857,6 +891,16 @@ impl Store { pub(crate) fn prepare_timeline_rebuild( &self, ticket: &RebuildTicket, + ) -> StoreResult { + self.prepare_timeline_rebuild_before(ticket, None) + } + + /// Foreground opens use a cooperative deadline between bounded source reads; + /// unfinished work remains queued for normal background reconstruction. + pub(crate) fn prepare_timeline_rebuild_before( + &self, + ticket: &RebuildTicket, + deadline: Option, ) -> StoreResult { let last_seq = self.read_bulk(|tx| { authorize_ticket(tx, ticket)?; @@ -870,6 +914,7 @@ impl Store { let mut after = 0; let mut compact = true; while after < last_seq { + check_rebuild_deadline(deadline)?; let entries = self.read_bulk(|tx| source_metadata_batch(tx, ticket.task, after, last_seq))?; if entries.is_empty() { @@ -909,6 +954,15 @@ impl Store { &self, prepared: &PreparedRebuild, ) -> StoreResult { + self.prepare_timeline_chunk_before(prepared, None) + } + + pub(crate) fn prepare_timeline_chunk_before( + &self, + prepared: &PreparedRebuild, + deadline: Option, + ) -> StoreResult { + check_rebuild_deadline(deadline)?; let cursor = self.read_bulk(|tx| { authorize_ticket(tx, &prepared.ticket)?; rebuild_cursor(tx, &prepared.ticket) @@ -919,7 +973,7 @@ impl Store { let end = (start + CHUNK_ENTRIES).min(outline.len()); (outline[start..end].to_vec(), end == outline.len()) } - Outline::Rescan => self.rescan_outline_chunk(prepared, cursor.0)?, + Outline::Rescan => self.rescan_outline_chunk(prepared, cursor.0, deadline)?, }; self.read_bulk(|tx| { authorize_ticket(tx, &prepared.ticket)?; @@ -946,10 +1000,12 @@ impl Store { fn walk_source_branch( &self, prepared: &PreparedRebuild, + deadline: Option, mut visit: impl FnMut(i64, &SessionEntry) -> StoreResult<()>, ) -> StoreResult<()> { let mut seq = prepared.last_seq; while seq > 0 { + check_rebuild_deadline(deadline)?; let entry = self.read_bulk(|tx| session::read_entry(tx, prepared.ticket.task, seq))?; visit(seq, &entry)?; seq = match entry.parent_id { @@ -978,12 +1034,13 @@ impl Store { &self, prepared: &PreparedRebuild, through: i64, + deadline: Option, ) -> StoreResult<(Vec, bool)> { use std::collections::VecDeque; let mut selected = VecDeque::<(i64, String)>::new(); let mut count = 0usize; let mut id_bytes = 0usize; - self.walk_source_branch(prepared, |seq, entry| { + self.walk_source_branch(prepared, deadline, |seq, entry| { if seq > through { count = count.saturating_add(1); id_bytes = id_bytes.saturating_add(entry.id.len()); @@ -1011,7 +1068,7 @@ impl Store { omitted: false, }) .collect(); - self.walk_source_branch(prepared, |_, entry| { + self.walk_source_branch(prepared, deadline, |_, entry| { if let EntryKind::ContextEdit { target_id, replacement: None, @@ -1026,6 +1083,13 @@ impl Store { } } +fn check_rebuild_deadline(deadline: Option) -> StoreResult<()> { + if deadline.is_some_and(|deadline| Instant::now() >= deadline) { + return Err(StoreError::Busy); + } + Ok(()) +} + struct CompactMeta { seq: i64, id: u32, diff --git a/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs index e08c907ec..dc4b8d2ec 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/projector/tests.rs @@ -414,7 +414,7 @@ fn timeline_format_pins_the_canonical_fixture() { assert_eq!( (TIMELINE_FORMAT, digest.as_str()), ( - 1, + 2, "bd6b8487d68c8376bcda6a6d7323a632c5bd942368dbe742a9de26440af852c2" ) ); diff --git a/apps/maple-agent/crates/maple-agent/src/store/schema.rs b/apps/maple-agent/crates/maple-agent/src/store/schema.rs index 87653ee6f..3e9d3e80d 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/schema.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/schema.rs @@ -6,7 +6,7 @@ use super::{StoreError, StoreResult}; pub(super) const APPLICATION_ID: i64 = 1_296_126_028; pub(super) const SCHEMA_EPOCH: i64 = 1; -pub(super) const VERSION: i64 = 2; +pub(super) const VERSION: i64 = 4; /// Original migration schema, including partial indexes and constraints. /// Validation executes the complete append-only migration chain as its reference. @@ -260,6 +260,40 @@ struct Migration { const PROJECT_NAME_REVISION_DDL: &str = "ALTER TABLE projects ADD COLUMN name_revision INTEGER NOT NULL DEFAULT 0 CHECK (name_revision >= 0);"; +// Version 3 is an intentional durable fence. Released readers stop here before +// any image source is rewritten; a new reader can resume the bounded retirement +// work with the original tables still present. Version 4 removes those tables +// only after source retirement and byte collection have completed. +const FILE_ATTACHMENTS_DDL: &str = r#" +CREATE TABLE file_attachments ( + public_id TEXT PRIMARY KEY CHECK (length(public_id) = 32), + account_id INTEGER NOT NULL REFERENCES accounts(id) ON DELETE CASCADE, + task_id INTEGER NOT NULL REFERENCES tasks(id) ON DELETE CASCADE, + operation_id TEXT NOT NULL, + name TEXT NOT NULL, + mime TEXT NOT NULL, + size INTEGER NOT NULL CHECK (size >= 0 AND size <= 31457280), + sha256 TEXT NOT NULL CHECK (length(sha256) = 64), + state TEXT NOT NULL CHECK (state IN ('uploading', 'ready', 'cancelled')), + bound INTEGER NOT NULL DEFAULT 0 CHECK (bound IN (0, 1)), + created_ms INTEGER NOT NULL, + updated_ms INTEGER NOT NULL, + UNIQUE (account_id, operation_id) +) STRICT; +CREATE INDEX file_attachments_task ON file_attachments (task_id); +CREATE INDEX file_attachments_unbound ON file_attachments (updated_ms) WHERE bound = 0; +INSERT INTO store_meta(key,value) VALUES ('attachment_migration_phase',0); +INSERT INTO store_meta(key,value) VALUES ('attachment_migration_cursor',0); +"#; + +const RETIRE_BLOBS_DDL: &str = r#" +DROP TABLE attachments; +DROP TABLE task_blobs; +DROP TABLE blob_chunks; +DROP TABLE blobs; +DELETE FROM store_meta WHERE key IN ('attachment_migration_phase','attachment_migration_cursor'); +"#; + const MIGRATIONS: &[Migration] = &[ Migration { version: 1, @@ -269,6 +303,14 @@ const MIGRATIONS: &[Migration] = &[ version: 2, sql: PROJECT_NAME_REVISION_DDL, }, + Migration { + version: 3, + sql: FILE_ATTACHMENTS_DDL, + }, + Migration { + version: 4, + sql: RETIRE_BLOBS_DDL, + }, ]; /// Check a declared schema using only reads before writer tuning or migrations. @@ -327,6 +369,9 @@ pub(super) fn migrate(connection: &mut Connection, now_ms: i64) -> StoreResult<( .iter() .filter(|migration| migration.version > version) { + if migration.version == 4 { + super::attachment_migration::retire_legacy_images(connection)?; + } let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; transaction.execute_batch(migration.sql)?; if migration.version == 1 { @@ -403,6 +448,42 @@ pub(super) fn mark_interrupted_receipts(connection: &mut Connection) -> StoreRes Ok(()) } +#[cfg(test)] +pub(super) fn migrate_to_version( + connection: &mut Connection, + now_ms: i64, + version: i64, +) -> StoreResult<()> { + if !(1..=VERSION).contains(&version) { + return Err(StoreError::InvalidInput); + } + connection.pragma_update(None, "application_id", APPLICATION_ID)?; + for migration in MIGRATIONS + .iter() + .filter(|migration| migration.version <= version) + { + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; + transaction.execute_batch(migration.sql)?; + if migration.version == 1 { + transaction.execute( + "INSERT INTO store_meta(key,value) VALUES ('store_id',?1)", + [rand::random::<[u8; 16]>().as_slice()], + )?; + transaction.execute( + "INSERT INTO store_meta(key,value) VALUES ('created_ms',?1)", + [now_ms], + )?; + transaction.execute( + "INSERT INTO store_meta(key,value) VALUES ('schema_epoch',?1)", + [SCHEMA_EPOCH], + )?; + } + transaction.pragma_update(None, "user_version", migration.version)?; + transaction.commit()?; + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -420,12 +501,12 @@ mod tests { } #[test] - fn migrations_validate_from_empty_and_keep_all_nineteen_tables() { + fn migrations_validate_from_empty_with_metadata_only_attachments() { let mut connection = empty(); migrate(&mut connection, 42).unwrap(); validate(&connection).unwrap(); let tables: i64 = connection.query_row("SELECT count(*) FROM sqlite_schema WHERE type='table' AND name NOT LIKE 'sqlite_%'", [], |r| r.get(0)).unwrap(); - assert_eq!(tables, 19); + assert_eq!(tables, 16); assert_eq!( connection .pragma_query_value::(None, "user_version", |r| r.get(0)) diff --git a/apps/maple-agent/crates/maple-agent/src/store/session.rs b/apps/maple-agent/crates/maple-agent/src/store/session.rs index 40df15abf..525e82efd 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/session.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/session.rs @@ -18,7 +18,7 @@ use pi_coding_agent::session::{ use pi_coding_agent::store::{MemoryStore, SessionStore}; use rusqlite::{OptionalExtension, Transaction, params}; -use super::blobs::PreparedEntryImages; +use super::attachments::PreparedEntryImages; use super::tasks::{TaskKey, task_key_tx, writable_account}; use super::{ AccountId, Durability, Store, StoreChange, StoreError, StoreResult, StoredRowChange, @@ -406,10 +406,8 @@ impl Store { if seq != next || flags & !1 != 0 { return Err(StoreError::InvalidData); } - let mut entry = decode_entry(encoded, &payload)?; - if flags & 1 != 0 { - self.rehydrate_entry_images(account, uuid, &mut entry)?; - } + let entry = decode_entry(encoded, &payload)?; + // Typed references stay on disk until the effective model context is selected. entries.push(entry); next += 1; } @@ -419,10 +417,13 @@ impl Store { } pub(super) fn decode_entry(encoded: i64, payload: &[u8]) -> StoreResult { - serde_json::from_slice(&codec::decode(encoded, payload)?).map_err(|_| StoreError::InvalidData) + let entry = serde_json::from_slice(&codec::decode(encoded, payload)?) + .map_err(|_| StoreError::InvalidData)?; + super::attachments::validate_stored_entry_images(&entry)?; + Ok(entry) } -/// Projector replay uses image markers plus blob-size metadata, never blob bytes. +/// Projector replay uses image references plus attachment metadata, never file bytes. pub(super) fn read_entries_from( tx: &Transaction<'_>, task: TaskKey, @@ -447,7 +448,7 @@ pub(super) fn read_entries_from( } /// A source entry for full row-field expansion. The projector restores image -/// summary lengths from blob metadata; timeline reads never fetch image bytes. +/// summary lengths from attachment metadata; timeline reads never fetch image bytes. pub(super) fn read_entry( tx: &Transaction<'_>, task: TaskKey, @@ -529,16 +530,29 @@ impl SqliteSessionStore { .prepare_entry_images(self.account, &self.uuid, entry)?; let json = serde_json::to_vec(&images.entry).map_err(|_| StoreError::InvalidInput)?; let (codec, payload) = codec::encode(&json)?; + let projection = self.prepare_projection(&images.entry, seq)?; Ok(( images, PreparedEntry { - projection: super::projector::prepare_entry(entry, seq)?, + projection, codec, payload, }, )) } + fn prepare_projection( + &self, + entry: &SessionEntry, + seq: i64, + ) -> StoreResult { + let sizes = self.store.read(|tx| { + self.guard(tx)?; + super::attachments::entry_image_sizes_tx(tx, self.task, entry) + })?; + super::projector::prepare_stored_entry(entry, seq, &sizes) + } + fn guard(&self, tx: &Transaction<'_>) -> StoreResult<()> { writable_account(tx, self.account)?; if task_key_tx(tx, self.account, &self.uuid)? != self.task { @@ -792,10 +806,7 @@ impl SqliteSessionStore { serde_json::to_vec(&images.entry).map_err(|_| StoreError::InvalidInput)?; let (codec, payload) = codec::encode(&raw)?; Ok(PreparedEntry { - projection: super::projector::prepare_entry( - &entries[offset], - start + offset as i64, - )?, + projection: self.prepare_projection(&images.entry, start + offset as i64)?, codec, payload, }) diff --git a/apps/maple-agent/crates/maple-agent/src/store/session/tests.rs b/apps/maple-agent/crates/maple-agent/src/store/session/tests.rs index dd702d6fd..b6ed99dbd 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/session/tests.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/session/tests.rs @@ -288,8 +288,21 @@ fn cache_parity(store: &Store, account: AccountId, manager: &SessionManager) { #[test] fn all_entry_kinds_and_message_variants_round_trip_exactly() { let (_directory, store, account) = fixture(); + let stored = store + .store_image_uploads( + account, + UUID, + &[crate::store::ImageUpload { + name: "fixture.png".into(), + data_url: format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(b"\x89PNG\r\n\x1a\nfixture") + ), + }], + ) + .unwrap(); let image = ImageContent { - data: base64::engine::general_purpose::STANDARD.encode(b"one copy of image bytes"), + data: stored[0].source.clone(), mime_type: "image/png".into(), }; let mut first_user = UserMessage { @@ -420,20 +433,22 @@ fn all_entry_kinds_and_message_variants_round_trip_exactly() { serde_json::to_vec(&source).unwrap() ); assert!(adapter.contains_entry("00000000").unwrap()); - let (compressed, blobs): (i64, i64) = store + let (compressed, files): (i64, i64) = store .read(|tx| { Ok(( tx.query_row("SELECT count(*) FROM entries WHERE codec=1", [], |r| { r.get(0) })?, - tx.query_row("SELECT count(*) FROM blobs WHERE complete=1", [], |r| { - r.get(0) - })?, + tx.query_row( + "SELECT count(*) FROM file_attachments WHERE state='ready' AND bound=1", + [], + |r| r.get(0), + )?, )) }) .unwrap(); assert!(compressed > 0); - assert_eq!(blobs, 1); + assert_eq!(files, 1); } #[test] @@ -449,7 +464,7 @@ fn third_party_images_survive_failed_append_retry_and_cold_plain_or_keyed_reopen let mut manager = loaded.into_manager("/projects/test"); manager.append_message(user("saved before a third-party tool result")); receipts.drain(); - let canonical = STANDARD.encode([0xfb, 0xff]); + let canonical = STANDARD.encode(b"\x89PNG\r\n\x1a\n\xfb\xff"); let mut content = vec![ Content::image("not base64!", "image/png"), Content::image("", "image/png"), @@ -464,7 +479,7 @@ fn third_party_images_survive_failed_append_retry_and_cold_plain_or_keyed_reopen ]; content.extend([ Content::image(canonical.clone(), "image/png"), - Content::image(canonical, "image/png"), + Content::image(canonical.clone(), "image/png"), ]); // Keep the exact formerly poisoning entry in Pi's retry slice while an // unrelated transaction failure forces write_all instead of append. @@ -504,15 +519,97 @@ fn third_party_images_survive_failed_append_retry_and_cold_plain_or_keyed_reopen ); assert_eq!(manager.entries(), retained); cache_parity(&store, account, &manager); - let (blobs, links): (i64, i64) = store + let persisted = store + .load_session(account, UUID, SessionLoad::Full) + .unwrap() + .entries; + let original = retained.iter().find(|entry| entry.id == image_id).unwrap(); + let saved = persisted.iter().find(|entry| entry.id == image_id).unwrap(); + assert_eq!( + (&saved.id, &saved.parent_id, saved.timestamp), + (&original.id, &original.parent_id, original.timestamp) + ); + let EntryKind::Message { + message: SessionMessage::Llm(Message::ToolResult(result)), + } = &saved.kind + else { + panic!("tool result must retain its entry kind"); + }; + assert!(!result.is_error); + assert_eq!(result.tool_call_id, "third-party-image"); + assert_eq!(result.tool_name, "mcp__fixture__screenshot"); + assert_eq!( + result.details, + Some(json!({"data":"maple-inline-image:opaque-details"})) + ); + assert_eq!(result.content.len(), 12); + for block in &result.content { + let Content::Image(image) = block else { + panic!("typed image marker must remain typed"); + }; + assert!( + image.data.starts_with("maple-attachment://") + || image.data == "maple-attachment-unavailable:retired" + ); + assert!( + !image.data.contains(&canonical), + "source rows must not retain image bytes" + ); + } + for index in [0, 1, 6, 7, 8, 9] { + let Content::Image(image) = &result.content[index] else { + unreachable!() + }; + assert_eq!(image.data, "maple-attachment-unavailable:retired"); + } + for index in [2, 3, 10, 11] { + let Content::Image(image) = &result.content[index] else { + unreachable!() + }; + assert!( + image.data.starts_with("maple-attachment://"), + "valid wrapped/native media must be retained" + ); + } + let mut model_messages = vec![Message::ToolResult(result.clone())]; + store + .materialize_message_images(account, UUID, &mut model_messages) + .unwrap(); + let Message::ToolResult(materialized) = &model_messages[0] else { + unreachable!() + }; + for index in [2, 3, 10, 11] { + let Content::Image(image) = &materialized.content[index] else { + panic!("valid image lost"); + }; + assert_eq!(image.data, canonical); + } + for index in [0, 1, 6, 7, 8, 9] { + assert!( + matches!(&materialized.content[index], Content::Text(text) if text.text.contains("no longer available")) + ); + } + let (files, bound): (i64, i64) = store .read(|tx| { Ok(( - tx.query_row("SELECT count(*) FROM blobs", [], |row| row.get(0))?, - tx.query_row("SELECT count(*) FROM task_blobs", [], |row| row.get(0))?, + tx.query_row( + "SELECT count(*) FROM file_attachments WHERE state='ready'", + [], + |row| row.get(0), + )?, + tx.query_row( + "SELECT count(*) FROM file_attachments WHERE state='ready' AND bound=1", + [], + |row| row.get(0), + )?, )) }) .unwrap(); - assert_eq!((blobs, links), (1, 1)); + assert!(files >= 4); + assert_eq!( + files, bound, + "retry must not leave duplicate unbound copies" + ); let page_request = crate::store::paging::OpenTaskPageRequest { task: UUID.into(), limit_rows: None, @@ -542,7 +639,7 @@ fn third_party_images_survive_failed_append_retry_and_cold_plain_or_keyed_reopen assert_eq!( items(&store), before_rebuild, - "metadata rebuild preserves raw image lengths" + "metadata rebuild preserves canonical image lengths" ); drop(manager); store.close_for_shutdown().unwrap(); @@ -551,7 +648,7 @@ fn third_party_images_survive_failed_append_retry_and_cold_plain_or_keyed_reopen let cold = reopened .load_session(account, UUID, SessionLoad::Full) .unwrap(); - assert_eq!(cold.entries, retained); + assert_eq!(cold.entries, persisted); let mut manager = cold.into_manager("/projects/test"); manager.append_message(user("continues after reopening")); assert!(manager.take_persist_error().is_none()); diff --git a/apps/maple-agent/crates/maple-agent/src/store/tasks.rs b/apps/maple-agent/crates/maple-agent/src/store/tasks.rs index dbc8c4cdd..957490b47 100644 --- a/apps/maple-agent/crates/maple-agent/src/store/tasks.rs +++ b/apps/maple-agent/crates/maple-agent/src/store/tasks.rs @@ -511,7 +511,7 @@ impl Store { AND deleted_ms IS NULL AND provisional=1 AND title_user_set=0 AND message_count=0 AND last_seq=0 AND NOT EXISTS (SELECT 1 FROM entries e WHERE e.task_id=tasks.id) - AND NOT EXISTS (SELECT 1 FROM task_blobs b WHERE b.task_id=tasks.id)", + AND NOT EXISTS (SELECT 1 FROM file_attachments a WHERE a.task_id=tasks.id)", params![account.0, uuid, now_ms], )? != 0; Ok((changed, changed.then(|| task_change(account, uuid)))) @@ -1295,7 +1295,7 @@ impl Store { } /// Remove only untouched provisional tasks. Such tasks contain no source - /// entries or blobs; deleting them does not become a payload-sized cascade. + /// entries or attachment uploads; deleting them does not orphan staged files. pub fn sweep_provisional_tasks( &self, account: AccountId, @@ -1309,7 +1309,7 @@ impl Store { WHERE t.account_id=?1 AND t.provisional=1 AND t.title_user_set=0 AND t.last_seq=0 AND t.message_count=0 AND t.updated_ms=40000)") if obs.get("attachment_bytes") != 10 * MIB or not integer(obs.get("oversized_row_total_bytes"), 20 * MIB): raise ValueError("attachment/oversized-row fixture size differs") + if obs.get("attachment_storage") != "plaintext session-owned file; database metadata only": + raise ValueError("attachment storage boundary differs") for name in ("model_context_full_cold", "model_context_boundary_cold"): counts = obs.get(f"{name}_entries", []) if len(counts) != samples[name]["samples"] or not all(integer(count, 1) for count in counts): @@ -224,12 +224,21 @@ def validate_observations(run, samples, full): stages.add(chunk.get("stage")) if not all(integer(chunk.get(key)) for key in ("rows", "bytes", "elapsed_micros")): raise ValueError("collection chunk measurements missing") + if chunk["stage"] == "Attachments": + # One owned attachment subtree: at most 128 filesystem mutations + # plus its metadata row. Bytes count the declared payload only on + # final removal, not adjacent model outputs or rewritten SQL bytes. + if not 1 <= chunk["rows"] <= 129 or chunk["bytes"] > 30 * MIB or chunk.get("oversized_row") is not False: + raise ValueError("attachment cleanup exceeds its bounded subtree step") + continue oversized = chunk["bytes"] > MIB if (type(chunk.get("oversized_row")) is not bool or chunk["oversized_row"] != oversized or chunk["rows"] > 1000 or (oversized and (chunk["rows"] != 1 or chunk["stage"] not in {"Timeline", "Entries", "Summaries"}))): raise ValueError("collection chunk exceeds caps without an isolated row") - if not {"Timeline", "Entries", "Blobs", "TaskMetadata"} <= stages or not any(chunk["oversized_row"] for chunk in collection): - raise ValueError("collection lost its oversized source-row/blob fixture") + if not {"Timeline", "Entries", "Attachments", "TaskMetadata"} <= stages or not any(chunk["oversized_row"] for chunk in collection): + raise ValueError("collection lost its oversized source-row/file fixture") + if sum(chunk["bytes"] for chunk in collection if chunk["stage"] == "Attachments") != 10 * MIB: + raise ValueError("collection lost its complete 10 MiB file fixture") summary = samples["oversized_task_collection_chunk"] if summary["samples"] != len(collection) or summary["maximum"] < max(chunk["elapsed_micros"] for chunk in collection): raise ValueError("collection summary lost a measured chunk") @@ -258,7 +267,7 @@ def check(mode, name, actual, budget, scope, statistic="median", oversized=False ("open_100_cold", 2000, "Store page and plan lookup; UI/live overlay unmeasured"), ("raw_connection_open", 1000, "one raw-key SQLite connection; excludes host startup/password KDF"), ("model_context_boundary_cold", 50000, "first partial context load including parse; OS cache warm"), - ("attachment_10mib_first_read", 25000, "first SQLite-reader range assembly; OS cache warm, no network")): + ("attachment_10mib_first_read", 25000, "first plaintext-file range assembly with keyed metadata; OS cache warm, no network")): check(mode, name, samples[name]["median"], budget, scope) if report["profile"] == "full": check(mode, "history_200_deep_cold", samples["history_200_deep_cold"]["median"], 3000, @@ -270,11 +279,11 @@ def check(mode, name, actual, budget, scope, statistic="median", oversized=False ("timeline_rebuild_writer_chunk", obs["timeline_rebuild_chunks"], "writer_micros"), ("old_epoch_collection_chunk", obs["old_epoch_collection_chunks"], "writer_micros"), ("oversized_task_collection_chunk", obs["oversized_collection_chunks"], "elapsed_micros")): - capped = [chunk[field] for chunk in chunks if chunk["bytes"] <= MIB] - oversized = [chunk[field] for chunk in chunks if chunk["bytes"] > MIB] + capped = [chunk[field] for chunk in chunks if chunk["bytes"] <= MIB or chunk.get("stage") == "Attachments"] + oversized = [chunk[field] for chunk in chunks if chunk["bytes"] > MIB and chunk.get("stage") != "Attachments"] if capped: check(mode, f"{name}/capped", max(capped), 50000, - "maximum capped chunk Store call; includes admission/commit, bounds writer hold", "maximum") + "maximum capped chunk or bounded attachment-subtree cleanup call; includes admission/commit, bounds writer hold", "maximum") if oversized: check(mode, f"{name}/isolated_oversized_row", max(oversized), 50000, "isolated oversized row; miss requires chunk storage or an approved measured alternative", diff --git a/apps/maple-agent/crates/store-bench/src/run.rs b/apps/maple-agent/crates/store-bench/src/run.rs index e16071b84..cbdb5487c 100644 --- a/apps/maple-agent/crates/store-bench/src/run.rs +++ b/apps/maple-agent/crates/store-bench/src/run.rs @@ -4,18 +4,19 @@ use crate::{ stats::Samples, }; use anyhow::{Context, Result, ensure}; -use base64::{Engine as _, engine::general_purpose::STANDARD}; use maple_agent::store::bench::{EffectivePragmas, FixtureOwner, PlanCheck, SeedEntry}; use maple_agent::store::keys::{KdfParams, KeyFile, LockedKey, generate_passphrase}; use maple_agent::store::{ - AccountId, HistoryPageRequest, ImageUpload, NavigationGroupQuery, NavigationSearchQuery, - NavigationTaskQuery, NewTask, OpenTaskPageRequest, PageTransport, PinFilter, RowRangeRequest, - RowReader, Store, StorePaths, TaskKind, TaskPatch, TaskSection, TaskState, + AccountId, AttachmentUploadMetadata, AttachmentUploadStart, HistoryPageRequest, + NavigationGroupQuery, NavigationSearchQuery, NavigationTaskQuery, NewTask, OpenTaskPageRequest, + PageTransport, PinFilter, RowRangeRequest, RowReader, Store, StorePaths, StoredAttachment, + TaskKind, TaskPatch, TaskSection, TaskState, }; use pi_coding_agent::session::{SESSION_VERSION, SessionHeader}; use pi_coding_agent::store::{JsonlStore, SessionStore}; use serde::Serialize; use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; use std::{ collections::BTreeMap, fs, @@ -433,6 +434,33 @@ fn seed( ) } +fn upload_fixture( + store: &Store, + account: AccountId, + task: &str, + operation: &str, + name: &str, + bytes: &[u8], +) -> Result { + let metadata = AttachmentUploadMetadata { + name: name.into(), + mime_type: "application/octet-stream".into(), + size: bytes.len() as u64, + sha256: format!("{:x}", Sha256::digest(bytes)), + }; + let attachment = match store.begin_attachment_upload(account, task, operation, &metadata)? { + AttachmentUploadStart::Complete(attachment) => attachment, + AttachmentUploadStart::Writing(mut writer) => { + for chunk in bytes.chunks(1024 * 1024) { + writer.write_chunk(chunk)?; + } + writer.finish(store)? + } + }; + store.pin_attachments(account, task, std::slice::from_ref(&attachment.id))?; + Ok(attachment) +} + fn attachment_scenarios( store: &Arc, account: AccountId, @@ -440,15 +468,14 @@ fn attachment_scenarios( scenarios: &mut Vec, observations: &mut BTreeMap, ) -> Result<()> { - let mut bytes = vec![0u8; 10 * 1024 * 1024]; - bytes[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); - let attachments = store.store_image_uploads( + let bytes = vec![0u8; 10 * 1024 * 1024]; + let attachment = upload_fixture( + store, account, &task.uuid, - &[ImageUpload { - name: "synthetic.png".into(), - data_url: format!("data:image/png;base64,{}", STANDARD.encode(&bytes)), - }], + "benchmark-read-file", + "synthetic.bin", + &bytes, )?; let read = || -> Result { let mut offset = 0; @@ -456,7 +483,7 @@ fn attachment_scenarios( let range = store.read_attachment( account, &task.uuid, - &attachments[0].id, + &attachment.id, offset as u64, 1024 * 1024, )?; @@ -473,6 +500,10 @@ fn attachment_scenarios( )); scenarios.push(measure("attachment_10mib_warm", 5, || read().map(|_| ()))?); observations.insert("attachment_bytes".into(), json!(bytes.len())); + observations.insert( + "attachment_storage".into(), + json!("plaintext session-owned file; database metadata only"), + ); Ok(()) } @@ -502,18 +533,20 @@ fn oversized_scenarios( created_ms: task.created_ms, }, )?; - // A distinct chunked blob makes this deletion exercise blob payload work, - // rather than only removing a shared attachment link from the hot task. - let mut blob = vec![0x55; 10 * 1024 * 1024]; - blob[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); - let upload = ImageUpload { - name: "synthetic-deleted.png".into(), - data_url: format!("data:image/png;base64,{}", STANDARD.encode(blob)), - }; + // Each task owns a separate file. Deletion measures file unlink and durable + // metadata removal, not a SQL payload rewrite or a shared blob reference. + let bytes = vec![0x55; 10 * 1024 * 1024]; let started = Instant::now(); - store.store_image_uploads(account, &task.uuid, &[upload])?; + upload_fixture( + store, + account, + &task.uuid, + "benchmark-deleted-file", + "synthetic-deleted.bin", + &bytes, + )?; scenarios.push(Samples::micros( - "durable_chunked_upload_10mib", + "durable_file_upload_10mib", vec![micros(started)], )); let entry = data::oversized_entry(&task, 0, 20 * 1024 * 1024); diff --git a/apps/maple-agent/crates/store-bench/tests/test_evaluate.py b/apps/maple-agent/crates/store-bench/tests/test_evaluate.py index b61d41e30..4da91effd 100644 --- a/apps/maple-agent/crates/store-bench/tests/test_evaluate.py +++ b/apps/maple-agent/crates/store-bench/tests/test_evaluate.py @@ -41,10 +41,10 @@ def report(full=False): rebuild = [{"prepare_micros": 100, "writer_micros": 100, "entries": 100, "bytes": 1000, "done": False, **wal()}, {"prepare_micros": 100, "writer_micros": 100, "entries": 0, "bytes": 0, "done": True, **wal()}] old = [{"writer_micros": 100, "rows": 100, "bytes": 1000, **wal()}] - collection = [{"stage": stage, "rows": 1, "bytes": size, "oversized_row": size > MODULE.MIB, + collection = [{"stage": stage, "rows": 3 if stage == "Attachments" else 1, "bytes": size, "oversized_row": size > MODULE.MIB and stage != "Attachments", "elapsed_micros": 100, "done": stage == "TaskMetadata", **wal()} for stage, size in (("Timeline", 20 * MODULE.MIB), ("Entries", 20 * MODULE.MIB), - ("Blobs", MODULE.MIB), ("TaskMetadata", 0))] + ("Attachments", 10 * MODULE.MIB), ("TaskMetadata", 0))] scenarios["timeline_rebuild_writer_chunk"] = sample("timeline_rebuild_writer_chunk", [100, 100]) scenarios["timeline_rebuild_prepare_chunk"] = sample("timeline_rebuild_prepare_chunk", [100, 100]) scenarios["oversized_task_collection_chunk"] = sample("oversized_task_collection_chunk", [100] * 4) @@ -54,7 +54,7 @@ def report(full=False): for role in ("writer", "reader-1", "reader-2", "bulk", "checkpointer")], "plans": [{"name": name, "statistics": stats, "expected_index": index, "passed": True, "details": [ - f"SEARCH attachments USING COVERING INDEX {index} (task_id=? AND blob_id=?)" + f"SEARCH file_attachments USING INDEX {index} (task_id=?)" if name == "attachment-collection" else f"SEARCH USING INDEX {index}"]} for name, index in MODULE.PLAN_INDEXES.items() for stats in (False, True)], "seed": {"layout": "globally timestamp-interleaved source and projector rows", @@ -63,6 +63,7 @@ def report(full=False): "durability": "FULL with fullfsync; synthetic fixture batches outside timed production source-append measurements"}, "scenarios": list(scenarios.values()), "observations": {"history_depth_rows": 40000 if full else 1800, "attachment_bytes": 10 * MODULE.MIB, + "attachment_storage": "plaintext session-owned file; database metadata only", "oversized_row_total_bytes": 20 * MODULE.MIB, "model_context_full_cold_entries": [scale["long_entries"]] * 5, "model_context_boundary_cold_entries": [100] * 5, @@ -267,29 +268,18 @@ def test_caps_and_observed_maxima_are_validated(self): with self.assertRaisesRegex(ValueError, "chunk quantile|chunk maximum"): MODULE.validate(actual) - def test_attachment_collection_accepts_both_covering_two_equality_indexes(self): - for index in ("attachments_task_blob", "attachments_blob"): - for constraints in ("task_id=? AND blob_id=?", "blob_id=? AND task_id=?"): - with self.subTest(index=index, constraints=constraints): - actual = report() - for run in actual["runs"]: - for plan in run["plans"]: - if plan["name"] == "attachment-collection": - plan["details"] = [f"SEARCH attachments USING COVERING INDEX {index} ({constraints})"] - MODULE.validate(actual) + def test_attachment_collection_seeks_task_metadata(self): + MODULE.validate(report()) - def test_attachment_collection_rejects_wrong_incomplete_or_noncovering_plans(self): + def test_attachment_collection_rejects_wrong_or_unbounded_plans(self): bad_details = ( - ["SEARCH attachments USING COVERING INDEX attachments_blob (blob_id=?)"], - ["SEARCH attachments USING COVERING INDEX attachments_task_blob (task_id=?)"], - ["SEARCH attachments USING COVERING INDEX attachments_blob (blob_id=? AND task_id>?)"], - ["SEARCH attachments USING INDEX attachments_blob (blob_id=? AND task_id=?)"], - ["SEARCH attachments USING COVERING INDEX unrelated (blob_id=? AND task_id=?)"], - ["SEARCH attachments USING COVERING INDEX attachments_blob_extra (blob_id=? AND task_id=?)"], - ["SCAN attachments USING COVERING INDEX attachments_blob (blob_id=? AND task_id=?)"], - ["SEARCH other_table USING COVERING INDEX attachments_blob (blob_id=? AND task_id=?)"], - ["SEARCH attachments USING COVERING INDEX attachments_blob (blob_id=? AND task_id=?)", "SCAN attachments"], - ["SEARCH attachments USING COVERING INDEX attachments_blob (blob_id=? AND task_id=?)", "USE TEMP B-TREE FOR ORDER BY"], + ["SEARCH file_attachments USING INDEX file_attachments_task (task_id>?)"], + ["SEARCH file_attachments USING INDEX unrelated (task_id=?)"], + ["SEARCH file_attachments USING INDEX file_attachments_task_extra (task_id=?)"], + ["SCAN file_attachments USING INDEX file_attachments_task"], + ["SEARCH other_table USING INDEX file_attachments_task (task_id=?)"], + ["SEARCH file_attachments USING INDEX file_attachments_task (task_id=?)", "SCAN file_attachments"], + ["SEARCH file_attachments USING INDEX file_attachments_task (task_id=?)", "USE TEMP B-TREE FOR ORDER BY"], ) for details in bad_details: with self.subTest(details=details): @@ -299,16 +289,57 @@ def test_attachment_collection_rejects_wrong_incomplete_or_noncovering_plans(sel with self.assertRaisesRegex(ValueError, "required index"): MODULE.validate(actual) - def test_attachment_collection_keeps_primary_expected_field_and_actual_pass_guard(self): - for changes in ({"expected_index": "attachments_blob"}, {"passed": False}): + def test_attachment_collection_keeps_expected_field_and_actual_pass_guard(self): + for changes in ({"expected_index": "wrong_index"}, {"passed": False}): with self.subTest(changes=changes): actual = report() plan = next(plan for plan in actual["runs"][1]["plans"] if plan["name"] == "attachment-collection") - plan["details"] = ["SEARCH attachments USING COVERING INDEX attachments_blob (blob_id=? AND task_id=?)"] plan.update(changes) with self.assertRaises(ValueError): MODULE.validate(actual) + def test_file_collection_is_bounded_without_an_oversized_row_exception(self): + for rows in (1, 3, 128, 129): + actual = report() + chunk = next(chunk for chunk in actual["runs"][0]["observations"]["oversized_collection_chunks"] + if chunk["stage"] == "Attachments") + chunk["rows"] = rows + MODULE.validate(actual) + for changes in ({"rows": 0}, {"rows": 130}, {"bytes": 30 * MODULE.MIB + 1}, {"oversized_row": True}): + with self.subTest(changes=changes): + actual = report() + chunk = next(chunk for chunk in actual["runs"][0]["observations"]["oversized_collection_chunks"] + if chunk["stage"] == "Attachments") + chunk.update(changes) + with self.assertRaisesRegex(ValueError, "bounded subtree step"): + MODULE.validate(actual) + actual = report() + chunk = next(chunk for chunk in actual["runs"][0]["observations"]["oversized_collection_chunks"] + if chunk["stage"] == "Attachments") + chunk["bytes"] = MODULE.MIB + with self.assertRaisesRegex(ValueError, "complete 10 MiB file"): + MODULE.validate(actual) + + def test_file_cleanup_keeps_the_normal_latency_gate(self): + actual = report(True) + chunk = next(chunk for chunk in actual["runs"][0]["observations"]["oversized_collection_chunks"] + if chunk["stage"] == "Attachments") + chunk["rows"] = 129 + chunk["elapsed_micros"] = 50001 + set_sample(actual, "plain", "oversized_task_collection_chunk", [100, 100, 100, 50001]) + MODULE.validate(actual) + outcome = MODULE.absolute_results(actual) + result = next(item for item in outcome["checks"] + if item["scenario"] == "plain/oversized_task_collection_chunk/capped") + self.assertEqual(result["status"], "failed") + self.assertEqual(result["actual_us"], 50001) + + def test_old_blob_storage_evidence_cannot_pass_as_file_storage(self): + actual = report() + del actual["runs"][0]["observations"]["attachment_storage"] + with self.assertRaisesRegex(ValueError, "storage boundary"): + MODULE.validate(actual) + def test_signed_wal_growth_is_observation_not_bytes_written(self): actual = report() actual["runs"][0]["observations"]["timeline_rebuild_chunks"][0].update(wal(1000, 100)) diff --git a/apps/maple-agent/docs/local-host.md b/apps/maple-agent/docs/local-host.md index 4dac11b7b..a9f294d43 100644 --- a/apps/maple-agent/docs/local-host.md +++ b/apps/maple-agent/docs/local-host.md @@ -132,8 +132,10 @@ database record admission and its result, so a lost reply can be queried without resending the prompt. A crash during admission produces an explicit unknown outcome: inspect the task before retrying. Receipts contain hashes and task/run IDs, not prompt text or -images. They are not a token event log. IPC frames are bounded at 160 MiB to fit -the existing attachment allowance; oversized requests fail explicitly. +images. They are not a token event log. IPC frames retain a 160 MiB ceiling for +bounded history and control replies; oversized messages fail explicitly. +Attachment imports use staged paths and bounded byte streams, not inline IPC +payloads. History pages can shorten an oversized display row without changing its saved source. Desktop **Show full output** reconstructs a row only within its 32 MiB diff --git a/apps/maple-agent/docs/local-store-encryption.md b/apps/maple-agent/docs/local-store-encryption.md index 7563f3e84..2153b7463 100644 --- a/apps/maple-agent/docs/local-store-encryption.md +++ b/apps/maple-agent/docs/local-store-encryption.md @@ -2,8 +2,12 @@ Agent uses a fresh `store-v1` namespace. Previous Agent/Research data is neither read, adopted, migrated nor deleted by this implementation. Plain mode remains -the default. Password mode encrypts the SQLite database, WAL, attachments and -other authoritative records with a random local DEK wrapped by the password. +the default. Password mode encrypts the SQLite database, WAL, attachment metadata +and other database records with a random local DEK wrapped by the password. +Attachment payloads, including images and tool screenshots, are ordinary +session-owned files outside that encryption boundary. They remain plaintext +even while the database is locked. See [local storage](local-store.md) for +their lifecycle, migration and backup rules. ## Release exposure diff --git a/apps/maple-agent/docs/local-store.md b/apps/maple-agent/docs/local-store.md index 4fbdc0794..a30d1c63b 100644 --- a/apps/maple-agent/docs/local-store.md +++ b/apps/maple-agent/docs/local-store.md @@ -48,6 +48,7 @@ store-v1/ keys.json password-wrapped data key, password mode only keys.backup.json second key-file copy, password mode only logs/maple-agent.log plaintext operational diagnostics + attachments/ plaintext session-owned files and images resources/ generated instruction Markdown, skills and prompts setup.json only while an erase/setup operation is unfinished db.tmp/ db.new/ db.old/ temporary generation replacement/recovery only @@ -58,11 +59,11 @@ ownership, permissions, links and nonregular database files. Windows creates a protected owner-only directory ACL. The authoritative lock is inside the store, so removal of a temporary IPC socket cannot create a second database owner. -New app state includes task metadata, complete source entries, derived timeline -rows, attachments/images, preferences and dictionary words, sign-in credentials, +New database state includes task metadata, complete source entries, derived timeline +rows, attachment metadata, preferences and dictionary words, sign-in credentials, MCP definitions/secrets, send receipts and remote identities/grants. Sign-out keeps local history. Account deletion removes only that account's new-store -records after server success is confirmed. If the outcome is ambiguous, Maple +records and owned attachment files after server success is confirmed. If the outcome is ambiguous, Maple keeps the data and pending deletion intent. Automatic reconciliation is not implemented in this version; ordinary sign-in may remain blocked by that intent. An authentication failure by itself does not prove that the server account was @@ -74,6 +75,41 @@ attachments, remote identities or temporary spill files. Existing installations therefore begin with fresh application state in this namespace. There is no migration or automatic cleanup of previous data. +Within `store-v1`, attachments use ordinary session-owned files. Attaching makes +a stable copy on the execution computer, including when the desktop is local. +Moving or editing the original does not change an accepted attachment. The +database keeps ownership and metadata; normal files and archives are given to +the model by their execution-side path. Archives are not automatically unpacked. +Images retain the model's existing vision and `read_image` behavior. + +Before sending, the desktop keeps a temporary plaintext copy under its profile's +`draft-attachments/` directory, outside `store-v1`. This lets a draft survive +changes to the selected original and lets an uncertain upload be retried with +the same operation identity. Live copies are released when the draft is removed, +sent, signed out or its view is torn down. A bounded sweep removes unlocked +crash leftovers after seven days. Store reset does not manage this desktop-owned +temporary directory, and database encryption does not protect these copies. + +Deleting a session schedules deletion of its owned files; interrupted cleanup +retries without starving other sessions. Ordinary read-only owned folders can be +reclaimed, but inaccessible permissions or foreign ownership still require repair; +Maple does not follow links or report incomplete erasure as success. +Compaction alone retains historical attachments. A stopped-store backup +must include both the database and the attachment tree. Unlock passwords do not +encrypt these files, and no password change converts them to encrypted payloads. + +The attachment schema migration retires the former SQLite image payloads instead +of exporting them as plaintext. Conversation text and tool results remain, with +explicit unavailable-image markers where needed. Existing captions are retained +when present; an image does not necessarily have a caption. This migration is +restartable and does not reset the database. It cannot be undone by launching an +older binary, and reclaimed SQLite pages may not immediately shrink the database +file. Existing backups and filesystem snapshots are not erased. +Opening or continuing an older task prioritizes rebuilding its derived history. +If that bounded preparation is still pending, Maple reports it instead of showing +an empty conversation. Missing historical image files become unavailable; new +outgoing images still require successful preparation. + General project/workspace Markdown, `AGENTS.md`/`CLAUDE.md`, `SYSTEM.md`, skills and prompt templates retain their resource-loading behavior. They remain files and are outside database encryption and application-data erasure. Newly @@ -123,6 +159,7 @@ local unlock password, separate from the Maple account password. It is not saved an OS keychain, passed in argv/environment variables, or sent to a remote peer. The password protects the current database while its host is stopped. It does +not encrypt attachment files, including images and tool screenshots. It does not protect against code running as the same user while the host is unlocked, root/administrator access, screenshots, clipboard contents, the project files tools edit, instruction files, or copies kept by external agents. File sizes, @@ -263,9 +300,9 @@ For a manual copy of this version's store: commands alone do not acquire this lock. 3. Copy the complete `store-v1/` directory to a separate destination. Include the complete `db/` generation, any WAL still present, both key files when - present, and `generation.json`. Never combine a database and key/generation + present, `generation.json`, and the complete `attachments/` tree. Never combine a database and key/generation files from different copies. Preserve private permissions and appropriate - ownership; instructions/resources and logs in that directory remain plaintext. + ownership; attachments, instructions/resources and logs remain plaintext. 4. Verify the copy independently while retaining the original. A file-transfer success is not an integrity or password-unlock check. Reopen the original only after the copy has finished. @@ -281,8 +318,8 @@ an ownership, generation or wrong-password refusal. Maple cannot reset a forgotten local unlock password. **Forgot password?** opens **Erase this version's local data**, requiring the exact text `ERASE`. This -removes this version's saved app-state database, including tasks/history, -attachments, settings, sign-in, MCP secrets and this store's local remote +removes this version's saved app-state database and owned attachment files, +including tasks/history, settings, sign-in, MCP secrets and this store's local remote pairings. It does not delete a server account, other computers' state, external agents' own copies, backups, snapshots, old app data, or the instruction-file exception above. Setup then requires an explicit new password or **Continue @@ -352,9 +389,11 @@ Append schema migrations; changes to item builders/derivation require a Generation changes revoke cached account/task-store handles; numeric account IDs cannot carry ownership into a replacement database. Credential writes compare the exact stored owner even when its row is hidden from the current API view. -Malformed or empty tool-result images remain lossless inline source when they -cannot be stored as blobs, so a blob conversion failure does not permanently -poison persistence retries. +Image payloads are never retained as inline database fallback data. Invalid or +unavailable media becomes an explicit unavailable marker without changing a tool +result's success status. Valid image references are resolved only when needed +for a preview or effective model context; ordinary history reads do not load +every image payload. Idle maintenance reclaims at most 64 freelist pages per chunk and periodically runs bounded `PRAGMA optimize`, yielding to foreground work. Idle full-integrity diff --git a/apps/maple-agent/docs/remote-access.md b/apps/maple-agent/docs/remote-access.md index 28ac67d48..3cf09aa25 100644 --- a/apps/maple-agent/docs/remote-access.md +++ b/apps/maple-agent/docs/remote-access.md @@ -77,9 +77,10 @@ when requested, and selecting a folder uses the existing project trust flow. The execution daemon resolves paths and performs filesystem operations. Folder lists are paged and bounded, and hidden folders can be shown explicitly. This is a folder picker, not a filesystem sandbox for an already paired controller. -Images picked -on the controller are explicit uploads. Remote frames are limited to 16 MiB, -including JSON/base64 overhead; reduce attachments if a send exceeds that bound. +Files picked on the controller are explicit uploads to the execution computer. +Remote control frames remain limited to 16 MiB, including encoding overhead; +file uploads use separate bounded raw-byte streams and do not count against a +send's control frame. Microphone, playback, billing and account operations stay on the controller. The Integrations pane manages the selected execution computer and labels it; account and application preference panes still use this computer. Native @@ -168,25 +169,48 @@ Maple explicitly clears public address lookup and builds Iroh without its port-mapping feature. The Minimal preset alone is not the port-mapping policy. There is no second encryption tunnel or custom routing layer. The remote protocol keeps the `maple-agent/remote/1` ALPN family and currently advertises -wire version 6 with minimum version 6, independent of application builds. Its +wire version 7 with minimum version 7, independent of application builds. Its shared project names, bounded navigation/history pages, stable row identities, attachment ranges, and per-connection task subscriptions remain unchanged. Project-change events carry only changed metadata; snapshots and project pages provide authoritative names, revisions, and default-workspace flags. -Version 5 and earlier peers receive an explicit version refusal. Bounded frames +Version 6 and earlier peers receive an explicit version refusal. Bounded frames and an explicit task-operation allowlist apply. Account operations and arbitrary whole-host administration cannot be forwarded. The narrow experimental update capability uses the host's own signed updater; it does not expose the ordinary host-stop command. -Optional discovery reports protocol version `6.2.0` and explicit capabilities -over separate request streams; the wire-6 handshake, acknowledgement shape and -ALPN remain unchanged. Bounded folder browsing and single-folder creation require +Optional discovery reports protocol version `7.0.0` and explicit capabilities +over separate request streams. The acknowledgement shape and ALPN remain +readable across the required version transition. Bounded folder browsing and single-folder creation require the `project_folders_v1` capability. The controller checks support before sending -either folder command. An older host without that capability keeps ordinary +either folder command. A compatible host without that capability keeps ordinary project/task access and its existing remote-update support; only the folder -browser reports that support is unavailable or an update is needed. Older -controllers can still use newer hosts through their existing commands. +browser reports that support is unavailable or an update is needed. + +Wire 7 requires attachment-reference sends. Images, ordinary files and archives +are copied into the execution daemon's session-owned plaintext file store; +SQLite contains metadata. Local imports use a controller-local staged path. +Remote imports send bounded metadata followed by raw bytes on a dedicated +low-priority stream, with admission before the body and a durable upload receipt +after finalization. The remote daemon never interprets a controller's local path. +The existing identified `send_message` then carries attachment IDs. An upload +retry cannot replay a prompt. Each authenticated connection has two bulk slots +separate from ordinary requests, range reads and urgent replies; generic files +are limited to 30 MiB each, while +composer images retain their 10 MiB limit. See [local storage](local-store.md). + +Both computers must use wire 7 for remote access, including remote update +management. Upgrade the execution computer first while the old pair is still +compatible, then update the controller locally. A wire-7 controller stops +automatic reconnection after an explicit version refusal and shows **Update +required**, without treating ordinary offline errors as version mismatches. +Released wire-6 controllers keep their shipped behavior during the gap: cold +computer discovery can retry, and Settings can report the peer as unreachable. +Existing pairing records are retained. A disconnect alone does not confirm an +installation; check the running identity and retained receipt after reconnecting. +If the controller was updated first, update the execution computer directly, +then select that computer again to reconnect. A bounded capability probe can fail without ending ordinary task requests. Responses tolerate future optional fields and capability names. Experimental @@ -203,11 +227,12 @@ separate from local IPC. Each control connection bounds ordinary calls and keeps separate capacity for cancellation/questions. Each request/reply uses an independent bidirectional QUIC stream; the authenticated initial stream carries ordered events. A bounded header admits each request before its attachment body -is read. Stop and question answers have reserved capacity and higher native stream +is read. Header parsing has an overall queue-and-read deadline. Stop and question +answers have reserved capacity and higher native stream priority, so an incomplete upload or unread image reply cannot hold their dispatch loop. Streams still share the connection's bandwidth; this is not a latency guarantee on a saturated or unavailable network. Transfers use bounded frames and -inactivity deadlines, allowing slow transfers that keep making progress. Image +inactivity deadlines. Uploads also have a ten-minute total deadline. Image readback uses base64 on the remote wire. A request failure or oversized reply returns a request error, while authority loss or incompatible framing ends the connection. A failed connection recovers from an @@ -216,7 +241,7 @@ the host or cause an automatic prompt replay; explicit host shutdown still stops its work. A stable send identity crosses both IPC and Iroh; only the execution daemon records its admission and result, account-bound in its application database. An uncertain send keeps its original -identity, text and draft images in the desktop view and queries the execution +identity, text and draft attachments in the desktop view and queries the execution host's receipt after reconnecting. It never replays the command or overwrites a newer draft. If the execution daemon restarts before recording a send's result, it marks that receipt interrupted, preserving the identity against replay. The diff --git a/apps/maple-agent/docs/updates.md b/apps/maple-agent/docs/updates.md index d59f8743b..7fc9f1bb5 100644 --- a/apps/maple-agent/docs/updates.md +++ b/apps/maple-agent/docs/updates.md @@ -39,9 +39,15 @@ An acknowledged request is consumed by the execution computer's running desktop and uses this same installer. Password-protected hosts are excluded even while unlocked; an independent host-side check runs at the final shutdown boundary. The controller's updater is never installed or restarted by this operation. -Daemon-only hosts require a local update. New capability discovery preserves the -existing wire-6 handshake and does not require older peers to update for ordinary -task access. +Daemon-only hosts require a local update. Remote management requires compatible +remote protocol majors, just like ordinary task access. The attachment change +requires wire 7 on both computers; it does not add a cross-version management +bridge. Update the execution computer first while both still use wire 6, then +update the controller through its local updater. During the expected mismatch, +retain the original install operation rather than resubmitting it; verify the +running identity and its receipt after both computers reconnect on wire 7. +Updating the controller first requires updating the execution computer directly. +A remote mismatch does not disable the controller's local update controls. The installer holds the per-application lock while staging a verified bundle beside the installed application and while that update is ready. It checks the