1use std::sync::Arc;
9
10use axum::extract::{Path, Query, State};
11use axum::Json;
12use serde::{Deserialize, Serialize};
13
14use tuitbot_core::context::retrieval::{self, VaultCitation};
15use tuitbot_core::storage::watchtower;
16
17use crate::account::AccountContext;
18use crate::error::ApiError;
19use crate::state::AppState;
20
21const SNIPPET_MAX_LEN: usize = 120;
23
24const DEFAULT_LIMIT: u32 = 20;
26
27const MAX_LIMIT: u32 = 100;
29
30fn clamp_limit(limit: Option<u32>) -> u32 {
31 limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT)
32}
33
34fn truncate_snippet(text: &str, max_len: usize) -> String {
35 if text.len() <= max_len {
36 text.to_string()
37 } else {
38 let mut end = max_len.saturating_sub(3);
39 while end > 0 && !text.is_char_boundary(end) {
40 end -= 1;
41 }
42 format!("{}...", &text[..end])
43 }
44}
45
46#[derive(Serialize)]
51pub struct VaultSourcesResponse {
52 pub sources: Vec<VaultSourceStatusItem>,
53}
54
55#[derive(Serialize)]
56pub struct VaultSourceStatusItem {
57 pub id: i64,
58 pub source_type: String,
59 pub status: String,
60 pub error_message: Option<String>,
61 pub node_count: i64,
62 pub updated_at: String,
63 #[serde(skip_serializing_if = "Option::is_none")]
66 pub path: Option<String>,
67}
68
69pub async fn vault_sources(
70 State(state): State<Arc<AppState>>,
71 ctx: AccountContext,
72) -> Result<Json<VaultSourcesResponse>, ApiError> {
73 let sources = watchtower::get_all_source_contexts_for(&state.db, &ctx.account_id).await?;
74
75 let mut items = Vec::with_capacity(sources.len());
76 for src in sources {
77 let count = watchtower::count_nodes_for_source(&state.db, &ctx.account_id, src.id)
78 .await
79 .unwrap_or(0);
80 let path = if src.source_type == "local_fs" {
81 serde_json::from_str::<serde_json::Value>(&src.config_json)
82 .ok()
83 .and_then(|v| v.get("path").and_then(|p| p.as_str().map(String::from)))
84 } else {
85 None
86 };
87 items.push(VaultSourceStatusItem {
88 id: src.id,
89 source_type: src.source_type,
90 status: src.status,
91 error_message: src.error_message,
92 node_count: count,
93 updated_at: src.updated_at,
94 path,
95 });
96 }
97
98 Ok(Json(VaultSourcesResponse { sources: items }))
99}
100
101#[derive(Deserialize)]
106pub struct SearchNotesQuery {
107 pub q: Option<String>,
108 pub source_id: Option<i64>,
109 pub limit: Option<u32>,
110}
111
112#[derive(Serialize)]
113pub struct SearchNotesResponse {
114 pub notes: Vec<VaultNoteItem>,
115}
116
117#[derive(Serialize)]
118pub struct VaultNoteItem {
119 pub node_id: i64,
120 pub source_id: i64,
121 pub title: Option<String>,
122 pub relative_path: String,
123 pub tags: Option<String>,
124 pub status: String,
125 pub chunk_count: i64,
126 pub updated_at: String,
127}
128
129pub async fn search_notes(
130 State(state): State<Arc<AppState>>,
131 ctx: AccountContext,
132 Query(params): Query<SearchNotesQuery>,
133) -> Result<Json<SearchNotesResponse>, ApiError> {
134 let limit = clamp_limit(params.limit);
135
136 let nodes = match (¶ms.q, params.source_id) {
137 (Some(q), _) if !q.is_empty() => {
138 watchtower::search_nodes_for(&state.db, &ctx.account_id, q, limit).await?
139 }
140 (_, Some(sid)) => {
141 watchtower::get_nodes_for_source_for(&state.db, &ctx.account_id, sid, limit).await?
142 }
143 _ => {
144 watchtower::search_nodes_for(&state.db, &ctx.account_id, "", limit).await?
146 }
147 };
148
149 let mut notes = Vec::with_capacity(nodes.len());
150 for node in nodes {
151 let chunk_count =
152 watchtower::count_chunks_for_node(&state.db, &ctx.account_id, node.id).await?;
153 notes.push(VaultNoteItem {
154 node_id: node.id,
155 source_id: node.source_id,
156 title: node.title,
157 relative_path: node.relative_path,
158 tags: node.tags,
159 status: node.status,
160 chunk_count,
161 updated_at: node.updated_at,
162 });
163 }
164
165 Ok(Json(SearchNotesResponse { notes }))
166}
167
168#[derive(Serialize)]
173pub struct VaultNoteDetail {
174 pub node_id: i64,
175 pub source_id: i64,
176 pub title: Option<String>,
177 pub relative_path: String,
178 pub tags: Option<String>,
179 pub status: String,
180 pub ingested_at: String,
181 pub updated_at: String,
182 pub chunks: Vec<VaultChunkSummary>,
183}
184
185#[derive(Serialize)]
186pub struct VaultChunkSummary {
187 pub chunk_id: i64,
188 pub heading_path: String,
189 pub snippet: String,
190 pub retrieval_boost: f64,
191}
192
193pub async fn note_detail(
194 State(state): State<Arc<AppState>>,
195 ctx: AccountContext,
196 Path(id): Path<i64>,
197) -> Result<Json<VaultNoteDetail>, ApiError> {
198 let node = watchtower::get_content_node_for(&state.db, &ctx.account_id, id)
199 .await?
200 .ok_or_else(|| ApiError::NotFound(format!("note {id} not found")))?;
201
202 let chunks = watchtower::get_chunks_for_node(&state.db, &ctx.account_id, id).await?;
203
204 let chunk_summaries: Vec<VaultChunkSummary> = chunks
205 .into_iter()
206 .map(|c| VaultChunkSummary {
207 chunk_id: c.id,
208 heading_path: c.heading_path,
209 snippet: truncate_snippet(&c.chunk_text, SNIPPET_MAX_LEN),
210 retrieval_boost: c.retrieval_boost,
211 })
212 .collect();
213
214 Ok(Json(VaultNoteDetail {
215 node_id: node.id,
216 source_id: node.source_id,
217 title: node.title,
218 relative_path: node.relative_path,
219 tags: node.tags,
220 status: node.status,
221 ingested_at: node.ingested_at,
222 updated_at: node.updated_at,
223 chunks: chunk_summaries,
224 }))
225}
226
227#[derive(Deserialize)]
232pub struct SearchFragmentsQuery {
233 pub q: String,
234 pub limit: Option<u32>,
235}
236
237#[derive(Serialize)]
238pub struct SearchFragmentsResponse {
239 pub fragments: Vec<VaultCitation>,
240}
241
242pub async fn search_fragments(
243 State(state): State<Arc<AppState>>,
244 ctx: AccountContext,
245 Query(params): Query<SearchFragmentsQuery>,
246) -> Result<Json<SearchFragmentsResponse>, ApiError> {
247 let limit = clamp_limit(params.limit);
248
249 if params.q.is_empty() {
250 return Ok(Json(SearchFragmentsResponse { fragments: vec![] }));
251 }
252
253 let keywords: Vec<String> = params.q.split_whitespace().map(|s| s.to_string()).collect();
254
255 let fragments =
256 retrieval::retrieve_vault_fragments(&state.db, &ctx.account_id, &keywords, None, limit)
257 .await?;
258
259 let citations = retrieval::build_citations(&fragments);
260
261 Ok(Json(SearchFragmentsResponse {
262 fragments: citations,
263 }))
264}
265
266#[derive(Deserialize)]
271pub struct ResolveRefsRequest {
272 pub node_ids: Vec<i64>,
273}
274
275#[derive(Serialize)]
276pub struct ResolveRefsResponse {
277 pub citations: Vec<VaultCitation>,
278}
279
280pub async fn resolve_refs(
281 State(state): State<Arc<AppState>>,
282 ctx: AccountContext,
283 Json(body): Json<ResolveRefsRequest>,
284) -> Result<Json<ResolveRefsResponse>, ApiError> {
285 if body.node_ids.is_empty() {
286 return Ok(Json(ResolveRefsResponse { citations: vec![] }));
287 }
288
289 let fragments = retrieval::retrieve_vault_fragments(
290 &state.db,
291 &ctx.account_id,
292 &[],
293 Some(&body.node_ids),
294 MAX_LIMIT,
295 )
296 .await?;
297
298 let citations = retrieval::build_citations(&fragments);
299
300 Ok(Json(ResolveRefsResponse { citations }))
301}
302
303#[cfg(test)]
308mod tests {
309 use super::*;
310
311 use std::collections::HashMap;
312 use std::path::PathBuf;
313
314 use axum::body::Body;
315 use axum::http::{Request, StatusCode};
316 use axum::routing::{get, post};
317 use axum::Router;
318 use tokio::sync::{broadcast, Mutex, RwLock};
319 use tower::ServiceExt;
320
321 use crate::ws::AccountWsEvent;
322
323 async fn test_state() -> Arc<AppState> {
324 let db = tuitbot_core::storage::init_test_db()
325 .await
326 .expect("init test db");
327 let (event_tx, _) = broadcast::channel::<AccountWsEvent>(16);
328 Arc::new(AppState {
329 db,
330 config_path: PathBuf::from("/tmp/test-config.toml"),
331 data_dir: PathBuf::from("/tmp"),
332 event_tx,
333 api_token: "test-token".to_string(),
334 passphrase_hash: RwLock::new(None),
335 passphrase_hash_mtime: RwLock::new(None),
336 bind_host: "127.0.0.1".to_string(),
337 bind_port: 3001,
338 login_attempts: Mutex::new(HashMap::new()),
339 runtimes: Mutex::new(HashMap::new()),
340 content_generators: Mutex::new(HashMap::new()),
341 circuit_breaker: None,
342 watchtower_cancel: RwLock::new(None),
343 content_sources: RwLock::new(Default::default()),
344 connector_config: Default::default(),
345 deployment_mode: Default::default(),
346 pending_oauth: Mutex::new(HashMap::new()),
347 token_managers: Mutex::new(HashMap::new()),
348 x_client_id: String::new(),
349 })
350 }
351
352 fn test_router(state: Arc<AppState>) -> Router {
353 Router::new()
354 .route("/vault/sources", get(vault_sources))
355 .route("/vault/notes", get(search_notes))
356 .route("/vault/notes/{id}", get(note_detail))
357 .route("/vault/search", get(search_fragments))
358 .route("/vault/resolve-refs", post(resolve_refs))
359 .with_state(state)
360 }
361
362 #[tokio::test]
363 async fn vault_sources_returns_empty_when_no_sources() {
364 let state = test_state().await;
365 let app = test_router(state);
366
367 let resp = app
368 .oneshot(
369 Request::builder()
370 .uri("/vault/sources")
371 .body(Body::empty())
372 .unwrap(),
373 )
374 .await
375 .unwrap();
376
377 assert_eq!(resp.status(), StatusCode::OK);
378 let body: serde_json::Value = serde_json::from_slice(
379 &axum::body::to_bytes(resp.into_body(), 1024 * 64)
380 .await
381 .unwrap(),
382 )
383 .unwrap();
384 assert_eq!(body["sources"].as_array().unwrap().len(), 0);
385 }
386
387 #[tokio::test]
388 async fn search_notes_returns_empty_for_no_matches() {
389 let state = test_state().await;
390 let app = test_router(state);
391
392 let resp = app
393 .oneshot(
394 Request::builder()
395 .uri("/vault/notes?q=nonexistent")
396 .body(Body::empty())
397 .unwrap(),
398 )
399 .await
400 .unwrap();
401
402 assert_eq!(resp.status(), StatusCode::OK);
403 let body: serde_json::Value = serde_json::from_slice(
404 &axum::body::to_bytes(resp.into_body(), 1024 * 64)
405 .await
406 .unwrap(),
407 )
408 .unwrap();
409 assert_eq!(body["notes"].as_array().unwrap().len(), 0);
410 }
411
412 #[tokio::test]
413 async fn note_detail_returns_404_for_missing_node() {
414 let state = test_state().await;
415 let app = test_router(state);
416
417 let resp = app
418 .oneshot(
419 Request::builder()
420 .uri("/vault/notes/999")
421 .body(Body::empty())
422 .unwrap(),
423 )
424 .await
425 .unwrap();
426
427 assert_eq!(resp.status(), StatusCode::NOT_FOUND);
428 }
429
430 #[tokio::test]
431 async fn search_fragments_returns_empty_for_no_chunks() {
432 let state = test_state().await;
433 let app = test_router(state);
434
435 let resp = app
436 .oneshot(
437 Request::builder()
438 .uri("/vault/search?q=nonexistent")
439 .body(Body::empty())
440 .unwrap(),
441 )
442 .await
443 .unwrap();
444
445 assert_eq!(resp.status(), StatusCode::OK);
446 let body: serde_json::Value = serde_json::from_slice(
447 &axum::body::to_bytes(resp.into_body(), 1024 * 64)
448 .await
449 .unwrap(),
450 )
451 .unwrap();
452 assert_eq!(body["fragments"].as_array().unwrap().len(), 0);
453 }
454
455 #[tokio::test]
456 async fn resolve_refs_returns_empty_for_empty_ids() {
457 let state = test_state().await;
458 let app = test_router(state);
459
460 let resp = app
461 .oneshot(
462 Request::builder()
463 .method("POST")
464 .uri("/vault/resolve-refs")
465 .header("content-type", "application/json")
466 .body(Body::from(r#"{"node_ids":[]}"#))
467 .unwrap(),
468 )
469 .await
470 .unwrap();
471
472 assert_eq!(resp.status(), StatusCode::OK);
473 let body: serde_json::Value = serde_json::from_slice(
474 &axum::body::to_bytes(resp.into_body(), 1024 * 64)
475 .await
476 .unwrap(),
477 )
478 .unwrap();
479 assert_eq!(body["citations"].as_array().unwrap().len(), 0);
480 }
481}