use super::*;
use crate::serve::{
frame,
protocol::{PROTOCOL, RECEIPT_CAPABILITY},
};
use std::{collections::BTreeMap, sync::Mutex};
pub type Rows = Arc<Mutex<BTreeMap<String, Value>>>;
pub struct Store {
pub calls: Arc<Mutex<Vec<Value>>>,
worker: std::thread::JoinHandle<Result>,
path: std::path::PathBuf,
}
impl Store {
pub fn finish(self) -> Result<Vec<Value>> {
self.worker
.join()
.map_err(|_| anyhow::anyhow!("store panic"))??;
std::fs::remove_file(self.path)?;
Ok(self
.calls
.lock()
.map_err(|_| anyhow::anyhow!("calls poisoned"))?
.clone())
}
}
pub fn rows(items: Vec<Value>) -> Rows {
Arc::new(Mutex::new(
items
.into_iter()
.map(|v| {
let conversation = v["conversation"].as_str().unwrap_or("conversation");
let suffix = v["seq"]
.as_u64()
.map_or_else(|| "meta".into(), |s| format!("{s:020}"));
(format!("{PERSON}/{conversation}/{suffix}"), v)
})
.collect(),
))
}
pub fn meta(conversation: &str) -> Value {
json!({"conversation":conversation,"title":conversation,"harness":"codex","created_at_ms":1,"updated_at_ms":2,"turns":3,"state":"live"})
}
pub fn start(
f: &Fixture,
records: Rows,
during: impl FnOnce(&Rows) -> Result + Send + 'static,
) -> Result<Store> {
with_frame(f, records, 16_777_216, during)
}
pub fn with_frame(
f: &Fixture,
records: Rows,
frame_bytes: u32,
during: impl FnOnce(&Rows) -> Result + Send + 'static,
) -> Result<Store> {
let path = f.files.path().join("store.sock");
let listener = std::os::unix::net::UnixListener::bind(&path)?;
let calls = Arc::new(Mutex::new(Vec::new()));
let captured = Arc::clone(&calls);
let worker = std::thread::spawn(move || -> Result {
let (mut socket, _) = listener.accept()?;
let cap =
std::num::NonZeroU32::new(frame_bytes).ok_or_else(|| anyhow::anyhow!("frame cap"))?;
let mut during = Some(during);
let mut versions = BTreeMap::new();
let mut head = 0u64;
while let Some(bytes) = frame::read(&mut socket, cap)? {
let request: Value = serde_json::from_slice(&bytes)?;
let method = request["operation"]["method"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("method"))?;
let p = &request["operation"]["params"];
captured
.lock()
.map_err(|_| anyhow::anyhow!("calls poisoned"))?
.push(request.clone());
let result = match method {
"hello" => {
json!({"protocol":PROTOCOL,"capabilities":[RECEIPT_CAPABILITY],"max_page":64,"max_frame":cap.get()})
}
"session.pass" => {
assert_eq!(p["pass"], "person-pass");
json!({"holder":PERSON,"expires":crate::serve::pass::now()+3600})
}
"documents.list" | "documents.get" => {
assert_eq!(p["selection"]["branch"], "main");
assert_eq!(p["collection"], "_assistant");
let version = if p["selection"]["version"].is_null() {
head += 1;
versions.insert(
head,
records
.lock()
.map_err(|_| anyhow::anyhow!("rows poisoned"))?
.clone(),
);
head
} else {
p["selection"]["version"]["sequence"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("version"))?
.parse::<u64>()?
};
if let Some(during) = during.take() {
during(&records)?;
}
let records = versions
.get(&version)
.ok_or_else(|| anyhow::anyhow!("unknown snapshot"))?;
if method == "documents.get" {
let id = p["document_id"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("id"))?;
let document = records.get(id).map(|v| json!({"id":id,"raw_json":v.to_string(),"revision":{"generation":"1","sequence":"1"}}));
json!({"document":document,"selection":{"branch":"main","version":{"generation":"1","sequence":version.to_string()}}})
} else {
page(records, p, version)?
}
}
_ => anyhow::bail!("unexpected method {method}"),
};
let mut reply = serde_json::to_vec(&json!({"id":request["id"],"result":result}))?;
if reply.len() > cap.get() as usize {
reply = serde_json::to_vec(
&json!({"id":request["id"],"error":{"code":"response_limit","message":"response exceeds frame limit; use a smaller page"}}),
)?;
}
frame::write(&mut socket, &reply, cap)?;
}
Ok(())
});
Ok(Store {
calls,
worker,
path,
})
}
fn page(records: &BTreeMap<String, Value>, p: &Value, version: u64) -> Result<Value> {
let prefix = p["prefix"].as_str().unwrap_or("");
let after = p["after"].as_str();
let limit = p["limit"]
.as_u64()
.ok_or_else(|| anyhow::anyhow!("limit"))?;
let limit = usize::try_from(limit)?;
let all: Vec<_> = records
.iter()
.filter(|(k, _)| k.starts_with(prefix) && after.is_none_or(|a| k.as_str() > a))
.collect();
let entries: Vec<_> = all
.iter()
.take(limit)
.map(|(id, v)| {
let mut entry = json!({"id":id,"revision":{"generation":"1","sequence":"1"}});
if p["content"] != "none" {
entry["raw_json"] = json!(v.to_string());
}
entry
})
.collect();
let next = if all.len() > limit {
entries.last().map(|e| e["id"].clone())
} else {
None
};
Ok(
json!({"entries":entries,"next_after":next,"version":{"generation":"1","sequence":version.to_string()}}),
)
}