memona-plugin-aindrive/src/test_host.rs

//! An in-memory aindrive behind the generated synchronous `Guest` API. It
//! answers tool calls the way the drive's MCP endpoint does, as one
//! server-sent event; live QA covers Memona's real HTTP host and the service.

pub use crate::sdk::host::{Header, HostError, Method, Request, Response};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use serde_json::{Value, json};
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;

pub struct File {
  pub bytes: Vec<u8>,
  /// Written into replies verbatim, the way the device's float prints.
  pub mtime: String,
}

impl File {
  pub fn revision(&self) -> String {
    format!("m{}-s{}", self.mtime, self.bytes.len())
  }
}

#[derive(Default)]
pub struct Drive {
  /// Drive-relative paths. A folder is any prefix of a file, or listed here.
  pub files: BTreeMap<String, File>,
  pub folders: Vec<String>,
  /// The tools `tools/list` offers; every tool by default.
  pub tools: Option<Vec<&'static str>>,
  /// The status every request answers with, when not 200.
  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>,
}

pub struct Seen {
  pub url: String,
  pub authorization: Option<String>,
  pub method: String,
  pub tool: Option<String>,
}

thread_local! {
  static ACTIVE: RefCell<Option<Rc<RefCell<Drive>>>> =
    const { RefCell::new(None) };
}

pub fn install(drive: Drive) -> Rc<RefCell<Drive>> {
  let drive = Rc::new(RefCell::new(drive));
  ACTIVE.set(Some(drive.clone()));
  drive
}

pub fn file(bytes: &[u8], mtime: &str) -> File {
  File {
    bytes: bytes.to_vec(),
    mtime: mtime.to_owned(),
  }
}

pub fn http(request: &Request) -> Result<Exchange, HostError> {
  if !matches!(request.method, Method::Post) {
    return Err(HostError::InvalidInput);
  }
  let drive = ACTIVE
    .with_borrow(Clone::clone)
    .ok_or(HostError::Internal)?;
  Ok(Exchange {
    drive,
    request: request.clone(),
    body: RefCell::new(Vec::new()),
    reply: RefCell::new(None),
  })
}

pub struct Exchange {
  drive: Rc<RefCell<Drive>>,
  request: Request,
  body: RefCell<Vec<u8>>,
  reply: RefCell<Option<Vec<u8>>>,
}

impl Exchange {
  pub fn push(&self, data: &[u8]) -> Result<(), HostError> {
    if data.len() > 1024 * 1024 || self.reply.borrow().is_some() {
      return Err(HostError::InvalidInput);
    }
    self.body.borrow_mut().extend_from_slice(data);
    Ok(())
  }

  pub fn respond(&self) -> Result<Response, HostError> {
    let call: Value = serde_json::from_slice(&self.body.borrow())
      .map_err(|_| HostError::InvalidInput)?;
    let mut drive = self.drive.borrow_mut();
    let authorization = self
      .request
      .headers
      .iter()
      .find(|header| header.name == "authorization")
      .map(|header| header.value.clone());
    drive.seen.push(Seen {
      url: self.request.url.clone(),
      authorization,
      method: call["method"].as_str().unwrap_or_default().to_owned(),
      tool: call["params"]["name"].as_str().map(str::to_owned),
    });
    let (status, reply) = match drive.status {
      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,
      headers: Vec::new(),
    })
  }

  /// Hands the reply over in pieces smaller than asked for, so a caller
  /// that stops at the first short piece loses the rest.
  pub fn pull(&self, limit: u32) -> Result<Vec<u8>, HostError> {
    let mut reply = self.reply.borrow_mut();
    let reply = reply.as_mut().ok_or(HostError::InvalidInput)?;
    let count = reply.len().min(limit as usize).min(4096);
    Ok(reply.drain(..count).collect())
  }
}

fn event(message: &Value) -> Vec<u8> {
  format!("event: message\ndata: {message}\n\n").into_bytes()
}

fn success(structured: Value) -> Value {
  json!({
    "result": {
      "_meta": { "ai.aindrive/a2ui": [] },
      "content": [{ "type": "text", "text": "ok" }],
      "structuredContent": structured,
    },
    "jsonrpc": "2.0",
    "id": 1,
  })
}

fn failure(text: &str) -> Value {
  json!({
    "result": {
      "_meta": {},
      "content": [{ "type": "text", "text": text }],
      "isError": true,
    },
    "jsonrpc": "2.0",
    "id": 1,
  })
}

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(|| {
        vec![
          "list_files",
          "read_file",
          "write_file",
          "delete_path",
          "stat",
        ]
      });
      let tools: Vec<Value> =
        tools.iter().map(|name| json!({ "name": name })).collect();
      return json!({ "result": { "tools": tools }, "jsonrpc": "2.0", "id": 1 });
    }
    let arguments = &call["params"]["arguments"];
    let path = arguments["path"].as_str().unwrap_or_default().to_owned();
    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) => {
          failure("[internal] ENOTDIR: not a directory, scandir '/Users/me'")
        }
        None => failure("[internal] ENOENT: no such file or directory"),
      },
      "stat" if path.is_empty() => failure("[invalid_params] path required"),
      "stat" => match self.entry(&path) {
        Some(entry) => success(entry),
        None => failure(&format!("[not_found] no entry at {path}")),
      },
      "read_file" => {
        if let Some(changed) = self.change_on_read.take() {
          self.files.insert(path.clone(), changed);
        }
        match self.files.get(&path) {
          Some(file) => success(json!({
            "method": "read",
            "content": STANDARD.encode(&file.bytes),
            "encoding": "base64",
            "truncated": false,
          })),
          None if self.is_folder(&path) => failure("[internal] is a directory"),
          None => failure("[internal] ENOENT: no such file or directory"),
        }
      }
      _ => failure("[internal] unknown tool"),
    }
  }

  fn is_folder(&self, path: &str) -> bool {
    path.is_empty()
      || self.folders.iter().any(|folder| folder == path)
      || self
        .files
        .keys()
        .any(|file| file.starts_with(&format!("{path}/")))
  }

  fn entry(&self, path: &str) -> Option<Value> {
    let name = path.rsplit('/').next().unwrap_or(path);
    if let Some(file) = self.files.get(path) {
      let entry = format!(
        r#"{{"name":{},"path":{},"isDir":false,"size":{},"mtimeMs":{}}}"#,
        json!(name),
        json!(path),
        file.bytes.len(),
        file.mtime,
      );
      return serde_json::from_str(&entry).ok();
    }
    self.is_folder(path).then(|| {
      json!({ "name": name, "path": path, "isDir": true, "size": 96,
        "mtimeMs": 1_791_726_544_422.142_f64 })
    })
  }

  fn children(&self, path: &str) -> Option<Vec<Value>> {
    if !self.is_folder(path) {
      return None;
    }
    let prefix = if path.is_empty() {
      String::new()
    } else {
      format!("{path}/")
    };
    let mut names: Vec<&str> = self
      .files
      .keys()
      .chain(self.folders.iter())
      .filter_map(|item| item.strip_prefix(&prefix))
      .filter_map(|rest| rest.split('/').next())
      .filter(|name| !name.is_empty())
      .collect();
    names.dedup();
    names.sort_unstable();
    names.dedup();
    Some(
      names
        .into_iter()
        .filter_map(|name| self.entry(&format!("{prefix}{name}")))
        .collect(),
    )
  }
}