use super::assistant_wire::{identifier, item};
use crate::{
client::{Client, Reply},
serve::storage::Failure,
};
use serde_json::{Value, json};
use std::{path::Path, time::Duration};
fn call(client: &mut Client, method: &str, params: &Value) -> Result<Value, Failure> {
let raw = serde_json::value::to_raw_value(params).map_err(|_| invalid())?;
match client
.call(method, Some(&raw))
.map_err(|error| Failure::new("assistant_record_unavailable", error.why()))?
{
Reply::Answered(raw) => serde_json::from_str(raw.get()).map_err(|_| invalid()),
Reply::Refused(error) => Err(Failure::new(&error.code, error.message)),
}
}
pub(super) struct Record<'a> {
pub socket: &'a Path,
pub timeout: Duration,
pub pass: &'a str,
pub person: &'a str,
pub bound: usize,
}
pub(super) struct Reader {
client: Client,
person: String,
bound: usize,
version: Value,
}
impl Record<'_> {
pub fn open(&self) -> Result<Reader, Failure> {
let mut client = Client::connect(self.socket, self.timeout)
.map_err(|error| Failure::new("assistant_record_unavailable", error.why()))?;
call(&mut client, "session.pass", &json!({"pass":self.pass}))?;
Ok(Reader {
client,
person: self.person.to_owned(),
bound: self.bound,
version: Value::Null,
})
}
}
impl Reader {
fn selection(&self) -> Value {
json!({"branch":"main","version":self.version})
}
fn capture(&mut self, version: &Value) -> Result<(), Failure> {
serde_json::from_value::<crate::serve::protocol::Version>(version.clone())
.map_err(|_| invalid())?;
if self.version.is_null() {
self.version = version.clone();
} else if &self.version != version {
return Err(invalid());
}
Ok(())
}
fn list(
&mut self,
prefix: &str,
after: Option<&str>,
limit: usize,
content: &str,
) -> Result<Value, Failure> {
let params = json!({"selection":self.selection(),"collection":"_assistant","prefix":prefix,"content":content,"after":after,"limit":limit});
let page = call(&mut self.client, "documents.list", ¶ms)?;
self.capture(page.get("version").ok_or_else(invalid)?)?;
let entries = page["entries"].as_array().ok_or_else(invalid)?;
if entries.len() > limit {
return Err(invalid());
}
let mut previous = after;
for entry in entries {
let id = entry["id"]
.as_str()
.filter(|id| id.starts_with(prefix) && previous.is_none_or(|old| *id > old))
.ok_or_else(invalid)?;
previous = Some(id);
}
let next = page.get("next_after").ok_or_else(invalid)?;
if !next.is_null() && (entries.is_empty() || next.as_str() != previous) {
return Err(invalid());
}
Ok(page)
}
pub fn replay(
&mut self,
conversation: &str,
after: u64,
mut visit: impl FnMut(Value) -> Result<(), Failure>,
) -> Result<(), Failure> {
self.version = Value::Null;
let prefix = format!("{}/{conversation}/", self.person);
let mut after = Some(format!("{prefix}{after:020}"));
let mut limit = 64usize.min(self.client.max_page().get() as usize);
loop {
let page = self.list(&prefix, after.as_deref(), limit, "full")?;
let entries = page["entries"].as_array().ok_or_else(invalid)?;
let bytes = entries.iter().try_fold(0usize, |bytes, entry| {
let raw = entry["raw_json"].as_str().ok_or_else(invalid)?;
Ok::<_, Failure>(bytes.saturating_add(raw.len()))
})?;
if bytes > self.bound {
if limit == 1 {
return Err(Failure::new(
"assistant_record_item_too_large",
format!("events_after_bytes={}; count={bytes}", self.bound),
));
}
limit = (limit / 2).max(1);
continue;
}
for entry in entries {
let key = entry["id"].as_str().ok_or_else(invalid)?;
if key.ends_with("/meta") {
continue;
}
let value: Value =
serde_json::from_str(entry["raw_json"].as_str().ok_or_else(invalid)?)
.map_err(|_| invalid())?;
item(&value)?;
let seq = value["seq"].as_u64().ok_or_else(invalid)?;
if value["conversation"] != conversation
|| value["durable"] != true
|| key != format!("{prefix}{seq:020}")
{
return Err(invalid());
}
visit(value)?;
}
if page["next_after"].is_null() {
return Ok(());
}
after = Some(page["next_after"].as_str().ok_or_else(invalid)?.to_owned());
}
}
fn conversations(&mut self) -> Result<Vec<Value>, Failure> {
self.version = Value::Null;
let prefix = format!("{}/", self.person);
let mut after: Option<String> = None;
let mut found = Vec::new();
let mut bytes = 0usize;
loop {
let page = self.list(&prefix, after.as_deref(), 1, "none")?;
let entries = page["entries"].as_array().ok_or_else(invalid)?;
let Some(entry) = entries.first() else {
break;
};
let key = entry["id"].as_str().ok_or_else(invalid)?;
let (id, _) = key
.strip_prefix(&prefix)
.and_then(|key| key.split_once('/'))
.filter(|(id, _)| identifier(id))
.ok_or_else(invalid)?;
let params = json!({"selection":self.selection(),"collection":"_assistant","document_id":format!("{prefix}{id}/meta")});
let document = call(&mut self.client, "documents.get", ¶ms)?;
self.capture(&document["selection"]["version"])?;
let raw = document["document"]["raw_json"]
.as_str()
.ok_or_else(invalid)?;
let mut meta: Value = serde_json::from_str(raw).map_err(|_| invalid())?;
if !meta["title"].is_string()
|| meta.get("harness").is_none_or(Value::is_null)
|| !meta["created_at_ms"].is_u64()
|| !meta["updated_at_ms"].is_u64()
|| !meta["turns"].is_u64()
|| !matches!(meta["state"].as_str(), Some("live" | "dormant" | "ended"))
{
return Err(invalid());
}
meta["id"] = json!(id);
bytes = bytes.saturating_add(meta.to_string().len());
if bytes > self.bound {
return Err(super::assistant_state::bound(
"events_after_bytes",
self.bound,
bytes,
));
}
found.push(meta);
after = Some(format!("{prefix}{id}0"));
}
found.sort_by(|a, b| {
b["updated_at_ms"]
.as_u64()
.cmp(&a["updated_at_ms"].as_u64())
});
Ok(found)
}
}
pub(super) fn conversations(
socket: &Path,
timeout: Duration,
pass: &str,
person: &str,
bound: usize,
) -> Result<Vec<Value>, Failure> {
Record {
socket,
timeout,
pass,
person,
bound,
}
.open()?
.conversations()
}
fn invalid() -> Failure {
Failure::new(
"assistant_record_invalid",
"the assistant record page violates its key, item or cursor contract",
)
}