memona-plugin-aindrive/src/session_tests.rs

//! 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, 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";

fn config() -> Vec<(String, Setting)> {
  vec![
    (
      "url".to_owned(),
      Setting::Text("https://aindrive.ainetwork.ai/d/KQHGpl8hFXq8".to_owned()),
    ),
    (
      "token".to_owned(),
      Setting::Secret(" aind_pat_secret\n".to_owned()),
    ),
  ]
}

fn at(path: &str) -> Locator {
  Locator::Path(path.to_owned())
}

fn sample() -> test_host::Drive {
  let mut files = BTreeMap::new();
  for (path, bytes, mtime) in [
    ("PAGE.md", &b"# Studio"[..], "1791731214739.939"),
    ("notes/작가노트_2026.md", &b"notes"[..], "1791726579100"),
    ("notes/odd\\name.md", &b"x"[..], "1"),
    ("artworks.tsv", &b"title\tyear"[..], "1791726579103.648"),
  ] {
    files.insert(path.to_owned(), test_host::file(bytes, mtime));
  }
  test_host::Drive {
    files,
    folders: vec!["outputs".to_owned()],
    ..test_host::Drive::default()
  }
}

#[tokio::test]
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);
  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);
  };
  assert_eq!(seen.url, ENDPOINT);
  assert_eq!(seen.method, "tools/list");
  assert_eq!(
    seen.authorization.as_deref(),
    Some("Bearer aind_pat_secret")
  );
  Ok(())
}

#[tokio::test]
async fn a_refused_token_or_a_missing_drive_leaves_no_connection() {
  for (status, expected) in
    [(401, ErrorCode::Denied), (404, ErrorCode::NotFound)]
  {
    let _drive = test_host::install(test_host::Drive {
      status: Some(status),
      ..sample()
    });
    assert_eq!(
      Aindrive::open("aindrive".to_owned(), config()).err(),
      Some(expected)
    );
    assert!(SESSION.with_borrow(Option::is_none));
  }
}

#[tokio::test]
async fn a_token_that_cannot_read_files_is_refused() {
  let _drive = test_host::install(test_host::Drive {
    tools: Some(vec!["search"]),
    ..sample()
  });
  assert_eq!(
    Aindrive::open("aindrive".to_owned(), config()).err(),
    Some(ErrorCode::Denied)
  );
  assert!(SESSION.with_borrow(Option::is_none));
}

#[tokio::test]
async fn a_config_without_a_token_or_drive_is_invalid() {
  let _drive = test_host::install(sample());
  let only_url = config().into_iter().take(1).collect();
  let only_token = config().into_iter().skip(1).collect();
  for config in [only_url, only_token] {
    assert_eq!(
      Aindrive::open("aindrive".to_owned(), config).err(),
      Some(ErrorCode::InvalidInput)
    );
  }
}

#[tokio::test]
async fn listing_pages_through_a_folder_and_hides_unshowable_names()
-> Result<(), ErrorCode> {
  let _drive = test_host::install(sample());
  Aindrive::open("aindrive".to_owned(), config())?;
  let first = Aindrive::list(at("/"), None, 2)?;
  let names: Vec<_> = first.entries.iter().map(|e| e.name.as_str()).collect();
  assert_eq!(names, ["PAGE.md", "artworks.tsv"]);
  assert_eq!(
    first.entries[0].info.revision.as_deref(),
    Some("m1791731214739.939-s8")
  );
  let rest = Aindrive::list(at("/"), first.cursor, 2)?;
  let names: Vec<_> = rest.entries.iter().map(|e| e.name.as_str()).collect();
  assert_eq!(names, ["notes", "outputs"]);
  assert!(rest.cursor.is_none());
  assert!(matches!(rest.entries[0].info.kind, EntryKind::Directory));
  assert_eq!(rest.entries[0].info.size, 0);
  assert!(rest.entries[0].info.revision.is_none());

  let notes = Aindrive::list(at("/notes/"), None, 10)?;
  let names: Vec<_> = notes.entries.iter().map(|e| e.name.as_str()).collect();
  assert_eq!(names, ["작가노트_2026.md"]);
  Ok(())
}

#[tokio::test]
async fn listing_a_file_or_a_missing_folder_says_which() -> Result<(), ErrorCode>
{
  let _drive = test_host::install(sample());
  Aindrive::open("aindrive".to_owned(), config())?;
  assert_eq!(
    Aindrive::list(at("/PAGE.md"), None, 10).err(),
    Some(ErrorCode::InvalidInput)
  );
  assert_eq!(
    Aindrive::list(at("/gone"), None, 10).err(),
    Some(ErrorCode::NotFound)
  );
  assert_eq!(
    Aindrive::list(at("/"), Some("x".to_owned()), 10).err(),
    Some(ErrorCode::InvalidInput)
  );
  Ok(())
}

