use std::sync::Arc;
use std::time::Duration;
use crate::server_client::{
RagClient,
sync::{RagSync, ServerCapability, ServerProbe},
};
use kimun_core::NoteVault;
use tokio::task::JoinHandle;
use super::RagStatus;
use super::client::server_config;
use crate::components::events::{AppEvent, AppTx};
use crate::settings::SharedSettings;
const SYNC_INTERVAL: Duration = Duration::from_secs(10);
const RECONCILE_EVERY_N_TICKS: u32 = 30;
#[derive(Debug, PartialEq, Eq)]
enum Plan {
Skip(RagStatus),
Wait { flash: Option<RagStatus> },
Run {
flash: Option<RagStatus>,
reconcile: bool,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Outcome {
Synced,
SkippedRebuild,
AuthRejected,
Failed,
}
struct Cadence {
ticks_since_reconcile: u32,
auth_failed: bool,
llm_available: bool,
}
impl Cadence {
fn new() -> Self {
Self {
ticks_since_reconcile: RECONCILE_EVERY_N_TICKS,
auth_failed: false,
llm_available: false,
}
}
fn force_reconcile(&mut self) {
self.ticks_since_reconcile = RECONCILE_EVERY_N_TICKS;
}
fn plan(&mut self, probe: Option<&ServerProbe>, has_token: bool, index_ready: bool) -> Plan {
let probe = match probe {
Some(p) => p,
None => {
self.force_reconcile();
self.auth_failed = false;
return Plan::Skip(RagStatus::Offline);
}
};
if probe.capability == ServerCapability::Unconfigured {
self.force_reconcile();
return Plan::Skip(RagStatus::NotConfigured);
}
if probe.auth_required && !has_token {
self.force_reconcile();
return Plan::Skip(RagStatus::Unauthorized);
}
let llm_available = probe.capability.llm_available();
self.llm_available = llm_available;
let flash = (!self.auth_failed).then_some(RagStatus::Syncing { llm_available });
if !index_ready {
self.force_reconcile();
return Plan::Wait { flash };
}
let reconcile = self.ticks_since_reconcile >= RECONCILE_EVERY_N_TICKS;
if reconcile {
self.ticks_since_reconcile = 0;
} else {
self.ticks_since_reconcile += 1;
}
Plan::Run { flash, reconcile }
}
fn settle(&mut self, outcome: Outcome) -> RagStatus {
let llm_available = self.llm_available;
match outcome {
Outcome::Synced => {
self.auth_failed = false;
RagStatus::Online { llm_available }
}
Outcome::SkippedRebuild => {
self.force_reconcile();
self.auth_failed = false;
RagStatus::Syncing { llm_available }
}
Outcome::AuthRejected => {
self.auth_failed = true;
RagStatus::Unauthorized
}
Outcome::Failed => {
self.auth_failed = false;
RagStatus::Offline
}
}
}
}
pub fn spawn_rag_sync(
vault: Arc<NoteVault>,
settings: &SharedSettings,
tx: AppTx,
) -> Option<JoinHandle<()>> {
let (url, token) = server_config(settings)?;
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(SYNC_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut sync: Option<RagSync> = None;
let mut cadence = Cadence::new();
loop {
interval.tick().await;
if sync.is_none() {
match vault.vault_id().await {
Ok(id) => {
let client = RagClient::new(url.clone(), token.clone(), id.to_string());
sync = Some(RagSync::new(vault.clone(), client));
}
Err(e) => {
log::warn!("RAG: cannot read vault id (will retry): {e}");
let _ = tx.send(AppEvent::RagStatus(RagStatus::Offline));
continue;
}
}
}
let sync = sync.as_ref().expect("sync established above");
let probe = sync.probe().await;
let reconcile = match cadence.plan(probe.as_ref(), token.is_some(), sync.index_ready())
{
Plan::Skip(status) => {
let _ = tx.send(AppEvent::RagStatus(status));
continue;
}
Plan::Wait { flash } => {
if let Some(status) = flash {
let _ = tx.send(AppEvent::RagStatus(status));
}
continue;
}
Plan::Run { flash, reconcile } => {
if let Some(status) = flash {
let _ = tx.send(AppEvent::RagStatus(status));
}
reconcile
}
};
let result = if reconcile {
sync.tick().await } else {
sync.drain().await };
let outcome = match &result {
Ok(true) => Outcome::Synced,
Ok(false) => Outcome::SkippedRebuild,
Err(e) if e.is_auth() => {
log::warn!("RAG server rejected the configured token: {e}");
Outcome::AuthRejected
}
Err(e) => {
log::debug!("RAG sync failed: {e}");
Outcome::Failed
}
};
let status = cadence.settle(outcome);
let _ = tx.send(AppEvent::RagStatus(status));
}
}))
}
#[cfg(test)]
mod tests {
use super::*;
fn probe(capability: ServerCapability, auth_required: bool) -> ServerProbe {
ServerProbe {
capability,
auth_required,
}
}
#[test]
fn offline_probe_reports_offline_and_forces_reconcile() {
let mut c = Cadence::new();
c.plan(Some(&probe(ServerCapability::Full, false)), true, true);
assert_eq!(c.plan(None, true, true), Plan::Skip(RagStatus::Offline));
assert_eq!(
c.plan(Some(&probe(ServerCapability::Full, false)), true, true),
Plan::Run {
flash: Some(RagStatus::Syncing {
llm_available: true
}),
reconcile: true,
}
);
}
#[test]
fn unconfigured_server_skips_sync() {
let mut c = Cadence::new();
assert_eq!(
c.plan(
Some(&probe(ServerCapability::Unconfigured, false)),
true,
true
),
Plan::Skip(RagStatus::NotConfigured)
);
}
#[test]
fn auth_required_without_token_reports_unauthorized_up_front() {
let mut c = Cadence::new();
assert_eq!(
c.plan(Some(&probe(ServerCapability::Full, true)), false, true),
Plan::Skip(RagStatus::Unauthorized)
);
}
#[test]
fn auth_required_with_token_syncs() {
let mut c = Cadence::new();
assert!(matches!(
c.plan(Some(&probe(ServerCapability::Full, true)), true, true),
Plan::Run { .. }
));
}
#[test]
fn index_not_ready_waits_and_reconciles_once_filled() {
let mut c = Cadence::new();
c.plan(Some(&probe(ServerCapability::Full, false)), true, true);
c.plan(Some(&probe(ServerCapability::Full, false)), true, true);
assert_eq!(
c.plan(Some(&probe(ServerCapability::Full, false)), true, false),
Plan::Wait {
flash: Some(RagStatus::Syncing {
llm_available: true
})
}
);
assert_eq!(
c.plan(Some(&probe(ServerCapability::Full, false)), true, true),
Plan::Run {
flash: Some(RagStatus::Syncing {
llm_available: true
}),
reconcile: true,
}
);
}
#[test]
fn reconcile_cadence_first_tick_then_drains_then_reconciles_again() {
let mut c = Cadence::new();
let p = probe(ServerCapability::Full, false);
assert!(matches!(
c.plan(Some(&p), true, true),
Plan::Run {
reconcile: true,
..
}
));
for _ in 0..RECONCILE_EVERY_N_TICKS {
assert!(matches!(
c.plan(Some(&p), true, true),
Plan::Run {
reconcile: false,
..
}
));
}
assert!(matches!(
c.plan(Some(&p), true, true),
Plan::Run {
reconcile: true,
..
}
));
}
#[test]
fn auth_rejection_is_sticky_and_suppresses_the_syncing_flash() {
let mut c = Cadence::new();
let p = probe(ServerCapability::Full, true);
assert!(matches!(c.plan(Some(&p), true, true), Plan::Run { .. }));
assert_eq!(c.settle(Outcome::AuthRejected), RagStatus::Unauthorized);
assert_eq!(
c.plan(Some(&p), true, true),
Plan::Run {
flash: None,
reconcile: false,
}
);
assert_eq!(
c.settle(Outcome::Synced),
RagStatus::Online {
llm_available: true
}
);
assert!(matches!(
c.plan(Some(&p), true, true),
Plan::Run { flash: Some(_), .. }
));
}
#[test]
fn skipped_pass_reports_syncing_and_forces_reconcile() {
let mut c = Cadence::new();
let p = probe(ServerCapability::SemanticOnly, false);
c.plan(Some(&p), true, true);
assert_eq!(
c.settle(Outcome::SkippedRebuild),
RagStatus::Syncing {
llm_available: false
}
);
assert!(matches!(
c.plan(Some(&p), true, true),
Plan::Run {
reconcile: true,
..
}
));
}
#[test]
fn sync_failure_reports_offline() {
let mut c = Cadence::new();
let p = probe(ServerCapability::Full, false);
c.plan(Some(&p), true, true);
assert_eq!(c.settle(Outcome::Failed), RagStatus::Offline);
}
}