Skip to main content

kimun_notes/server_client/
mod.rs

1//! The component inside Kimün that owns every dealing with the RAG server:
2//! connection/capability probing, pushing note changes, and hash-diff
3//! reconciliation (see CONTEXT.md, "Server client"). Core stays network-free;
4//! it feeds this module only through the [`observer`] seam, which reports a
5//! path, a content hash and upsert-or-delete, and knows nothing of RAG.
6//!
7//! Published as the `kimun_server_client` crate until kimun-notes went to
8//! crates.io — which forced every path dependency of a published crate to be
9//! published too, and this one was on crates.io for no reason of its own. Like
10//! the `ropetext` module it stays self-contained, so it can go back out if a
11//! second consumer (a GUI) ever wants it: nothing here may name `crate::`
12//! outside `crate::server_client::`, checked in CI. `kimun_core` is the one
13//! kimün dependency it keeps, and it would keep that as a crate too.
14
15use std::collections::HashMap;
16use std::time::Duration;
17
18pub mod dto;
19pub mod observer;
20pub mod reconcile;
21
22pub mod sync;
23
24use async_trait::async_trait;
25use dto::{
26    AnswerResult, DeleteRequest, EmbeddingsResponse, Health, HistoryTurn, IndexDocsRequest,
27    JobAccepted, JobStatus, QueryRequest, WireDoc,
28};
29
30pub use dto::{ChunkResult, WireSection};
31pub use observer::{DirtyOp, DirtySet, RagObserver};
32pub use reconcile::{ReconcilePlan, diff as reconcile_diff};
33
34#[derive(Debug, thiserror::Error)]
35pub enum RagError {
36    #[error("http error: {0}")]
37    Http(#[from] reqwest::Error),
38    #[error("server returned {status}: {body}")]
39    Status { status: u16, body: String },
40    /// A well-formed HTTP exchange that violated the expected protocol (job
41    /// failed/timed out, unparseable result).
42    #[error("{0}")]
43    Protocol(String),
44}
45
46impl RagError {
47    /// Whether this is an authentication/authorization rejection (401/403) —
48    /// a token problem, not an unreachable server. Callers should surface it
49    /// as such instead of folding it into a generic "offline".
50    pub fn is_auth(&self) -> bool {
51        matches!(
52            self,
53            RagError::Status {
54                status: 401 | 403,
55                ..
56            }
57        )
58    }
59}
60
61/// The subset of server operations the sync orchestration depends on, behind a
62/// trait so it can be exercised with a fake in tests. [`RagClient`] is the real
63/// implementation.
64#[async_trait]
65pub trait RagTransport: Send + Sync {
66    async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError>;
67    async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError>;
68    async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError>;
69}
70
71#[async_trait]
72impl RagTransport for RagClient {
73    async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<(), RagError> {
74        // The server indexes in the background; an accepted push is enough for
75        // the drain path (reconciliation catches any server-side failure).
76        RagClient::push_docs(self, docs).await.map(|_job_id| ())
77    }
78    async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
79        RagClient::delete_paths(self, paths).await
80    }
81    async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
82        RagClient::server_hashes(self).await
83    }
84}
85
86/// The single place a note's content hash is turned into its wire/reconcile
87/// string. Both the pushed [`WireDoc::hash`](dto::WireDoc) and the reconcile
88/// `local` hash set MUST go through this — string equality in
89/// [`reconcile::diff`] only holds if the two are byte-identical.
90pub fn hash_string(hash: u64) -> String {
91    hash.to_string()
92}
93
94/// Number of results to request. Maps to the server's `context_size` variants,
95/// so a caller can't send an invalid string (which the server rejects with 400).
96#[derive(Debug, Clone, Copy)]
97pub enum ContextSize {
98    Small,
99    Medium,
100    Large,
101}
102
103impl ContextSize {
104    fn as_str(self) -> &'static str {
105        match self {
106            ContextSize::Small => "small",
107            ContextSize::Medium => "medium",
108            ContextSize::Large => "large",
109        }
110    }
111}
112
113/// Bound on establishing a TCP/TLS connection. Without it, a black-holing
114/// host (firewalled port, sleeping machine) hangs a probe for the OS default
115/// (minutes) instead of failing over to "offline" promptly.
116const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
117
118/// Bound on a whole request/response exchange. reqwest has NO default here —
119/// a server that accepts the connection but never answers would hang forever.
120const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
121
122/// Looser bound for [`RagClient::push_docs`]: a first sync of a large vault
123/// uploads every document in one request, which can legitimately outlast
124/// [`REQUEST_TIMEOUT`] on a slow link.
125const PUSH_TIMEOUT: Duration = Duration::from_secs(120);
126
127/// One process-wide HTTP client, so every [`RagClient`] — however short-lived —
128/// shares the same connection pool and keep-alive connections instead of
129/// paying a fresh TCP+TLS handshake per construction.
130fn shared_http() -> reqwest::Client {
131    static HTTP: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
132    HTTP.get_or_init(|| {
133        reqwest::Client::builder()
134            .connect_timeout(CONNECT_TIMEOUT)
135            .timeout(REQUEST_TIMEOUT)
136            .build()
137            // Only fails on broken TLS backend/system config; no meaningful
138            // recovery, and it would fail identically for every request.
139            .expect("build HTTP client")
140    })
141    .clone()
142}
143
144/// HTTP client for one vault's collection on a RAG server.
145#[derive(Clone)]
146pub struct RagClient {
147    http: reqwest::Client,
148    base_url: String,
149    token: Option<String>,
150    vault_id: String,
151}
152
153impl RagClient {
154    /// `base_url` like `http://host:7573`; `token` is the bearer token if the
155    /// server requires one; `vault_id` selects this vault's collection.
156    pub fn new(
157        base_url: impl Into<String>,
158        token: Option<String>,
159        vault_id: impl Into<String>,
160    ) -> Self {
161        let base_url = base_url.into().trim_end_matches('/').to_string();
162        Self {
163            http: shared_http(),
164            base_url,
165            token,
166            vault_id: vault_id.into(),
167        }
168    }
169
170    fn url(&self, path: &str) -> String {
171        format!("{}{}", self.base_url, path)
172    }
173
174    /// Attaches the bearer token when configured.
175    fn auth(&self, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
176        match &self.token {
177            Some(token) => req.bearer_auth(token),
178            None => req,
179        }
180    }
181
182    /// Turns a non-2xx response into a [`RagError::Status`], else yields the
183    /// response for JSON decoding.
184    async fn ok(resp: reqwest::Response) -> Result<reqwest::Response, RagError> {
185        if resp.status().is_success() {
186            Ok(resp)
187        } else {
188            let status = resp.status().as_u16();
189            let body = resp.text().await.unwrap_or_default();
190            Err(RagError::Status { status, body })
191        }
192    }
193
194    /// Probes `GET /health` for reachability + capabilities.
195    pub async fn health(&self) -> Result<Health, RagError> {
196        let resp = self.auth(self.http.get(self.url("/health"))).send().await?;
197        Ok(Self::ok(resp).await?.json::<Health>().await?)
198    }
199
200    /// Pushes documents to this vault's collection; returns the server's job id.
201    pub async fn push_docs(&self, docs: Vec<WireDoc>) -> Result<String, RagError> {
202        let body = IndexDocsRequest {
203            vault_id: self.vault_id.clone(),
204            docs,
205        };
206        let resp = self
207            .auth(self.http.post(self.url("/api/index/docs")).json(&body))
208            .timeout(PUSH_TIMEOUT)
209            .send()
210            .await?;
211        Ok(Self::ok(resp).await?.json::<JobAccepted>().await?.job_id)
212    }
213
214    /// Deletes notes by path from this vault's collection.
215    pub async fn delete_paths(&self, paths: Vec<String>) -> Result<(), RagError> {
216        let body = DeleteRequest {
217            vault_id: self.vault_id.clone(),
218            paths,
219        };
220        let resp = self
221            .auth(self.http.post(self.url("/api/index/delete")).json(&body))
222            .send()
223            .await?;
224        Self::ok(resp).await?;
225        Ok(())
226    }
227
228    /// The server's `{note-path → hash}` set for this vault (reconcile input).
229    ///
230    /// `vault_id` is interpolated into the URL path un-encoded; this is safe
231    /// because it is always a UUID (from `.kimun/vault-id`) and thus
232    /// URL-safe. If that ever changes, percent-encode the segment here.
233    pub async fn server_hashes(&self) -> Result<HashMap<String, String>, RagError> {
234        let path = format!("/api/collections/{}/hashes", self.vault_id);
235        let resp = self.auth(self.http.get(self.url(&path))).send().await?;
236        Ok(Self::ok(resp)
237            .await?
238            .json::<HashMap<String, String>>()
239            .await?)
240    }
241
242    /// Semantic search: returns the matching chunks (no LLM). `context_size`
243    /// omitted uses the server's configured default.
244    pub async fn search(
245        &self,
246        query: &str,
247        context_size: Option<ContextSize>,
248    ) -> Result<Vec<ChunkResult>, RagError> {
249        let body = QueryRequest {
250            vault_id: self.vault_id.clone(),
251            query: query.to_string(),
252            context_size: context_size.map(|c| c.as_str().to_string()),
253            history: vec![],
254        };
255        let resp = self
256            .auth(self.http.post(self.url("/api/embeddings")).json(&body))
257            .send()
258            .await?;
259        Ok(Self::ok(resp)
260            .await?
261            .json::<EmbeddingsResponse>()
262            .await?
263            .chunks)
264    }
265
266    /// Submits a question, polls the job to completion, and returns the LLM
267    /// answer plus its cited source chunks.
268    pub async fn ask(
269        &self,
270        query: &str,
271        history: &[(String, String)],
272        context_size: Option<ContextSize>,
273    ) -> Result<AnswerResult, RagError> {
274        let body = QueryRequest {
275            vault_id: self.vault_id.clone(),
276            query: query.to_string(),
277            context_size: context_size.map(|c| c.as_str().to_string()),
278            history: history
279                .iter()
280                .map(|(q, a)| HistoryTurn {
281                    question: q.clone(),
282                    answer: a.clone(),
283                })
284                .collect(),
285        };
286        let resp = self
287            .auth(self.http.post(self.url("/api/answer")).json(&body))
288            .send()
289            .await?;
290        let job_id = Self::ok(resp).await?.json::<JobAccepted>().await?.job_id;
291        self.poll_answer(&job_id).await
292    }
293
294    /// Poll iterations before giving up on an answer job, at ~1s each: ~12
295    /// minutes, safely inside the server's 15-minute job retention (see
296    /// `server_state.rs`) so a slow LLM (large context, CPU-only local model
297    /// — legitimately several minutes) isn't cut off with a false timeout
298    /// while its result is still coming. Keep this below the retention window:
299    /// past it the job is swept and polling can never succeed.
300    const ANSWER_POLL_ATTEMPTS: u32 = 720;
301
302    /// Polls `/api/job/{id}` until the answer job completes or fails.
303    async fn poll_answer(&self, job_id: &str) -> Result<AnswerResult, RagError> {
304        let path = format!("/api/job/{job_id}");
305        let mut consecutive_errors = 0u32;
306        for _ in 0..Self::ANSWER_POLL_ATTEMPTS {
307            // A transient poll error (network blip, server briefly busy) must not
308            // abort an answer that is still being generated — retry, and only
309            // give up after a run of consecutive failures.
310            let status = match self.poll_once(&path).await {
311                Ok(s) => {
312                    consecutive_errors = 0;
313                    s
314                }
315                Err(e) => {
316                    consecutive_errors += 1;
317                    if consecutive_errors >= 15 {
318                        return Err(e);
319                    }
320                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
321                    continue;
322                }
323            };
324            match status.status.as_str() {
325                "completed" => {
326                    let result = status
327                        .result
328                        .ok_or_else(|| RagError::Protocol("completed job had no result".into()))?;
329                    return serde_json::from_value::<AnswerResult>(result)
330                        .map_err(|e| RagError::Protocol(format!("bad answer result: {e}")));
331                }
332                "failed" => {
333                    return Err(RagError::Protocol(
334                        status.error.unwrap_or_else(|| "answer job failed".into()),
335                    ));
336                }
337                _ => tokio::time::sleep(std::time::Duration::from_secs(1)).await,
338            }
339        }
340        Err(RagError::Protocol("answer job timed out".into()))
341    }
342
343    /// One poll of the job status endpoint.
344    async fn poll_once(&self, path: &str) -> Result<JobStatus, RagError> {
345        let resp = self.auth(self.http.get(self.url(path))).send().await?;
346        Ok(Self::ok(resp).await?.json::<JobStatus>().await?)
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn is_auth_matches_only_credential_rejections() {
356        let status = |status| RagError::Status {
357            status,
358            body: String::new(),
359        };
360        assert!(status(401).is_auth());
361        assert!(status(403).is_auth());
362        assert!(!status(500).is_auth());
363        assert!(!status(404).is_auth());
364        assert!(!RagError::Protocol("boom".into()).is_auth());
365    }
366}