1use std::collections::HashMap;
8use std::sync::Arc;
9
10use kimun_core::{IndexObserver, NoteVault, error::VaultError, nfs::VaultPath};
11
12use crate::server_client::dto::{WireDoc, WireSection};
13use crate::server_client::{
14 DirtyOp, DirtySet, RagClient, RagError, RagObserver, RagTransport, hash_string, reconcile_diff,
15};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ServerCapability {
21 Unconfigured,
24 SemanticOnly,
26 Full,
28}
29
30impl ServerCapability {
31 pub fn from_health(health: &crate::server_client::dto::Health) -> Self {
33 match (health.embedder.is_some(), health.llm_provider.is_some()) {
34 (false, _) => ServerCapability::Unconfigured,
35 (true, false) => ServerCapability::SemanticOnly,
36 (true, true) => ServerCapability::Full,
37 }
38 }
39
40 pub fn llm_available(self) -> bool {
42 matches!(self, ServerCapability::Full)
43 }
44}
45
46#[derive(Debug, Clone, PartialEq, Eq)]
48pub enum ServerUpdate {
49 Newer(String),
51 Legacy,
54}
55
56impl ServerUpdate {
57 pub fn summary(&self) -> String {
59 match self {
60 ServerUpdate::Newer(v) => format!("server {v} available"),
61 ServerUpdate::Legacy => "server update available".to_string(),
62 }
63 }
64
65 pub fn from_health(health: &crate::server_client::dto::Health) -> Option<Self> {
69 match &health.version {
70 None => Some(ServerUpdate::Legacy),
71 Some(_) => health.latest_version.clone().map(ServerUpdate::Newer),
72 }
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct ServerProbe {
84 pub capability: ServerCapability,
85 pub auth_required: bool,
86 pub server_update: Option<ServerUpdate>,
87}
88
89pub struct RagSync {
94 vault: Arc<NoteVault>,
95 dirty: Arc<DirtySet>,
96 observer: Arc<dyn IndexObserver>,
99 client: RagClient,
100}
101
102impl RagSync {
103 pub fn new(vault: Arc<NoteVault>, client: RagClient) -> Self {
109 let dirty = Arc::new(DirtySet::default());
110 let observer: Arc<dyn IndexObserver> = Arc::new(RagObserver::new(dirty.clone()));
111 vault.set_index_observer(observer.clone());
112 Self {
113 vault,
114 dirty,
115 observer,
116 client,
117 }
118 }
119
120 pub async fn probe(&self) -> Option<ServerProbe> {
126 self.client.health().await.ok().map(|h| ServerProbe {
127 capability: ServerCapability::from_health(&h),
128 auth_required: h.auth_required,
129 server_update: ServerUpdate::from_health(&h),
130 })
131 }
132
133 pub fn index_ready(&self) -> bool {
137 self.vault.index_ready()
138 }
139
140 pub async fn tick(&self) -> Result<bool, RagError> {
144 let drained = drain(&self.vault, &self.dirty, &self.client).await?;
145 let reconciled = reconcile(&self.vault, &self.client).await?;
146 Ok(drained && reconciled)
147 }
148
149 pub async fn drain(&self) -> Result<bool, RagError> {
154 drain(&self.vault, &self.dirty, &self.client).await
155 }
156
157 pub async fn reconcile(&self) -> Result<bool, RagError> {
161 reconcile(&self.vault, &self.client).await
162 }
163
164 pub fn client(&self) -> &RagClient {
166 &self.client
167 }
168}
169
170impl Drop for RagSync {
171 fn drop(&mut self) {
172 self.vault.clear_index_observer_if(&self.observer);
176 }
177}
178
179pub async fn build_doc(
185 vault: &NoteVault,
186 path: &VaultPath,
187 hash: u64,
188) -> Result<Option<WireDoc>, VaultError> {
189 let chunks = vault.get_note_chunks(path).await?;
190 let sections: Vec<WireSection> = chunks
191 .into_values()
192 .flatten()
193 .map(|c| WireSection {
194 title: c.get_breadcrumb().to_string(),
195 text: c.get_text().to_string(),
196 })
197 .collect();
198 if sections.is_empty() {
199 return Ok(None);
200 }
201 Ok(Some(WireDoc {
202 path: path.to_string(),
203 hash: hash_string(hash),
204 sections,
205 }))
206}
207
208pub async fn drain<T: RagTransport>(
216 vault: &NoteVault,
217 dirty: &DirtySet,
218 transport: &T,
219) -> Result<bool, RagError> {
220 if !vault.index_ready() {
221 return Ok(false);
222 }
223 let ops = dirty.drain();
224 if ops.is_empty() {
225 return Ok(true);
226 }
227
228 let mut upserts: Vec<(VaultPath, u64)> = Vec::new();
229 let mut deletes: Vec<String> = Vec::new();
230 for (path, op) in ops {
231 match op {
232 DirtyOp::Upsert(hash) => upserts.push((path, hash)),
233 DirtyOp::Delete => deletes.push(path.to_string()),
234 }
235 }
236
237 let mut docs = Vec::new();
240 let mut built: Vec<(VaultPath, u64)> = Vec::new();
241 for (path, hash) in upserts {
242 match build_doc(vault, &path, hash).await {
243 Ok(Some(doc)) => {
244 docs.push(doc);
245 built.push((path, hash));
246 }
247 Ok(None) => deletes.push(path.to_string()),
250 Err(_) => dirty.requeue([(path, DirtyOp::Upsert(hash))]),
251 }
252 }
253
254 let mut first_err: Option<RagError> = None;
255 if !docs.is_empty()
256 && let Err(e) = transport.push_docs(docs).await
257 {
258 dirty.requeue(built.into_iter().map(|(p, h)| (p, DirtyOp::Upsert(h))));
259 first_err = Some(e);
260 }
261 if !deletes.is_empty() {
262 let paths_for_requeue: Vec<VaultPath> = deletes.iter().map(VaultPath::new).collect();
263 if let Err(e) = transport.delete_paths(deletes).await {
264 dirty.requeue(paths_for_requeue.into_iter().map(|p| (p, DirtyOp::Delete)));
265 first_err = first_err.or(Some(e));
266 }
267 }
268
269 match first_err {
270 Some(e) => Err(e),
271 None => Ok(true),
272 }
273}
274
275pub async fn reconcile<T: RagTransport>(
283 vault: &NoteVault,
284 transport: &T,
285) -> Result<bool, RagError> {
286 if !vault.index_ready() {
287 return Ok(false);
288 }
289 let notes = vault
290 .get_all_notes()
291 .await
292 .map_err(|e| RagError::Protocol(format!("read vault notes: {e}")))?;
293
294 let local_hashes: HashMap<String, u64> = notes
295 .into_iter()
296 .map(|(entry, content)| (entry.path.to_string(), content.hash))
297 .collect();
298 let local_str: HashMap<String, String> = local_hashes
299 .iter()
300 .map(|(p, h)| (p.clone(), hash_string(*h)))
301 .collect();
302
303 let server = transport.server_hashes().await?;
304 let plan = reconcile_diff(&local_str, &server);
305
306 let mut docs = Vec::new();
307 let mut to_delete = plan.to_delete;
308 for path_str in &plan.to_push {
309 let hash = local_hashes[path_str];
310 match build_doc(vault, &VaultPath::new(path_str), hash)
311 .await
312 .map_err(|e| RagError::Protocol(format!("build doc {path_str}: {e}")))?
313 {
314 Some(doc) => docs.push(doc),
315 None => {
319 if server.contains_key(path_str) {
320 to_delete.push(path_str.clone());
321 }
322 }
323 }
324 }
325 if !docs.is_empty() {
326 transport.push_docs(docs).await?;
327 }
328 if !to_delete.is_empty() {
329 transport.delete_paths(to_delete).await?;
330 }
331 Ok(true)
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337 use async_trait::async_trait;
338
339 #[test]
340 fn server_update_from_health_fields() {
341 use crate::server_client::dto::Health;
342 let h = |version: Option<&str>, latest: Option<&str>| Health {
343 status: "ok".into(),
344 reranker: false,
345 embedder: None,
346 llm_provider: None,
347 auth_required: false,
348 version: version.map(str::to_string),
349 latest_version: latest.map(str::to_string),
350 };
351 assert_eq!(ServerUpdate::from_health(&h(Some("0.4.4"), None)), None);
352 assert_eq!(
353 ServerUpdate::from_health(&h(Some("0.4.4"), Some("0.4.5"))),
354 Some(ServerUpdate::Newer("0.4.5".into()))
355 );
356 let legacy: Health = serde_json::from_str(
358 r#"{"status":"ok","reranker":false,"embedder":"fastembed","llm_provider":null,"auth_required":false,"degraded":null}"#,
359 )
360 .unwrap();
361 assert_eq!(
362 ServerUpdate::from_health(&legacy),
363 Some(ServerUpdate::Legacy)
364 );
365 }
366
367 #[test]
368 fn capability_from_health_fields() {
369 use crate::server_client::dto::Health;
370 let h = |embedder: Option<&str>, llm: Option<&str>| Health {
371 status: "ok".into(),
372 reranker: false,
373 embedder: embedder.map(str::to_string),
374 llm_provider: llm.map(str::to_string),
375 auth_required: false,
376 version: Some("0.4.4".into()),
377 latest_version: None,
378 };
379 assert_eq!(
380 ServerCapability::from_health(&h(None, None)),
381 ServerCapability::Unconfigured
382 );
383 assert_eq!(
384 ServerCapability::from_health(&h(None, Some("gemini"))),
386 ServerCapability::Unconfigured
387 );
388 assert_eq!(
389 ServerCapability::from_health(&h(Some("fastembed"), None)),
390 ServerCapability::SemanticOnly
391 );
392 assert_eq!(
393 ServerCapability::from_health(&h(Some("fastembed"), Some("gemini"))),
394 ServerCapability::Full
395 );
396 }
397 use kimun_core::VaultConfig;
398 use std::sync::Mutex;
399 use tempfile::TempDir;
400
401 #[derive(Default)]
402 struct FakeTransport {
403 pushed: Mutex<Vec<WireDoc>>,
404 deleted: Mutex<Vec<String>>,
405 server: Mutex<HashMap<String, String>>,
406 fail_push: Mutex<bool>,
407 }
408
409 #[async_trait]
410 impl RagTransport for FakeTransport {
411 async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError> {
412 if *self.fail_push.lock().unwrap() {
413 return Err(RagError::Protocol("boom".into()));
414 }
415 self.pushed.lock().unwrap().extend(docs);
416 Ok(())
417 }
418 async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
419 self.deleted.lock().unwrap().extend(paths);
420 Ok(())
421 }
422 async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
423 Ok(self.server.lock().unwrap().clone())
424 }
425 }
426
427 fn register(vault: &NoteVault) -> Arc<DirtySet> {
432 let dirty = Arc::new(DirtySet::default());
433 vault.set_index_observer(Arc::new(RagObserver::new(dirty.clone())));
434 dirty
435 }
436
437 fn sys(path: impl AsRef<std::path::Path>) -> kimun_core::SystemPath {
441 kimun_core::SystemPath::try_absolute(path).expect("test path must be absolute")
442 }
443
444 async fn vault(dir: &std::path::Path) -> NoteVault {
445 let vault = NoteVault::new(VaultConfig::new(sys(dir))).await.unwrap();
446 vault.validate_and_init().await.unwrap();
449 vault
450 }
451
452 #[tokio::test]
453 async fn drain_pushes_created_note_and_deletes_removed() {
454 let dir = TempDir::new().unwrap();
455 let vault = vault(dir.path()).await;
456 let dirty = register(&vault);
457 let transport = FakeTransport::default();
458
459 vault
460 .create_note(&VaultPath::new("a.md"), "# Title\n\nbody")
461 .await
462 .unwrap();
463 drain(&vault, &dirty, &transport).await.unwrap();
464
465 {
468 let pushed = transport.pushed.lock().unwrap();
469 assert_eq!(pushed.len(), 1);
470 assert_eq!(pushed[0].path, "/a.md"); assert!(!pushed[0].sections.is_empty());
472 assert!(dirty.is_empty());
473 }
474
475 vault.delete_note(&VaultPath::new("a.md")).await.unwrap();
476 drain(&vault, &dirty, &transport).await.unwrap();
477 assert_eq!(
478 *transport.deleted.lock().unwrap(),
479 vec!["/a.md".to_string()]
480 );
481 }
482
483 #[tokio::test]
484 async fn failed_push_requeues() {
485 let dir = TempDir::new().unwrap();
486 let vault = vault(dir.path()).await;
487 let dirty = register(&vault);
488 let transport = FakeTransport::default();
489 *transport.fail_push.lock().unwrap() = true;
490
491 vault
492 .create_note(&VaultPath::new("a.md"), "body")
493 .await
494 .unwrap();
495 assert!(drain(&vault, &dirty, &transport).await.is_err());
496 assert_eq!(dirty.len(), 1);
498 }
499
500 #[tokio::test]
501 async fn reconcile_pushes_missing_and_deletes_stale() {
502 let dir = TempDir::new().unwrap();
503 let vault = vault(dir.path()).await;
504 let _dirty = register(&vault);
505 let transport = FakeTransport::default();
506
507 vault
508 .create_note(&VaultPath::new("keep.md"), "kept")
509 .await
510 .unwrap();
511 transport
513 .server
514 .lock()
515 .unwrap()
516 .insert("/gone.md".to_string(), "oldhash".to_string());
517
518 assert!(reconcile(&vault, &transport).await.unwrap());
519
520 let pushed = transport.pushed.lock().unwrap();
521 assert!(pushed.iter().any(|d| d.path == "/keep.md"));
522 assert_eq!(
523 *transport.deleted.lock().unwrap(),
524 vec!["/gone.md".to_string()]
525 );
526 }
527
528 #[tokio::test]
529 async fn reconcile_skipped_while_index_not_ready() {
530 let dir = TempDir::new().unwrap();
531 let vault = NoteVault::new(VaultConfig::new(sys(dir.path())))
535 .await
536 .unwrap();
537 assert!(!vault.index_ready());
538 let transport = FakeTransport::default();
539 transport
540 .server
541 .lock()
542 .unwrap()
543 .insert("/precious.md".to_string(), "hash".to_string());
544
545 assert!(!reconcile(&vault, &transport).await.unwrap());
546 assert!(transport.deleted.lock().unwrap().is_empty());
547 assert!(transport.pushed.lock().unwrap().is_empty());
548
549 let dirty = register(&vault);
553 dirty.record(&kimun_core::NoteChange::Upsert {
554 path: VaultPath::new("precious.md"),
555 hash: 1,
556 });
557 assert!(!drain(&vault, &dirty, &transport).await.unwrap());
558 assert_eq!(dirty.len(), 1);
559 assert!(transport.deleted.lock().unwrap().is_empty());
560 }
561}