use std::collections::HashMap;
use std::sync::Arc;
use kimun_core::{IndexObserver, NoteVault, error::VaultError, nfs::VaultPath};
use crate::server_client::dto::{WireDoc, WireSection};
use crate::server_client::{
DirtyOp, DirtySet, RagClient, RagError, RagObserver, RagTransport, hash_string, reconcile_diff,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServerCapability {
Unconfigured,
SemanticOnly,
Full,
}
impl ServerCapability {
pub fn from_health(health: &crate::server_client::dto::Health) -> Self {
match (health.embedder.is_some(), health.llm_provider.is_some()) {
(false, _) => ServerCapability::Unconfigured,
(true, false) => ServerCapability::SemanticOnly,
(true, true) => ServerCapability::Full,
}
}
pub fn llm_available(self) -> bool {
matches!(self, ServerCapability::Full)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ServerProbe {
pub capability: ServerCapability,
pub auth_required: bool,
}
pub struct RagSync {
vault: Arc<NoteVault>,
dirty: Arc<DirtySet>,
observer: Arc<dyn IndexObserver>,
client: RagClient,
}
impl RagSync {
pub fn new(vault: Arc<NoteVault>, client: RagClient) -> Self {
let dirty = Arc::new(DirtySet::default());
let observer: Arc<dyn IndexObserver> = Arc::new(RagObserver::new(dirty.clone()));
vault.set_index_observer(observer.clone());
Self {
vault,
dirty,
observer,
client,
}
}
pub async fn probe(&self) -> Option<ServerProbe> {
self.client.health().await.ok().map(|h| ServerProbe {
capability: ServerCapability::from_health(&h),
auth_required: h.auth_required,
})
}
pub fn index_ready(&self) -> bool {
self.vault.index_ready()
}
pub async fn tick(&self) -> Result<bool, RagError> {
let drained = drain(&self.vault, &self.dirty, &self.client).await?;
let reconciled = reconcile(&self.vault, &self.client).await?;
Ok(drained && reconciled)
}
pub async fn drain(&self) -> Result<bool, RagError> {
drain(&self.vault, &self.dirty, &self.client).await
}
pub async fn reconcile(&self) -> Result<bool, RagError> {
reconcile(&self.vault, &self.client).await
}
pub fn client(&self) -> &RagClient {
&self.client
}
}
impl Drop for RagSync {
fn drop(&mut self) {
self.vault.clear_index_observer_if(&self.observer);
}
}
pub async fn build_doc(
vault: &NoteVault,
path: &VaultPath,
hash: u64,
) -> Result<Option<WireDoc>, VaultError> {
let chunks = vault.get_note_chunks(path).await?;
let sections: Vec<WireSection> = chunks
.into_values()
.flatten()
.map(|c| WireSection {
title: c.get_breadcrumb().to_string(),
text: c.get_text().to_string(),
})
.collect();
if sections.is_empty() {
return Ok(None);
}
Ok(Some(WireDoc {
path: path.to_string(),
hash: hash_string(hash),
sections,
}))
}
pub async fn drain<T: RagTransport>(
vault: &NoteVault,
dirty: &DirtySet,
transport: &T,
) -> Result<bool, RagError> {
if !vault.index_ready() {
return Ok(false);
}
let ops = dirty.drain();
if ops.is_empty() {
return Ok(true);
}
let mut upserts: Vec<(VaultPath, u64)> = Vec::new();
let mut deletes: Vec<String> = Vec::new();
for (path, op) in ops {
match op {
DirtyOp::Upsert(hash) => upserts.push((path, hash)),
DirtyOp::Delete => deletes.push(path.to_string()),
}
}
let mut docs = Vec::new();
let mut built: Vec<(VaultPath, u64)> = Vec::new();
for (path, hash) in upserts {
match build_doc(vault, &path, hash).await {
Ok(Some(doc)) => {
docs.push(doc);
built.push((path, hash));
}
Ok(None) => deletes.push(path.to_string()),
Err(_) => dirty.requeue([(path, DirtyOp::Upsert(hash))]),
}
}
let mut first_err: Option<RagError> = None;
if !docs.is_empty()
&& let Err(e) = transport.push_docs(docs).await
{
dirty.requeue(built.into_iter().map(|(p, h)| (p, DirtyOp::Upsert(h))));
first_err = Some(e);
}
if !deletes.is_empty() {
let paths_for_requeue: Vec<VaultPath> = deletes.iter().map(VaultPath::new).collect();
if let Err(e) = transport.delete_paths(deletes).await {
dirty.requeue(paths_for_requeue.into_iter().map(|p| (p, DirtyOp::Delete)));
first_err = first_err.or(Some(e));
}
}
match first_err {
Some(e) => Err(e),
None => Ok(true),
}
}
pub async fn reconcile<T: RagTransport>(
vault: &NoteVault,
transport: &T,
) -> Result<bool, RagError> {
if !vault.index_ready() {
return Ok(false);
}
let notes = vault
.get_all_notes()
.await
.map_err(|e| RagError::Protocol(format!("read vault notes: {e}")))?;
let local_hashes: HashMap<String, u64> = notes
.into_iter()
.map(|(entry, content)| (entry.path.to_string(), content.hash))
.collect();
let local_str: HashMap<String, String> = local_hashes
.iter()
.map(|(p, h)| (p.clone(), hash_string(*h)))
.collect();
let server = transport.server_hashes().await?;
let plan = reconcile_diff(&local_str, &server);
let mut docs = Vec::new();
let mut to_delete = plan.to_delete;
for path_str in &plan.to_push {
let hash = local_hashes[path_str];
match build_doc(vault, &VaultPath::new(path_str), hash)
.await
.map_err(|e| RagError::Protocol(format!("build doc {path_str}: {e}")))?
{
Some(doc) => docs.push(doc),
None => {
if server.contains_key(path_str) {
to_delete.push(path_str.clone());
}
}
}
}
if !docs.is_empty() {
transport.push_docs(docs).await?;
}
if !to_delete.is_empty() {
transport.delete_paths(to_delete).await?;
}
Ok(true)
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
#[test]
fn capability_from_health_fields() {
use crate::server_client::dto::Health;
let h = |embedder: Option<&str>, llm: Option<&str>| Health {
status: "ok".into(),
reranker: false,
embedder: embedder.map(str::to_string),
llm_provider: llm.map(str::to_string),
auth_required: false,
};
assert_eq!(
ServerCapability::from_health(&h(None, None)),
ServerCapability::Unconfigured
);
assert_eq!(
ServerCapability::from_health(&h(None, Some("gemini"))),
ServerCapability::Unconfigured
);
assert_eq!(
ServerCapability::from_health(&h(Some("fastembed"), None)),
ServerCapability::SemanticOnly
);
assert_eq!(
ServerCapability::from_health(&h(Some("fastembed"), Some("gemini"))),
ServerCapability::Full
);
}
use kimun_core::VaultConfig;
use std::sync::Mutex;
use tempfile::TempDir;
#[derive(Default)]
struct FakeTransport {
pushed: Mutex<Vec<WireDoc>>,
deleted: Mutex<Vec<String>>,
server: Mutex<HashMap<String, String>>,
fail_push: Mutex<bool>,
}
#[async_trait]
impl RagTransport for FakeTransport {
async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError> {
if *self.fail_push.lock().unwrap() {
return Err(RagError::Protocol("boom".into()));
}
self.pushed.lock().unwrap().extend(docs);
Ok(())
}
async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
self.deleted.lock().unwrap().extend(paths);
Ok(())
}
async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
Ok(self.server.lock().unwrap().clone())
}
}
fn register(vault: &NoteVault) -> Arc<DirtySet> {
let dirty = Arc::new(DirtySet::default());
vault.set_index_observer(Arc::new(RagObserver::new(dirty.clone())));
dirty
}
fn sys(path: impl AsRef<std::path::Path>) -> kimun_core::SystemPath {
kimun_core::SystemPath::try_absolute(path).expect("test path must be absolute")
}
async fn vault(dir: &std::path::Path) -> NoteVault {
let vault = NoteVault::new(VaultConfig::new(sys(dir))).await.unwrap();
vault.validate_and_init().await.unwrap();
vault
}
#[tokio::test]
async fn drain_pushes_created_note_and_deletes_removed() {
let dir = TempDir::new().unwrap();
let vault = vault(dir.path()).await;
let dirty = register(&vault);
let transport = FakeTransport::default();
vault
.create_note(&VaultPath::new("a.md"), "# Title\n\nbody")
.await
.unwrap();
drain(&vault, &dirty, &transport).await.unwrap();
{
let pushed = transport.pushed.lock().unwrap();
assert_eq!(pushed.len(), 1);
assert_eq!(pushed[0].path, "/a.md"); assert!(!pushed[0].sections.is_empty());
assert!(dirty.is_empty());
}
vault.delete_note(&VaultPath::new("a.md")).await.unwrap();
drain(&vault, &dirty, &transport).await.unwrap();
assert_eq!(
*transport.deleted.lock().unwrap(),
vec!["/a.md".to_string()]
);
}
#[tokio::test]
async fn failed_push_requeues() {
let dir = TempDir::new().unwrap();
let vault = vault(dir.path()).await;
let dirty = register(&vault);
let transport = FakeTransport::default();
*transport.fail_push.lock().unwrap() = true;
vault
.create_note(&VaultPath::new("a.md"), "body")
.await
.unwrap();
assert!(drain(&vault, &dirty, &transport).await.is_err());
assert_eq!(dirty.len(), 1);
}
#[tokio::test]
async fn reconcile_pushes_missing_and_deletes_stale() {
let dir = TempDir::new().unwrap();
let vault = vault(dir.path()).await;
let _dirty = register(&vault);
let transport = FakeTransport::default();
vault
.create_note(&VaultPath::new("keep.md"), "kept")
.await
.unwrap();
transport
.server
.lock()
.unwrap()
.insert("/gone.md".to_string(), "oldhash".to_string());
assert!(reconcile(&vault, &transport).await.unwrap());
let pushed = transport.pushed.lock().unwrap();
assert!(pushed.iter().any(|d| d.path == "/keep.md"));
assert_eq!(
*transport.deleted.lock().unwrap(),
vec!["/gone.md".to_string()]
);
}
#[tokio::test]
async fn reconcile_skipped_while_index_not_ready() {
let dir = TempDir::new().unwrap();
let vault = NoteVault::new(VaultConfig::new(sys(dir.path())))
.await
.unwrap();
assert!(!vault.index_ready());
let transport = FakeTransport::default();
transport
.server
.lock()
.unwrap()
.insert("/precious.md".to_string(), "hash".to_string());
assert!(!reconcile(&vault, &transport).await.unwrap());
assert!(transport.deleted.lock().unwrap().is_empty());
assert!(transport.pushed.lock().unwrap().is_empty());
let dirty = register(&vault);
dirty.record(&kimun_core::NoteChange::Upsert {
path: VaultPath::new("precious.md"),
hash: 1,
});
assert!(!drain(&vault, &dirty, &transport).await.unwrap());
assert_eq!(dirty.len(), 1);
assert!(transport.deleted.lock().unwrap().is_empty());
}
}