#[tokio::test]
async fn stat_reports_the_root_files_and_folders() -> Result<(), ErrorCode> {
  let _drive = test_host::install(sample());
  Aindrive::open("aindrive".to_owned(), config())?;
  let root = Aindrive::stat(at("/"))?;
  assert!(matches!(root.kind, EntryKind::Directory));
  let page = Aindrive::stat(at("/PAGE.md"))?;
  assert!(matches!(page.kind, EntryKind::File));
  assert_eq!(page.size, 8);
  assert_eq!(page.revision.as_deref(), Some("m1791731214739.939-s8"));
  assert!(matches!(
    Aindrive::stat(at("/outputs"))?.kind,
    EntryKind::Directory
  ));
  assert_eq!(
    Aindrive::stat(at("/missing.md")).err(),
    Some(ErrorCode::NotFound)
  );
  Ok(())
}

/// Reads all of a file from `offset`, as the host does, in pulls of
/// `limit` bytes that may come back shorter.
fn read_from(
  path: &str,
  offset: u64,
  limit: u32,
) -> Result<(Vec<u8>, Option<String>), ErrorCode> {
  let reader = download(&at(path), offset)?;
  let mut read = Vec::new();
  loop {
    let piece = reader.pull(limit)?;
    assert!(piece.len() <= limit as usize);
    if piece.is_empty() {
      return Ok((read, reader.revision()));
    }
    read.extend_from_slice(&piece);
  }
}

#[tokio::test]
async fn reads_start_anywhere_and_never_exceed_the_limit()
-> Result<(), ErrorCode> {
  // Longer than one of the test host's 4 KiB reply pieces, so groups of
  // base64 characters are split between pieces.
  let bytes: Vec<u8> = (0..=255).cycle().take(20_000).collect();
  let mut drive = sample();
  drive
    .files
    .insert("data.bin".to_owned(), test_host::file(&bytes, "5"));
  let _drive = test_host::install(drive);
  Aindrive::open("aindrive".to_owned(), config())?;
  for offset in [0_u64, 1, 2, 3, 3070, 3071, 10_000, 19_998, 19_999] {
    for limit in [1, 7, 4096, 1024 * 1024] {
      let (read, revision) = read_from("/data.bin", offset, limit)?;
      assert_eq!(read, bytes[offset as usize..], "{offset} by {limit}");
      assert_eq!(revision.as_deref(), Some("m5-s20000"));
    }
  }
  assert!(read_from("/data.bin", 20_000, 10)?.0.is_empty());
  assert!(read_from("/data.bin", 50_000, 10)?.0.is_empty());
  Ok(())
}

#[tokio::test]
async fn files_shorter_than_a_group_and_empty_files_read_whole()
-> Result<(), ErrorCode> {
  let mut drive = sample();
  for (name, bytes) in [("one", &b"a"[..]), ("two", b"ab"), ("empty", b"")] {
    drive
      .files
      .insert(name.to_owned(), test_host::file(bytes, "1"));
  }
  let _drive = test_host::install(drive);
  Aindrive::open("aindrive".to_owned(), config())?;
  assert_eq!(read_from("/one", 0, 10)?.0, b"a");
  assert_eq!(read_from("/two", 1, 10)?.0, b"b");
  assert_eq!(read_from("/empty", 0, 10)?.0, b"");
  Ok(())
}

#[tokio::test]
async fn a_file_changed_between_stat_and_read_is_a_conflict()
-> Result<(), ErrorCode> {
  let _drive = test_host::install(test_host::Drive {
    change_on_read: Some(test_host::file(b"# Studio, longer now", "2")),
    ..sample()
  });
  Aindrive::open("aindrive".to_owned(), config())?;
  assert_eq!(
    download(&at("/PAGE.md"), 0).err(),
    Some(ErrorCode::Conflict)
  );
  Ok(())
}

#[tokio::test]
async fn folders_missing_files_and_oversized_files_are_not_read()
-> Result<(), ErrorCode> {
  let mut drive = sample();
  drive.files.insert(
    "large.bin".to_owned(),
    test_host::file(&vec![0; 8 * 1024 * 1024 + 1], "9"),
  );
  let drive = test_host::install(drive);
  Aindrive::open("aindrive".to_owned(), config())?;
  assert_eq!(
    download(&at("/notes"), 0).err(),
    Some(ErrorCode::InvalidInput)
  );
  assert_eq!(
    download(&at("/missing.md"), 0).err(),
    Some(ErrorCode::NotFound)
  );
  drive.borrow_mut().seen.clear();
  assert_eq!(
    download(&at("/large.bin"), 0).err(),
    Some(ErrorCode::Unsupported)
  );
  let tools: Vec<_> = drive
    .borrow()
    .seen
    .iter()
    .map(|seen| seen.tool.clone())
    .collect();
  assert_eq!(tools, [Some("stat".to_owned())]);
  Ok(())
}

#[tokio::test]
async fn nothing_is_requested_before_a_connection_opens() {
  let drive = test_host::install(sample());
  Aindrive::close();
  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(())
}