Save and delete files in an aindrive drive without overwriting another change #2

Merged
RaymondCho merged 2 commits from save-and-delete into main 2026-10-13 21:54:07 +00:00
6 changed files with 615 additions and 54 deletions
Showing only changes of commit f95c58f9e9 - Show all commits
+21 -8
View File
@@ -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
+84 -7
View File
@@ -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;
@@ -76,7 +76,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<Item = &'a [u8]>,
out: &mut [u8],
) -> Result<(), ErrorCode> {
@@ -134,9 +134,72 @@ pub fn last(chars: [u8; 4]) -> Result<Vec<u8>, 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;
@@ -175,6 +238,20 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn every_length_encodes_as_a_general_encoder_would() {
let bytes: Vec<u8> = (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"] {
+150 -32
View File
@@ -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<Self, ErrorCode> {
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<Vec<Vec<u8>>, ErrorCode> {
fn post(body: &[u8], effect: Effect) -> Result<Vec<Vec<u8>>, ErrorCode> {
let (url, authorization) = SESSION.with_borrow(|session| {
session
.as_ref()
@@ -116,12 +130,27 @@ fn post(body: &[u8]) -> Result<Vec<Vec<u8>>, 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<u8>>) -> Vec<u8> {
}
}
/// 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<T: DeserializeOwned>(
name: &str,
arguments: &serde_json::Value,
) -> Result<T, ErrorCode> {
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<T: DeserializeOwned>(body: &[u8]) -> Result<T, ErrorCode> {
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<Download, ErrorCode> {
} 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<Download, ErrorCode> {
})
}
/// 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<String>,
bytes: RefCell<Vec<u8>>,
finished: Cell<bool>,
}
impl GuestWriter for Upload {
fn push(&self, _data: Vec<u8>) -> Result<(), ErrorCode> {
match *self {}
fn push(&self, data: Vec<u8>) -> 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<Option<FileInfo>, 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<u64>,
) -> Result<Upload, ErrorCode> {
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<Capabilities, ErrorCode> {
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<u64>,
target: Locator,
condition: WriteCondition,
size_hint: Option<u64>,
) -> Result<Writer, ErrorCode> {
Err(ErrorCode::Unsupported)
Ok(Writer::new(upload(&target, condition, size_hint)?))
}
fn mkdir(_target: Locator) -> Result<FileInfo, ErrorCode> {
@@ -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<String>) -> Result<ChangePage, ErrorCode> {
+41
View File
@@ -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<u8> {
.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<u8> {
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<usize> {
.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<String>,
}
/// 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<u8>],
@@ -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()
+242 -3
View File
@@ -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<Option<String>, 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<u8> = (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(())
}
+77 -4
View File
@@ -13,7 +13,13 @@ use std::rc::Rc;
pub struct File {
pub bytes: Vec<u8>,
/// 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<u16>,
/// Replaces a file between a stat and the read that follows it.
pub change_on_read: Option<File>,
/// 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<Seen>,
}
@@ -48,10 +59,10 @@ pub fn install(drive: Drive) -> Rc<RefCell<Drive>> {
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) => {