From adf14cc2a67b0275bf01722737515e5a8619e192 Mon Sep 17 00:00:00 2001 From: RaymondCho Date: Mon, 12 Oct 2026 23:01:34 +0900 Subject: [PATCH] Save and delete files in an aindrive drive without overwriting another change A token that may write opens a writable folder. A save sends the file whole through write_file with the revision it was read at, or "none" for create-only, so aindrive refuses it as a conflict when the file moved on; saves are capped at one 1 MiB transfer, which a table-driven encoder fits into one call's fuel. A save whose answer is lost, or that the drive's device timed out on, is reported with an unknown outcome. Deleting takes one file; folders are refused because delete_path is recursive and listings hide entries such as .git. Co-Authored-By: Claude Opus 5.5 --- README.md | 29 +++-- src/{decode.rs => codec.rs} | 91 ++++++++++++-- src/lib.rs | 182 ++++++++++++++++++++++----- src/mcp.rs | 41 ++++++ src/session_tests.rs | 245 +++++++++++++++++++++++++++++++++++- src/test_host.rs | 81 +++++++++++- 6 files changed, 615 insertions(+), 54 deletions(-) rename src/{decode.rs => codec.rs} (66%) diff --git a/README.md b/README.md index 757adc0..0a08322 100644 --- a/README.md +++ b/README.md @@ -27,15 +27,28 @@ that machine runs aindrive; until then they report the folder unavailable. ## What works -Browsing folders and reading files. aindrive reads a file whole, at most -8 MiB at a time, so a larger file is refused as unsupported rather than cut -short. A read is checked against the size the drive reported a moment before, -so a file that changes in between is a conflict rather than a mix of two -versions. Names the drive's machine allows but Memona cannot show, such as one -containing a backslash, are left out of listings. +Browsing folders, and reading, saving and deleting files. A token generated +as **Read only** opens the drive read-only. -Saving, creating folders, renaming and deleting are not supported yet, and -neither is change polling. +aindrive reads a file whole, at most 8 MiB at a time, so a larger file is +refused as unsupported rather than cut short. A read is checked against the +size the drive reported a moment before, so a file that changes in between is +a conflict rather than a mix of two versions. + +A save sends the file whole, up to 1 MiB, the most Memona saves in one +transfer. It names the revision it was read at, and aindrive compares and +writes under one lock, so a file someone changed since is reported as a +conflict instead of being overwritten. A save into a folder that does not +exist yet creates that folder, as aindrive does. + +Deleting takes one file. A folder is refused: aindrive deletes a folder with +everything in it, and leaves some entries, such as `.git`, out of every +listing, so a folder cannot be shown to be empty first. + +Creating folders, renaming and change polling are not supported, since the +drive's MCP endpoint offers no tool for them. Names the drive's machine allows +but Memona cannot show, such as one containing a backslash, are left out of +listings. ## Building diff --git a/src/decode.rs b/src/codec.rs similarity index 66% rename from src/decode.rs rename to src/codec.rs index 528b30f..6e8329d 100644 --- a/src/decode.rs +++ b/src/codec.rs @@ -1,8 +1,8 @@ -//! Base64 decoding cheap enough for Memona's per-call fuel. A read hands over -//! up to 1 MiB in one call, and a general decoder spends more instructions on -//! that than the call is given. Here every two characters cost one table -//! lookup, and the characters are read where they arrived, never copied: -//! the host charges fuel for each byte a plugin copies. +//! Base64 cheap enough for Memona's per-call fuel. A read or a save moves up +//! to 1 MiB in one call, and a general codec spends more instructions on that +//! than the call is given. Here every two characters cost one table lookup +//! either way, and decoded characters are read where they arrived, never +//! copied: the host charges fuel for each byte a plugin copies. use crate::sdk::ErrorCode; @@ -82,7 +82,7 @@ fn groups(chars: &[u8], out: &mut [u8]) -> u32 { /// Decodes standard base64 that arrived in several pieces, read as one run, /// into `out`. A group split between two pieces is joined on its own. The /// run holds whole groups and no padding; `last` decodes the final group. -pub fn run<'a>( +pub fn decode<'a>( pieces: impl Iterator, out: &mut [u8], ) -> Result<(), ErrorCode> { @@ -140,9 +140,72 @@ pub fn last(chars: [u8; 4]) -> Result, ErrorCode> { Ok(word.to_le_bytes()[..3 - padding].to_vec()) } +/// Two characters for each half of a group of three bytes, indexed straight +/// from the group read as a little-endian word, so no byte is moved first: +/// the first table by the first byte and the high nibble of the second, the +/// other by the second byte's low nibble and the third byte. +static HALVES: [[u16; 1 << 16]; 2] = halves(); + +const fn chars_of(bits: usize) -> u16 { + ALPHABET[bits >> 6] as u16 | (ALPHABET[bits & 63] as u16) << 8 +} + +const fn halves() -> [[u16; 1 << 16]; 2] { + let mut tables = [[0; 1 << 16]; 2]; + let mut index = 0; + while index < 1 << 16 { + tables[0][index] = chars_of((index & 0xff) << 4 | (index >> 12) & 0xf); + tables[1][index] = chars_of((index & 0xf) << 8 | (index >> 8) & 0xff); + index += 1; + } + tables +} + +/// The four characters of the three bytes at the bottom of `word`, read as a +/// little-endian word. +fn quad(word: u64) -> u32 { + u32::from(HALVES[0][(word & 0xf0ff) as usize]) + | u32::from(HALVES[1][((word >> 8) & 0xff0f) as usize]) << 16 +} + +/// How many characters `bytes` bytes encode to, padding included. +pub fn encoded_length(bytes: usize) -> usize { + bytes.div_ceil(3) * 4 +} + +/// Encodes `bytes` as standard, padded base64 into `out`, which must hold +/// exactly `encoded_length` characters. Twelve bytes go through each step +/// of the main loop, for the same reason sixteen characters do in `decode`. +pub fn encode(bytes: &[u8], out: &mut [u8]) { + let wide = bytes.chunks_exact(12); + let (wide_out, rest_out) = out.split_at_mut(bytes.len() / 12 * 16); + let rest = wide.remainder(); + for (group, chars) in wide.zip(wide_out.chunks_exact_mut(16)) { + let (Some(low), Some(high)) = + (group.first_chunk::<8>(), group.last_chunk::<8>()) + else { + continue; + }; + let low = u64::from_le_bytes(*low); + let high = u64::from_le_bytes(*high); + let front = u64::from(quad(low)) | u64::from(quad(low >> 24)) << 32; + let back = u64::from(quad(high >> 16)) | u64::from(quad(high >> 40)) << 32; + chars[..8].copy_from_slice(&front.to_le_bytes()); + chars[8..].copy_from_slice(&back.to_le_bytes()); + } + for (group, chars) in rest.chunks(3).zip(rest_out.chunks_exact_mut(4)) { + let mut padded = [0; 8]; + padded[..group.len()].copy_from_slice(group); + chars.copy_from_slice(&quad(u64::from_le_bytes(padded)).to_le_bytes()); + for char in &mut chars[group.len() + 1..] { + *char = b'='; + } + } +} + #[cfg(test)] mod tests { - use super::{last, run}; + use super::{decode as run, encode, encoded_length, last}; use crate::sdk::ErrorCode; use base64::Engine; use base64::engine::general_purpose::STANDARD; @@ -181,6 +244,20 @@ mod tests { Ok(()) } + #[tokio::test] + async fn every_length_encodes_as_a_general_encoder_would() { + let bytes: Vec = (0..=255).cycle().take(1000).collect(); + for length in 0..1000 { + let mut out = vec![0; encoded_length(length)]; + encode(&bytes[..length], &mut out); + assert_eq!( + out, + STANDARD.encode(&bytes[..length]).as_bytes(), + "{length}" + ); + } + } + #[tokio::test] async fn anything_but_standard_base64_is_refused() { for bad in ["QUJ-QUJD", "QUJ_QUJD", "QU\nJQUJD", "QU=DQUJD", "Q===QUJD"] { diff --git a/src/lib.rs b/src/lib.rs index 7ff4aa2..ca81307 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,7 +2,7 @@ //! remote MCP endpoint through the host's HTTP, with the drive's access token. mod address; -mod decode; +mod codec; mod mcp; mod sdk; #[cfg(test)] @@ -31,14 +31,28 @@ const MAX_CHUNK: u32 = 1024 * 1024; const MAX_READ: u64 = 8 * 1024 * 1024; /// A whole 8 MiB file in base64, with room for the reply around it. const MAX_REPLY: usize = 12 * 1024 * 1024; +/// The largest file a save sends. `write_file` takes a file whole, and its +/// base64 has to be encoded within one call's fuel, which a single transfer, +/// the most Memona saves a remote page in, leaves room for. +const MAX_WRITE: usize = MAX_CHUNK as usize; /// Without these the folder cannot even be browsed. const REQUIRED_TOOLS: [&str; 3] = ["list_files", "stat", "read_file"]; +/// A token offers these only when it may write. +const WRITE_TOOLS: [&str; 2] = ["write_file", "delete_path"]; struct Session { endpoint: String, authorization: String, } +/// Whether a request may change the drive. Once one has gone out whole, a +/// lost answer leaves the change unknown rather than failed. +#[derive(Clone, Copy, PartialEq)] +enum Effect { + Read, + Change, +} + impl Session { fn from_config(config: &[(String, Setting)]) -> Result { let lookup = |wanted: &str| { @@ -96,7 +110,7 @@ fn header(name: &str, value: &str) -> host::Header { /// the host handed over. They are not joined here: Memona's host charges /// fuel for every byte a plugin copies, and a file's reply joined whole would /// spend most of a call's budget before any of it was decoded. -fn post(body: &[u8]) -> Result>, ErrorCode> { +fn post(body: &[u8], effect: Effect) -> Result>, ErrorCode> { let (url, authorization) = SESSION.with_borrow(|session| { session .as_ref() @@ -116,12 +130,27 @@ fn post(body: &[u8]) -> Result>, ErrorCode> { for piece in body.chunks(MAX_CHUNK as usize) { call.push(piece).map_err(transport)?; } - let head = call.respond().map_err(transport)?; - mcp::status(head.status)?; + let unsettled = |error: ErrorCode| match (effect, error) { + (Effect::Change, ErrorCode::Unavailable | ErrorCode::Cancelled) => { + ErrorCode::OutcomeUnknown + } + _ => error, + }; + let head = call + .respond() + .map_err(|error| unsettled(transport(error)))?; + // A server error can come after the drive's server already acted, while a + // refused token, permission, drive or rate comes before it does anything. + mcp::status(head.status).map_err(|error| match head.status { + 500.. => unsettled(error), + _ => error, + })?; let mut pieces = Vec::new(); let mut total = 0; loop { - let piece = call.pull(MAX_CHUNK).map_err(transport)?; + let piece = call + .pull(MAX_CHUNK) + .map_err(|error| unsettled(transport(error)))?; if piece.is_empty() { return Ok(pieces); } @@ -141,12 +170,22 @@ fn joined(mut pieces: Vec>) -> Vec { } } -/// Calls one of the drive's tools and reads its structured answer. +/// Calls one of the drive's tools that only reads, and reads its structured +/// answer. A call that timed out on the drive's device changed nothing. fn tool( name: &str, arguments: &serde_json::Value, ) -> Result { - mcp::reply(&joined(post(&mcp::call(name, arguments))?)) + let pieces = post(&mcp::call(name, arguments), Effect::Read)?; + mcp::reply(&joined(pieces)).map_err(|error| match error { + ErrorCode::OutcomeUnknown => ErrorCode::Unavailable, + error => error, + }) +} + +/// Sends a call that changes the drive and reads its structured answer. +fn change(body: &[u8]) -> Result { + mcp::reply(&joined(post(body, Effect::Change)?)) } fn info_of(entry: &mcp::Entry) -> FileInfo { @@ -188,7 +227,7 @@ impl Download { chars[filled..filled + slice.len()].copy_from_slice(slice); filled += slice.len(); } - let bytes = decode::last(chars)?; + let bytes = codec::last(chars)?; let final_group = self.encoded.len() / 4 - 1; if filled != 4 || (index != final_group && bytes.len() != 3) { return Err(ErrorCode::Unavailable); @@ -214,7 +253,7 @@ impl Download { let mut bytes = Vec::with_capacity((end - first) * 3); bytes.resize((whole - first) * 3, 0); let chars = self.chars(first)..self.chars(whole); - decode::run(mcp::slices(&self.reply, chars), &mut bytes)?; + codec::decode(mcp::slices(&self.reply, chars), &mut bytes)?; if end > final_group { bytes.extend(self.group(final_group)?); } @@ -258,10 +297,10 @@ fn download(target: &Locator, offset: u64) -> Result { } else { let length = usize::try_from(entry.size.div_ceil(3) * 4) .map_err(|_| ErrorCode::Unsupported)?; - let reply = post(&mcp::call( - "read_file", - &json!({ "path": path, "encoding": "base64" }), - ))?; + let reply = post( + &mcp::call("read_file", &json!({ "path": path, "encoding": "base64" })), + Effect::Read, + )?; let encoded = mcp::contents(&reply, length)?; (reply, encoded) }; @@ -274,19 +313,80 @@ fn download(target: &Locator, offset: u64) -> Result { }) } -/// Writing is not offered yet, so no transfer is ever started. -enum Upload {} +/// One file being saved. `write_file` takes a file whole, so its bytes are +/// held until `finish` and nothing reaches the drive before then: a save +/// abandoned on the way leaves no partial file. +struct Upload { + path: String, + /// What `write_file` compares before writing, as its `base_revision`. + expected: Option, + bytes: RefCell>, + finished: Cell, +} impl GuestWriter for Upload { - fn push(&self, _data: Vec) -> Result<(), ErrorCode> { - match *self {} + fn push(&self, data: Vec) -> Result<(), ErrorCode> { + if self.finished.get() || data.len() > MAX_CHUNK as usize { + return Err(ErrorCode::InvalidInput); + } + let mut bytes = self.bytes.borrow_mut(); + if bytes.len() + data.len() > MAX_WRITE { + return Err(ErrorCode::Unsupported); + } + // A save of one transfer is kept as it arrived rather than copied. + if bytes.is_empty() { + *bytes = data; + } else { + bytes.extend_from_slice(&data); + } + Ok(()) } + /// One call is enough: the file goes out whole, and aindrive compares the + /// expected revision and writes under one lock, then answers with the new + /// revision a later conditional save compares against. fn finish(&self) -> Result, ErrorCode> { - match *self {} + if self.finished.replace(true) { + return Err(ErrorCode::InvalidInput); + } + let bytes = self.bytes.take(); + let body = mcp::write(&self.path, self.expected.as_deref(), &bytes); + let written: mcp::Written = change(&body)?; + Ok(Some(FileInfo { + kind: EntryKind::File, + size: bytes.len() as u64, + revision: written.revision, + })) } } +/// Starts saving one file under `condition`. aindrive spells create-only as +/// the revision `none`, and an unconditional save as no revision at all. +fn upload( + target: &Locator, + condition: WriteCondition, + size_hint: Option, +) -> Result { + let path = route(target); + if path.is_empty() { + return Err(ErrorCode::InvalidInput); + } + if size_hint.is_some_and(|size| size > MAX_WRITE as u64) { + return Err(ErrorCode::Unsupported); + } + let expected = match condition { + WriteCondition::CreateOnly => Some("none".to_owned()), + WriteCondition::IfMatch(revision) => Some(revision), + WriteCondition::Overwrite => None, + }; + Ok(Upload { + path, + expected, + bytes: RefCell::new(Vec::new()), + finished: Cell::new(false), + }) +} + impl Guest for Aindrive { type Reader = Download; type Writer = Upload; @@ -300,7 +400,8 @@ impl Guest for Aindrive { ) -> Result { SESSION.set(None); SESSION.set(Some(Session::from_config(&config)?)); - let probe = post(mcp::TOOLS).and_then(|pieces| mcp::tools(&joined(pieces))); + let probe = post(mcp::TOOLS, Effect::Read) + .and_then(|pieces| mcp::tools(&joined(pieces))); let tools = match probe { Ok(tools) => tools, Err(error) => { @@ -308,19 +409,23 @@ impl Guest for Aindrive { return Err(error); } }; - if !REQUIRED_TOOLS - .iter() - .all(|required| tools.iter().any(|tool| tool == required)) - { + let offers = |wanted: &[&str]| { + wanted + .iter() + .all(|wanted| tools.iter().any(|tool| tool == wanted)) + }; + if !offers(&REQUIRED_TOOLS) { SESSION.set(None); return Err(ErrorCode::Denied); } + // A read-only token is a read-only folder, not one whose saves fail. + let writable = offers(&WRITE_TOOLS); Ok(Capabilities { - write: false, + write: writable, mkdir: false, rename: false, - remove: false, - conditional_write: false, + remove: writable, + conditional_write: writable, watch: false, virtual_directories: false, }) @@ -380,11 +485,11 @@ impl Guest for Aindrive { } fn open_write( - _target: Locator, - _condition: WriteCondition, - _size_hint: Option, + target: Locator, + condition: WriteCondition, + size_hint: Option, ) -> Result { - Err(ErrorCode::Unsupported) + Ok(Writer::new(upload(&target, condition, size_hint)?)) } fn mkdir(_target: Locator) -> Result { @@ -398,8 +503,21 @@ impl Guest for Aindrive { Err(ErrorCode::Unsupported) } - fn remove(_target: Locator) -> Result<(), ErrorCode> { - Err(ErrorCode::Unsupported) + /// `delete_path` removes a folder with everything in it, and aindrive + /// leaves some entries, such as `.git`, out of every listing, so a folder + /// cannot be shown to be empty. Only files are removed. + fn remove(target: Locator) -> Result<(), ErrorCode> { + let path = route(&target); + if path.is_empty() { + return Err(ErrorCode::InvalidInput); + } + let entry: mcp::Entry = tool("stat", &json!({ "path": path }))?; + if entry.is_dir { + return Err(ErrorCode::Unsupported); + } + let _: mcp::Deleted = + change(&mcp::call("delete_path", &json!({ "path": path })))?; + Ok(()) } fn poll(_cursor: Option) -> Result { diff --git a/src/mcp.rs b/src/mcp.rs index f026b2f..31e52ad 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -1,6 +1,7 @@ //! The drive's remote MCP endpoint: one JSON-RPC `tools/call` per request, //! answered either as JSON or as a single server-sent event. +use crate::codec; use crate::sdk::ErrorCode; use serde::Deserialize; use serde::de::DeserializeOwned; @@ -25,6 +26,32 @@ pub fn call(tool: &str, arguments: &serde_json::Value) -> Vec { .into_bytes() } +/// The body of a `write_file` call carrying `bytes` as base64. It is written +/// out by hand around the encoded contents rather than serialized, which +/// would copy them again and scan every character for escaping. +pub fn write(path: &str, expected: Option<&str>, bytes: &[u8]) -> Vec { + const TAIL: &[u8] = br#""}}}"#; + let mut head = format!( + r#"{{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{{"name":"write_file","arguments":{{"path":{},"encoding":"base64","#, + serde_json::Value::from(path), + ); + if let Some(revision) = expected { + head.push_str(&format!( + r#""base_revision":{},"#, + serde_json::Value::from(revision) + )); + } + head.push_str(r#""content":""#); + let start = head.len(); + let length = codec::encoded_length(bytes.len()); + let mut body = Vec::with_capacity(start + length + TAIL.len()); + body.extend_from_slice(head.as_bytes()); + body.resize(start + length, 0); + codec::encode(bytes, &mut body[start..]); + body.extend_from_slice(TAIL); + body +} + /// What an HTTP status means before any body is read. A 404 is the drive /// itself missing: paths that are not there arrive as tool errors instead. pub fn status(code: u16) -> Result<(), ErrorCode> { @@ -190,6 +217,17 @@ fn find(haystack: &[u8], needle: &[u8]) -> Option { .position(|window| window == needle) } +/// What `write_file` answers: the file's revision once written, when the +/// drive could stat it again. +#[derive(Deserialize)] +pub struct Written { + pub revision: Option, +} + +/// What `delete_path` answers; only that it answered matters. +#[derive(Deserialize)] +pub struct Deleted {} + /// The parts of `pieces`, read as one run, that `range` covers. pub fn slices( pieces: &[Vec], @@ -285,6 +323,9 @@ fn device(message: &str) -> ErrorCode { ("ENOTDIR", ErrorCode::InvalidInput), ("ENAMETOOLONG", ErrorCode::InvalidInput), ("path too long", ErrorCode::InvalidInput), + // The call reached the device, which may have acted before going quiet. + ("agent timeout", ErrorCode::OutcomeUnknown), + ("agent disconnected", ErrorCode::OutcomeUnknown), ]; KNOWN .iter() diff --git a/src/session_tests.rs b/src/session_tests.rs index 8a97957..c5f1d6c 100644 --- a/src/session_tests.rs +++ b/src/session_tests.rs @@ -1,8 +1,11 @@ //! The provider against an in-memory aindrive. Live QA separately runs the //! compiled component through Memona's host against the real service. -use super::{Aindrive, SESSION, download, test_host}; -use crate::sdk::{EntryKind, ErrorCode, Guest, GuestReader, Locator, Setting}; +use super::{Aindrive, SESSION, download, test_host, upload}; +use crate::sdk::{ + EntryKind, ErrorCode, Guest, GuestReader, GuestWriter, Locator, Setting, + WriteCondition, +}; use std::collections::BTreeMap; const ENDPOINT: &str = "https://aindrive.ainetwork.ai/mcp/d/KQHGpl8hFXq8"; @@ -46,7 +49,9 @@ async fn opening_checks_the_token_against_the_drive_endpoint() -> Result<(), ErrorCode> { let drive = test_host::install(sample()); let capabilities = Aindrive::open("aindrive".to_owned(), config())?; - assert!(!capabilities.write && !capabilities.remove && !capabilities.watch); + assert!(capabilities.write && capabilities.remove); + assert!(capabilities.conditional_write); + assert!(!capabilities.mkdir && !capabilities.rename && !capabilities.watch); let drive = drive.borrow(); let [seen] = drive.seen.as_slice() else { return Err(ErrorCode::Internal); @@ -285,3 +290,237 @@ async fn nothing_is_requested_before_a_connection_opens() { assert_eq!(Aindrive::stat(at("/")).err(), Some(ErrorCode::Internal)); assert!(drive.borrow().seen.is_empty()); } + +/// Saves `bytes` the way the host does: one push, then one finish. +fn save( + path: &str, + condition: WriteCondition, + bytes: &[u8], +) -> Result, ErrorCode> { + let writer = upload(&at(path), condition, Some(bytes.len() as u64))?; + writer.push(bytes.to_vec())?; + let info = writer.finish()?.ok_or(ErrorCode::Internal)?; + assert_eq!(info.size, bytes.len() as u64); + Ok(info.revision) +} + +#[tokio::test] +async fn a_read_only_token_opens_a_read_only_folder() -> Result<(), ErrorCode> { + let _drive = test_host::install(test_host::Drive { + tools: Some(vec!["list_files", "read_file", "stat", "search"]), + ..sample() + }); + let capabilities = Aindrive::open("aindrive".to_owned(), config())?; + assert!(!capabilities.write && !capabilities.remove); + assert!(!capabilities.conditional_write); + Ok(()) +} + +#[tokio::test] +async fn a_save_on_the_revision_read_lands_and_reads_back() +-> Result<(), ErrorCode> { + let _drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + let (_, read) = read_from("/PAGE.md", 0, 4096)?; + let read = read.ok_or(ErrorCode::Internal)?; + let saved = save("/PAGE.md", WriteCondition::IfMatch(read), b"# Saved")?; + let (bytes, revision) = read_from("/PAGE.md", 0, 4096)?; + assert_eq!(bytes, b"# Saved"); + assert_eq!(revision, saved); + assert_eq!(Aindrive::stat(at("/PAGE.md"))?.revision, saved); + Ok(()) +} + +#[tokio::test] +async fn a_save_on_a_stale_revision_is_a_conflict_and_changes_nothing() +-> Result<(), ErrorCode> { + let _drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + let stale = WriteCondition::IfMatch("m1-s8".to_owned()); + assert_eq!(save("/PAGE.md", stale, b"lost"), Err(ErrorCode::Conflict)); + let gone = WriteCondition::IfMatch("m1-s8".to_owned()); + assert_eq!(save("/missing.md", gone, b"x"), Err(ErrorCode::Conflict)); + assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"# Studio"); + Ok(()) +} + +#[tokio::test] +async fn create_only_refuses_anything_already_there() -> Result<(), ErrorCode> { + let _drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + for taken in ["/PAGE.md", "/notes"] { + assert_eq!( + save(taken, WriteCondition::CreateOnly, b"new"), + Err(ErrorCode::Conflict), + "{taken}" + ); + } + assert!( + save("/outputs/new.md", WriteCondition::CreateOnly, b"new")?.is_some() + ); + assert_eq!(read_from("/outputs/new.md", 0, 4096)?.0, b"new"); + Ok(()) +} + +#[tokio::test] +async fn names_that_need_escaping_and_every_byte_survive_a_save() +-> Result<(), ErrorCode> { + let _drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + let bytes: Vec = (0..=255).rev().cycle().take(70_001).collect(); + let path = "/notes/\"quoted\" 노트 {x}.bin"; + save(path, WriteCondition::Overwrite, &bytes)?; + assert_eq!(read_from(path, 0, 1024 * 1024)?.0, bytes); + save("/empty.md", WriteCondition::Overwrite, b"")?; + assert_eq!(read_from("/empty.md", 0, 10)?.0, b""); + Ok(()) +} + +#[tokio::test] +async fn a_save_larger_than_one_transfer_is_refused_before_it_is_sent() +-> Result<(), ErrorCode> { + let drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + let over = Some(1024 * 1024 + 1); + assert_eq!( + upload(&at("/big.bin"), WriteCondition::Overwrite, over).err(), + Some(ErrorCode::Unsupported) + ); + let writer = upload(&at("/big.bin"), WriteCondition::Overwrite, None)?; + writer.push(vec![1; 1024 * 1024])?; + assert_eq!(writer.push(vec![1]), Err(ErrorCode::Unsupported)); + let tools: Vec<_> = + drive.borrow().seen.iter().map(|s| s.tool.clone()).collect(); + assert_eq!(tools, [None]); + Ok(()) +} + +#[tokio::test] +async fn an_abandoned_or_finished_save_sends_nothing_more() +-> Result<(), ErrorCode> { + let drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + let abandoned = upload(&at("/PAGE.md"), WriteCondition::Overwrite, None)?; + abandoned.push(b"half".to_vec())?; + drop(abandoned); + let writer = upload(&at("/PAGE.md"), WriteCondition::Overwrite, None)?; + writer.push(b"whole".to_vec())?; + writer.finish()?; + assert_eq!(writer.finish().err(), Some(ErrorCode::InvalidInput)); + assert_eq!(writer.push(b"more".to_vec()), Err(ErrorCode::InvalidInput)); + let writes = drive + .borrow() + .seen + .iter() + .filter(|seen| seen.tool.as_deref() == Some("write_file")) + .count(); + assert_eq!(writes, 1); + assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"whole"); + Ok(()) +} + +#[tokio::test] +async fn a_save_whose_answer_is_lost_has_an_unknown_outcome() +-> Result<(), ErrorCode> { + let _drive = test_host::install(test_host::Drive { + lose_reply: Some("write_file"), + ..sample() + }); + Aindrive::open("aindrive".to_owned(), config())?; + assert_eq!( + save("/PAGE.md", WriteCondition::Overwrite, b"landed"), + Err(ErrorCode::OutcomeUnknown) + ); + // It did land, which is why a retry must look first. + assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"landed"); + Ok(()) +} + +#[tokio::test] +async fn a_device_that_goes_quiet_leaves_saves_unknown_and_reads_unavailable() +-> Result<(), ErrorCode> { + for tool in ["write_file", "stat"] { + let _drive = test_host::install(test_host::Drive { + fail: Some((tool, "[internal] agent timeout")), + ..sample() + }); + Aindrive::open("aindrive".to_owned(), config())?; + let expected = match tool { + "write_file" => ErrorCode::OutcomeUnknown, + _ => ErrorCode::Unavailable, + }; + assert_eq!( + save("/PAGE.md", WriteCondition::Overwrite, b"x").err(), + (tool == "write_file").then_some(expected) + ); + assert_eq!( + Aindrive::stat(at("/PAGE.md")).err(), + (tool == "stat").then_some(expected) + ); + } + let _drive = test_host::install(test_host::Drive { + fail: Some(("write_file", "[internal] agent offline")), + ..sample() + }); + Aindrive::open("aindrive".to_owned(), config())?; + assert_eq!( + save("/PAGE.md", WriteCondition::Overwrite, b"x"), + Err(ErrorCode::Unavailable) + ); + Ok(()) +} + +#[tokio::test] +async fn files_are_removed_and_folders_are_left_alone() -> Result<(), ErrorCode> +{ + let drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + Aindrive::remove(at("/artworks.tsv"))?; + assert_eq!( + Aindrive::stat(at("/artworks.tsv")).err(), + Some(ErrorCode::NotFound) + ); + for folder in ["/notes", "/outputs"] { + assert_eq!( + Aindrive::remove(at(folder)).err(), + Some(ErrorCode::Unsupported), + "{folder}" + ); + } + assert!(drive.borrow().files.contains_key("notes/작가노트_2026.md")); + assert_eq!( + Aindrive::remove(at("/missing.md")).err(), + Some(ErrorCode::NotFound) + ); + assert_eq!( + Aindrive::remove(at("/")).err(), + Some(ErrorCode::InvalidInput) + ); + Ok(()) +} + +#[tokio::test] +async fn only_a_server_error_leaves_a_save_unknown() -> Result<(), ErrorCode> { + for (status, expected) in [ + (502, ErrorCode::OutcomeUnknown), + (429, ErrorCode::Unavailable), + (403, ErrorCode::Denied), + ] { + let drive = test_host::install(sample()); + Aindrive::open("aindrive".to_owned(), config())?; + drive.borrow_mut().status = Some(status); + assert_eq!( + save("/PAGE.md", WriteCondition::Overwrite, b"x"), + Err(expected), + "{status}" + ); + assert_eq!( + Aindrive::stat(at("/PAGE.md")).err(), + Some(match status { + 502 | 429 => ErrorCode::Unavailable, + _ => expected, + }) + ); + } + Ok(()) +} diff --git a/src/test_host.rs b/src/test_host.rs index b47c0c8..765495e 100644 --- a/src/test_host.rs +++ b/src/test_host.rs @@ -13,7 +13,13 @@ use std::rc::Rc; pub struct File { pub bytes: Vec, /// Written into replies verbatim, the way the device's float prints. - pub mtime: &'static str, + pub mtime: String, +} + +impl File { + pub fn revision(&self) -> String { + format!("m{}-s{}", self.mtime, self.bytes.len()) + } } #[derive(Default)] @@ -27,6 +33,11 @@ pub struct Drive { pub status: Option, /// Replaces a file between a stat and the read that follows it. pub change_on_read: Option, + /// A tool that answers with this error text instead of acting. + pub fail: Option<(&'static str, &'static str)>, + /// A tool that acts, but whose reply is lost on the way back. + pub lose_reply: Option<&'static str>, + pub writes: usize, pub seen: Vec, } @@ -48,10 +59,10 @@ pub fn install(drive: Drive) -> Rc> { drive } -pub fn file(bytes: &[u8], mtime: &'static str) -> File { +pub fn file(bytes: &[u8], mtime: &str) -> File { File { bytes: bytes.to_vec(), - mtime, + mtime: mtime.to_owned(), } } @@ -106,6 +117,11 @@ impl Exchange { Some(status) => (status, b"{\"error\":\"refused\"}".to_vec()), None => (200, event(&drive.answer(&call))), }; + if drive.lose_reply.is_some() + && drive.lose_reply == call["params"]["name"].as_str() + { + return Err(HostError::Unavailable); + } *self.reply.borrow_mut() = Some(reply); Ok(Response { status, @@ -152,6 +168,57 @@ fn failure(text: &str) -> Value { } impl Drive { + /// As `write_file` does: compare the expected revision, then write. + fn write(&mut self, path: &str, arguments: &Value) -> Value { + let Ok(bytes) = + STANDARD.decode(arguments["content"].as_str().unwrap_or("!")) + else { + return failure("[invalid_params] content (string) required"); + }; + let current = self.files.get(path).map(File::revision); + let folder = current.is_none() && self.is_folder(path); + let conflict = match arguments["base_revision"].as_str() { + None => false, + Some("none") => current.is_some() || folder, + Some(expected) => current.as_deref() != Some(expected), + }; + if conflict { + let current = current.as_deref().unwrap_or("none"); + return failure(&format!( + "[invalid_params] conflict: the file changed since that revision (current revision: {current})" + )); + } + if folder { + return failure("[internal] EISDIR: illegal operation on a directory"); + } + self.writes += 1; + let written = file(&bytes, &format!("17918{:08}.5", self.writes)); + let revision = written.revision(); + self.files.insert(path.to_owned(), written); + success(json!({ + "method": "write", "ok": true, "bytes": bytes.len(), "revision": revision, + })) + } + + /// As `delete_path` does: a folder goes with everything in it. + fn delete(&mut self, path: &str) -> Value { + let kind = if self.files.remove(path).is_some() { + "file" + } else if !path.is_empty() && self.is_folder(path) { + let prefix = format!("{path}/"); + self.files.retain(|file, _| !file.starts_with(&prefix)); + self + .folders + .retain(|folder| folder != path && !folder.starts_with(&prefix)); + "folder" + } else { + return failure(&format!("[not_found] no entry at {path}")); + }; + success( + json!({ "method": "delete", "ok": true, "path": path, "kind": kind }), + ) + } + fn answer(&mut self, call: &Value) -> Value { if call["method"] == "tools/list" { let tools = self.tools.clone().unwrap_or_else(|| { @@ -169,7 +236,13 @@ impl Drive { } let arguments = &call["params"]["arguments"]; let path = arguments["path"].as_str().unwrap_or_default().to_owned(); - match call["params"]["name"].as_str().unwrap_or_default() { + let tool = call["params"]["name"].as_str().unwrap_or_default(); + if let Some((_, text)) = self.fail.filter(|(failing, _)| *failing == tool) { + return failure(text); + } + match tool { + "write_file" => self.write(&path, arguments), + "delete_path" => self.delete(&path), "list_files" => match self.children(&path) { Some(entries) => success(json!({ "entries": entries })), None if self.files.contains_key(&path) => {