feat(desktop): embed BoosAPI Desktop client as a subproject

Vendor the rebranded BoosAPI Desktop client (Tauri + Vue) into metapi as
a monorepo subdirectory under desktop/. Source only — node_modules,
dist, Cargo target and the upstream .git history are excluded.

desktop/ keeps its own independent Cargo workspace (desktop/src-tauri
with crates/*); it is not part of the backend boosapi crate, so a
`cargo check` at the metapi root still compiles only the backend. Build
the desktop app from desktop/src-tauri per desktop/README.md.

Note: app icons are still placeholder assets pending the final BoosAPI
logo. The updater signing keys live outside the repo and are not
included.
This commit is contained in:
boos
2026-06-22 11:32:57 +08:00
parent d9c78ecc73
commit e0a5f57793
2305 changed files with 449355 additions and 0 deletions
+41
View File
@@ -0,0 +1,41 @@
[package]
name = "boosapi-core"
version = "0.0.1"
edition = "2021"
[dependencies]
serde = { workspace = true }
serde_json = { workspace = true }
sea-orm = { workspace = true }
tokio = { workspace = true }
thiserror = { workspace = true }
uuid = { workspace = true }
chrono = { workspace = true }
tracing = { workspace = true }
anyhow = { workspace = true }
aes-gcm = "0.10"
async-trait = "0.1"
aws-config = { version = "1", features = ["behavior-version-latest"] }
aws-credential-types = "1"
aws-sdk-s3 = "1"
sha2 = "0.10"
base64 = "0.22"
reqwest = { version = "0.13", features = ["json", "stream"] }
futures = "0.3"
rmcp = { version = "1.2", features = ["client", "transport-child-process", "transport-streamable-http-client-reqwest"] }
regex = "1"
hex = "0.4"
toml_edit = "0.22"
sqlite-vec = { workspace = true }
libsqlite3-sys = { workspace = true }
pdf-extract = "0.7"
zip = "2"
litchi = "0.0.1"
dirs = "5"
boosapi-migration = { path = "../migration" }
[target.'cfg(target_os = "windows")'.dependencies]
windows-sys = { version = "0.59", features = ["Win32_System_Threading"] }
[dev-dependencies]
tempfile = "3"
@@ -0,0 +1,572 @@
use crate::error::{BoosApiError, Result};
use crate::mcp_client::McpToolResult;
use regex::Regex;
use serde_json::Value;
/// Dispatch a tool call to the appropriate built-in implementation.
pub async fn dispatch(server_name: &str, tool_name: &str, args: Value) -> Result<McpToolResult> {
match server_name {
"@boosapi/fetch" => match tool_name {
"fetch_url" => {
let url = args.get("url").and_then(|v| v.as_str()).unwrap_or_default();
let max_length = args
.get("max_length")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
fetch_url(url, max_length).await
}
"fetch_markdown" => {
let url = args.get("url").and_then(|v| v.as_str()).unwrap_or_default();
let max_length = args
.get("max_length")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
fetch_markdown(url, max_length).await
}
_ => Err(BoosApiError::Gateway(format!(
"Unknown fetch tool: {}",
tool_name
))),
},
"@boosapi/search-file" => match tool_name {
"read_file" => {
let path = args
.get("path")
.and_then(|v| v.as_str())
.unwrap_or_default();
read_file(path).await
}
"list_directory" => {
let path = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
list_directory(path).await
}
"search_files" => {
let path = args.get("path").and_then(|v| v.as_str()).unwrap_or(".");
let pattern = args.get("pattern").and_then(|v| v.as_str()).unwrap_or("*");
let max_results = args
.get("max_results")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
search_files(path, pattern, max_results).await
}
_ => Err(BoosApiError::Gateway(format!(
"Unknown search-file tool: {}",
tool_name
))),
},
_ => Err(BoosApiError::Gateway(format!(
"Unknown builtin server: {}",
server_name
))),
}
}
// ---------------------------------------------------------------------------
// Fetch tools
// ---------------------------------------------------------------------------
async fn fetch_url(url: &str, max_length: Option<usize>) -> Result<McpToolResult> {
if url.is_empty() {
return Ok(McpToolResult {
content: "Error: url parameter is required".into(),
is_error: true,
});
}
let client = reqwest::Client::builder()
.user_agent("BoosApi/1.0 (Web Fetch Tool)")
.timeout(std::time::Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::limited(5))
.build()
.map_err(|e| BoosApiError::Gateway(format!("HTTP client error: {}", e)))?;
let resp = client
.get(url)
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to fetch {}: {}", url, e)))?;
let status = resp.status();
if !status.is_success() {
return Ok(McpToolResult {
content: format!(
"HTTP error: {} {}",
status.as_u16(),
status.canonical_reason().unwrap_or("Unknown")
),
is_error: true,
});
}
let body = resp
.text()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read response body: {}", e)))?;
let text = html_to_text(&body);
let max = max_length.unwrap_or(5000);
let content = truncate_text(&text, max);
Ok(McpToolResult {
content,
is_error: false,
})
}
async fn fetch_markdown(url: &str, max_length: Option<usize>) -> Result<McpToolResult> {
if url.is_empty() {
return Ok(McpToolResult {
content: "Error: url parameter is required".into(),
is_error: true,
});
}
let client = reqwest::Client::builder()
.user_agent("BoosApi/1.0 (Web Fetch Tool)")
.timeout(std::time::Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::limited(5))
.build()
.map_err(|e| BoosApiError::Gateway(format!("HTTP client error: {}", e)))?;
let resp = client
.get(url)
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to fetch {}: {}", url, e)))?;
let status = resp.status();
if !status.is_success() {
return Ok(McpToolResult {
content: format!(
"HTTP error: {} {}",
status.as_u16(),
status.canonical_reason().unwrap_or("Unknown")
),
is_error: true,
});
}
let body = resp
.text()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read response body: {}", e)))?;
let markdown = html_to_markdown(&body);
let max = max_length.unwrap_or(10000);
let content = truncate_text(&markdown, max);
Ok(McpToolResult {
content,
is_error: false,
})
}
// ---------------------------------------------------------------------------
// File tools
// ---------------------------------------------------------------------------
async fn read_file(path: &str) -> Result<McpToolResult> {
if path.is_empty() {
return Ok(McpToolResult {
content: "Error: path parameter is required".into(),
is_error: true,
});
}
match tokio::fs::read_to_string(path).await {
Ok(content) => {
let truncated = truncate_text(&content, 50000);
Ok(McpToolResult {
content: truncated,
is_error: false,
})
}
Err(e) => Ok(McpToolResult {
content: format!("Error reading file '{}': {}", path, e),
is_error: true,
}),
}
}
async fn list_directory(path: &str) -> Result<McpToolResult> {
let mut entries = match tokio::fs::read_dir(path).await {
Ok(rd) => rd,
Err(e) => {
return Ok(McpToolResult {
content: format!("Error listing directory '{}': {}", path, e),
is_error: true,
});
}
};
let mut items = Vec::new();
while let Ok(Some(entry)) = entries.next_entry().await {
let name = entry.file_name().to_string_lossy().to_string();
let is_dir = entry
.file_type()
.await
.map(|ft| ft.is_dir())
.unwrap_or(false);
let meta = entry.metadata().await.ok();
let size = meta.as_ref().map(|m| m.len()).unwrap_or(0);
if is_dir {
items.push(format!("📁 {}/", name));
} else {
items.push(format!("📄 {} ({})", name, human_size(size)));
}
}
items.sort();
let content = if items.is_empty() {
format!("Directory '{}' is empty", path)
} else {
format!("Contents of '{}':\n{}", path, items.join("\n"))
};
Ok(McpToolResult {
content,
is_error: false,
})
}
async fn search_files(
path: &str,
pattern: &str,
max_results: Option<usize>,
) -> Result<McpToolResult> {
let max = max_results.unwrap_or(50);
let mut results = Vec::new();
let pattern_lower = pattern.to_lowercase();
walk_dir_search(
std::path::Path::new(path),
&pattern_lower,
&mut results,
max,
)
.await;
let content = if results.is_empty() {
format!("No files matching '{}' found in '{}'", pattern, path)
} else {
format!(
"Found {} file(s) matching '{}':\n{}",
results.len(),
pattern,
results.join("\n")
)
};
Ok(McpToolResult {
content,
is_error: false,
})
}
async fn walk_dir_search(
root: &std::path::Path,
pattern: &str,
results: &mut Vec<String>,
max: usize,
) {
let mut stack = vec![root.to_path_buf()];
while let Some(dir) = stack.pop() {
if results.len() >= max {
return;
}
let mut entries = match tokio::fs::read_dir(&dir).await {
Ok(rd) => rd,
Err(_) => continue,
};
while let Ok(Some(entry)) = entries.next_entry().await {
if results.len() >= max {
return;
}
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with('.') {
continue;
}
let path = entry.path();
let is_dir = entry
.file_type()
.await
.map(|ft| ft.is_dir())
.unwrap_or(false);
if name.to_lowercase().contains(pattern) {
results.push(path.to_string_lossy().to_string());
}
if is_dir {
stack.push(path);
}
}
}
}
// ---------------------------------------------------------------------------
// HTML processing
// ---------------------------------------------------------------------------
fn html_to_text(html: &str) -> String {
let mut text = html.to_string();
remove_blocks(&mut text, "script");
remove_blocks(&mut text, "style");
remove_blocks(&mut text, "noscript");
remove_blocks(&mut text, "nav");
remove_blocks(&mut text, "footer");
remove_blocks(&mut text, "header");
// Replace block-level tags with newlines
let re_block = Regex::new(r"(?i)</?(p|div|section|article|main|aside|blockquote|pre|table|tr|ul|ol|dl|dt|dd|figcaption|figure)\s*[^>]*>").unwrap();
text = re_block.replace_all(&text, "\n").to_string();
let re_br = Regex::new(r"(?i)<br\s*/?>").unwrap();
text = re_br.replace_all(&text, "\n").to_string();
let re_hr = Regex::new(r"(?i)<hr\s*/?>").unwrap();
text = re_hr.replace_all(&text, "\n---\n").to_string();
let re_li = Regex::new(r"(?i)<li\s*[^>]*>").unwrap();
text = re_li.replace_all(&text, "\n• ").to_string();
// Strip remaining tags
let re_tags = Regex::new(r"<[^>]+>").unwrap();
text = re_tags.replace_all(&text, "").to_string();
decode_entities(&mut text);
collapse_whitespace(&text)
}
fn html_to_markdown(html: &str) -> String {
let mut text = html.to_string();
remove_blocks(&mut text, "script");
remove_blocks(&mut text, "style");
remove_blocks(&mut text, "noscript");
remove_blocks(&mut text, "nav");
remove_blocks(&mut text, "footer");
// Headings
for level in 1..=6 {
let prefix = "#".repeat(level);
let re_h = Regex::new(&format!(r"(?is)<h{}\s*[^>]*>(.*?)</h{}>", level, level)).unwrap();
text = re_h
.replace_all(&text, |caps: &regex::Captures| {
let inner = strip_tags(caps.get(1).map_or("", |m| m.as_str()));
format!("\n\n{} {}\n\n", prefix, inner.trim())
})
.to_string();
}
// Links
let re_a = Regex::new(r#"(?is)<a\s[^>]*href\s*=\s*["']([^"']+)["'][^>]*>(.*?)</a>"#).unwrap();
text = re_a
.replace_all(&text, |caps: &regex::Captures| {
let href = caps.get(1).map_or("", |m| m.as_str());
let inner = strip_tags(caps.get(2).map_or("", |m| m.as_str()));
format!("[{}]({})", inner.trim(), href)
})
.to_string();
// Images
let re_img = Regex::new(r#"(?i)<img\s[^>]*src\s*=\s*["']([^"']+)["'][^>]*/?\s*>"#).unwrap();
text = re_img
.replace_all(&text, |caps: &regex::Captures| {
let src = caps.get(1).map_or("", |m| m.as_str());
format!("![image]({})", src)
})
.to_string();
// Bold
let re_b = Regex::new(r"(?is)<(b|strong)\s*[^>]*>(.*?)</(b|strong)>").unwrap();
text = re_b
.replace_all(&text, |caps: &regex::Captures| {
let inner = caps.get(2).map_or("", |m| m.as_str());
format!("**{}**", inner.trim())
})
.to_string();
// Italic
let re_i = Regex::new(r"(?is)<(i|em)\s*[^>]*>(.*?)</(i|em)>").unwrap();
text = re_i
.replace_all(&text, |caps: &regex::Captures| {
let inner = caps.get(2).map_or("", |m| m.as_str());
format!("*{}*", inner.trim())
})
.to_string();
// Code
let re_code = Regex::new(r"(?is)<code\s*[^>]*>(.*?)</code>").unwrap();
text = re_code.replace_all(&text, "`$1`").to_string();
// Pre blocks
let re_pre = Regex::new(r"(?is)<pre\s*[^>]*>(.*?)</pre>").unwrap();
text = re_pre
.replace_all(&text, |caps: &regex::Captures| {
let inner = strip_tags(caps.get(1).map_or("", |m| m.as_str()));
format!("\n```\n{}\n```\n", inner.trim())
})
.to_string();
// List items
let re_li = Regex::new(r"(?i)<li\s*[^>]*>").unwrap();
text = re_li.replace_all(&text, "\n- ").to_string();
// Block elements → newlines
let re_block = Regex::new(
r"(?i)</?(p|div|section|article|main|aside|blockquote|table|tr|ul|ol|dl)\s*[^>]*>",
)
.unwrap();
text = re_block.replace_all(&text, "\n").to_string();
let re_br = Regex::new(r"(?i)<br\s*/?>").unwrap();
text = re_br.replace_all(&text, "\n").to_string();
let re_hr = Regex::new(r"(?i)<hr\s*/?>").unwrap();
text = re_hr.replace_all(&text, "\n---\n").to_string();
// Strip remaining tags
let re_tags = Regex::new(r"<[^>]+>").unwrap();
text = re_tags.replace_all(&text, "").to_string();
decode_entities(&mut text);
collapse_whitespace(&text)
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
fn remove_blocks(html: &mut String, tag: &str) {
let open = format!("<{}", tag);
let close = format!("</{}>", tag);
loop {
let lower = html.to_lowercase();
let start = match lower.find(&open) {
Some(s) => s,
None => break,
};
let search_from = start + open.len();
let end = match lower[search_from..].find(&close) {
Some(offset) => search_from + offset + close.len(),
None => break,
};
html.replace_range(start..end, "");
}
}
fn strip_tags(s: &str) -> String {
let re = Regex::new(r"<[^>]+>").unwrap();
re.replace_all(s, "").to_string()
}
fn decode_entities(text: &mut String) {
let replacements = [
("&amp;", "&"),
("&lt;", "<"),
("&gt;", ">"),
("&quot;", "\""),
("&#39;", "'"),
("&apos;", "'"),
("&nbsp;", " "),
("&#x27;", "'"),
("&#x2F;", "/"),
("&mdash;", "—"),
("&ndash;", "–"),
("&laquo;", "«"),
("&raquo;", "»"),
("&bull;", "•"),
("&hellip;", "…"),
("&copy;", "©"),
("&reg;", "®"),
("&trade;", "™"),
];
for (entity, replacement) in &replacements {
*text = text.replace(entity, replacement);
}
// Numeric entities: &#NNN;
let re_dec = Regex::new(r"&#(\d+);").unwrap();
*text = re_dec
.replace_all(text, |caps: &regex::Captures| {
caps.get(1)
.and_then(|m| m.as_str().parse::<u32>().ok())
.and_then(char::from_u32)
.map(|c| c.to_string())
.unwrap_or_default()
})
.to_string();
// Hex entities: &#xHHH;
let re_hex = Regex::new(r"(?i)&#x([0-9a-f]+);").unwrap();
*text = re_hex
.replace_all(text, |caps: &regex::Captures| {
caps.get(1)
.and_then(|m| u32::from_str_radix(m.as_str(), 16).ok())
.and_then(char::from_u32)
.map(|c| c.to_string())
.unwrap_or_default()
})
.to_string();
}
fn collapse_whitespace(text: &str) -> String {
let mut result = String::with_capacity(text.len());
let mut consecutive_newlines = 0u32;
let mut last_was_space = false;
for c in text.chars() {
if c == '\n' || c == '\r' {
if c == '\r' {
continue;
}
consecutive_newlines += 1;
if consecutive_newlines <= 2 {
result.push('\n');
}
last_was_space = true;
} else if c.is_whitespace() {
consecutive_newlines = 0;
if !last_was_space {
result.push(' ');
last_was_space = true;
}
} else {
consecutive_newlines = 0;
last_was_space = false;
result.push(c);
}
}
result.trim().to_string()
}
fn truncate_text(text: &str, max: usize) -> String {
if text.len() <= max {
return text.to_string();
}
// Find a clean break point near max
let boundary = text[..max].rfind('\n').unwrap_or(max);
format!(
"{}\n\n... (truncated, showing first {} of {} characters)",
&text[..boundary],
boundary,
text.len()
)
}
fn human_size(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
let mut size = bytes as f64;
for unit in UNITS {
if size < 1024.0 {
return format!("{:.1} {}", size, unit);
}
size /= 1024.0;
}
format!("{:.1} TB", size)
}
@@ -0,0 +1,78 @@
use aes_gcm::{
aead::{rand_core::RngCore, Aead, KeyInit, OsRng},
Aes256Gcm, Nonce,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use sha2::{Digest, Sha256};
use crate::error::{BoosApiError, Result};
const NONCE_SIZE: usize = 12;
pub fn generate_master_key() -> [u8; 32] {
let mut key = [0u8; 32];
OsRng.fill_bytes(&mut key);
key
}
pub fn encrypt_key(plaintext: &str, master_key: &[u8; 32]) -> Result<String> {
let cipher = Aes256Gcm::new_from_slice(master_key)
.map_err(|e| BoosApiError::Crypto(format!("Failed to create cipher: {}", e)))?;
let mut nonce_bytes = [0u8; NONCE_SIZE];
OsRng.fill_bytes(&mut nonce_bytes);
let nonce = Nonce::from_slice(&nonce_bytes);
let ciphertext = cipher
.encrypt(nonce, plaintext.as_bytes())
.map_err(|e| BoosApiError::Crypto(format!("Encryption failed: {}", e)))?;
let mut combined = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
combined.extend_from_slice(&nonce_bytes);
combined.extend_from_slice(&ciphertext);
Ok(BASE64.encode(&combined))
}
pub fn decrypt_key(encrypted: &str, master_key: &[u8; 32]) -> Result<String> {
let combined = BASE64
.decode(encrypted)
.map_err(|e| BoosApiError::Crypto(format!("Base64 decode failed: {}", e)))?;
if combined.len() < NONCE_SIZE {
return Err(BoosApiError::Crypto("Invalid encrypted data".to_string()));
}
let (nonce_bytes, ciphertext) = combined.split_at(NONCE_SIZE);
let nonce = Nonce::from_slice(nonce_bytes);
let cipher = Aes256Gcm::new_from_slice(master_key)
.map_err(|e| BoosApiError::Crypto(format!("Failed to create cipher: {}", e)))?;
let plaintext = cipher
.decrypt(nonce, ciphertext)
.map_err(|e| BoosApiError::Crypto(format!("Decryption failed: {}", e)))?;
String::from_utf8(plaintext)
.map_err(|e| BoosApiError::Crypto(format!("UTF-8 decode failed: {}", e)))
}
pub fn sha256_hash(input: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(input.as_bytes());
format!("{:x}", hasher.finalize())
}
pub fn key_prefix(key: &str) -> String {
if key.len() > 8 {
format!("{}...{}", &key[..4], &key[key.len() - 4..])
} else {
"****".to_string()
}
}
pub fn generate_gateway_key() -> String {
let mut bytes = [0u8; 32];
OsRng.fill_bytes(&mut bytes);
format!("aq-{}", hex::encode(bytes))
}
+745
View File
@@ -0,0 +1,745 @@
use boosapi_migration::MigratorTrait;
use sea_orm::{
ConnectOptions, ConnectionTrait, Database, DatabaseConnection, DbBackend, Statement,
};
use tracing::info;
use crate::error::Result;
use crate::types::*;
pub struct DbHandle {
pub conn: DatabaseConnection,
}
pub async fn create_pool(db_path: &str) -> Result<DbHandle> {
let url = if db_path.starts_with("sqlite:") {
format!("{}?mode=rwc", db_path)
} else {
format!("sqlite:{}?mode=rwc", db_path)
};
let mut opt = ConnectOptions::new(&url);
opt.max_connections(5)
.min_connections(1)
.sqlx_logging(false);
let conn = Database::connect(opt).await?;
// Enable WAL journal mode and foreign keys via PRAGMA
conn.execute(Statement::from_string(
DbBackend::Sqlite,
"PRAGMA journal_mode=WAL;",
))
.await?;
conn.execute(Statement::from_string(
DbBackend::Sqlite,
"PRAGMA foreign_keys=ON;",
))
.await?;
// Run SeaORM migrations
boosapi_migration::Migrator::up(&conn, None).await?;
info!("Database initialized at {}", db_path);
Ok(DbHandle { conn })
}
pub struct BuiltinProvider {
pub builtin_id: &'static str,
pub name: &'static str,
pub provider_type: ProviderType,
pub api_host: &'static str,
pub models: Vec<BuiltinModel>,
}
pub struct BuiltinModel {
pub model_id: &'static str,
pub name: &'static str,
pub group_name: Option<&'static str>,
pub model_type: Option<ModelType>,
pub capabilities: Vec<ModelCapability>,
pub max_tokens: Option<u32>,
pub enabled: bool,
pub param_overrides: Option<ModelParamOverrides>,
}
impl BuiltinModel {
fn chat(
model_id: &'static str,
name: &'static str,
capabilities: Vec<ModelCapability>,
max_tokens: Option<u32>,
) -> Self {
Self {
model_id,
name,
group_name: None,
model_type: Some(ModelType::Chat),
capabilities,
max_tokens,
enabled: true,
param_overrides: None,
}
}
fn image(model_id: &'static str, name: &'static str) -> Self {
Self {
model_id,
name,
group_name: Some("gpt-image"),
model_type: Some(ModelType::Image),
capabilities: vec![],
max_tokens: None,
enabled: true,
param_overrides: None,
}
}
fn rerank(model_id: &'static str, name: &'static str) -> Self {
Self {
model_id,
name,
group_name: None,
model_type: Some(ModelType::Rerank),
capabilities: vec![],
max_tokens: None,
enabled: true,
param_overrides: None,
}
}
fn disabled(mut self) -> Self {
self.enabled = false;
self
}
fn with_param_overrides(mut self, param_overrides: ModelParamOverrides) -> Self {
self.param_overrides = Some(param_overrides);
self
}
pub(crate) fn to_model(&self, provider_id: &str) -> Model {
Model {
provider_id: provider_id.to_string(),
model_id: self.model_id.to_string(),
name: self.name.to_string(),
group_name: self.group_name.map(str::to_string),
model_type: self
.model_type
.clone()
.unwrap_or_else(|| ModelType::detect(self.model_id)),
capabilities: self.capabilities.clone(),
max_tokens: self.max_tokens,
enabled: self.enabled,
param_overrides: self.param_overrides.clone(),
}
}
}
fn empty_param_overrides() -> ModelParamOverrides {
ModelParamOverrides {
temperature: None,
max_tokens: None,
top_p: None,
frequency_penalty: None,
use_max_completion_tokens: None,
no_system_role: None,
force_max_tokens: None,
thinking_param_style: None,
reasoning_profile: None,
reasoning_options: None,
reasoning_default: None,
extra_body: None,
}
}
fn reasoning_profile(profile: &'static str) -> ModelParamOverrides {
let mut overrides = empty_param_overrides();
overrides.reasoning_profile = Some(profile.to_string());
overrides
}
fn openai_reasoning_profile() -> ModelParamOverrides {
let mut overrides = reasoning_profile("openai_reasoning_effort");
overrides.use_max_completion_tokens = Some(true);
overrides
}
fn minimax_m2_profile() -> ModelParamOverrides {
let mut overrides = empty_param_overrides();
overrides.max_tokens = Some(2048);
overrides.use_max_completion_tokens = Some(true);
overrides
}
pub fn get_builtin_providers() -> Vec<BuiltinProvider> {
use ModelCapability::*;
vec![
// BoosApi's default gateway. OpenAI-compatible aggregator that
// proxies multiple upstream providers behind a single key. Served via
// the Custom adapter (base URL is configurable in provider settings).
// The model list below is a starter set; users can refresh the live
// catalog from the gateway's /v1/models endpoint.
BuiltinProvider {
builtin_id: "boosapi",
name: "BoosAPI",
provider_type: ProviderType::Custom,
api_host: "https://api.boos.lat",
models: vec![
BuiltinModel::chat(
"claude-opus-4-8",
"Claude Opus 4.8",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
),
BuiltinModel::chat(
"gpt-5.5",
"GPT-5.5",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
),
BuiltinModel::chat(
"gemini-3-pro",
"Gemini 3 Pro",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
),
BuiltinModel::chat(
"deepseek-v4-pro",
"DeepSeek v4 Pro",
vec![TextChat, Reasoning, FunctionCalling],
Some(1_000_000),
),
],
},
BuiltinProvider {
builtin_id: "openai",
name: "OpenAI",
provider_type: ProviderType::OpenAI,
api_host: "https://api.openai.com",
models: vec![
BuiltinModel::chat(
"gpt-5.5",
"GPT-5.5",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
)
.with_param_overrides(openai_reasoning_profile()),
BuiltinModel::chat(
"gpt-5.4",
"GPT-5.4",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
)
.with_param_overrides(openai_reasoning_profile()),
BuiltinModel::chat(
"gpt-5.4-mini",
"GPT-5.4 Mini",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(400_000),
)
.with_param_overrides(openai_reasoning_profile()),
BuiltinModel::chat(
"gpt-5.4-nano",
"GPT-5.4 Nano",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(400_000),
)
.with_param_overrides(openai_reasoning_profile())
.disabled(),
BuiltinModel::chat(
"gpt-4.1",
"GPT-4.1",
vec![TextChat, Vision, FunctionCalling],
Some(1_047_576),
),
BuiltinModel::chat(
"gpt-4.1-mini",
"GPT-4.1 Mini",
vec![TextChat, Vision, FunctionCalling],
Some(1_047_576),
),
BuiltinModel::chat(
"gpt-4.1-nano",
"GPT-4.1 Nano",
vec![TextChat, Vision, FunctionCalling],
Some(1_047_576),
)
.disabled(),
BuiltinModel::chat(
"gpt-4o",
"GPT-4o",
vec![TextChat, Vision, FunctionCalling],
Some(128_000),
)
.disabled(),
BuiltinModel::chat(
"gpt-4o-mini",
"GPT-4o Mini",
vec![TextChat, Vision, FunctionCalling],
Some(128_000),
)
.disabled(),
BuiltinModel::chat(
"o3",
"o3",
vec![TextChat, Reasoning, FunctionCalling],
Some(200_000),
)
.with_param_overrides(openai_reasoning_profile()),
BuiltinModel::chat(
"o4-mini",
"o4-mini",
vec![TextChat, Reasoning, FunctionCalling],
Some(200_000),
)
.with_param_overrides(openai_reasoning_profile()),
BuiltinModel::image("gpt-image-2", "gpt-image-2"),
BuiltinModel::image("gpt-image-1.5", "gpt-image-1.5"),
BuiltinModel::image("gpt-image-1", "gpt-image-1").disabled(),
BuiltinModel::image("gpt-image-1-mini", "gpt-image-1-mini").disabled(),
],
},
BuiltinProvider {
builtin_id: "openai_responses",
name: "OpenAI Responses",
provider_type: ProviderType::OpenAIResponses,
api_host: "https://api.openai.com",
models: vec![
BuiltinModel::chat(
"gpt-5.5",
"GPT-5.5",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
)
.with_param_overrides(reasoning_profile("openai_responses_reasoning")),
BuiltinModel::chat(
"gpt-5.4",
"GPT-5.4",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_000_000),
)
.with_param_overrides(reasoning_profile("openai_responses_reasoning")),
BuiltinModel::chat(
"gpt-5.4-mini",
"GPT-5.4 Mini",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(400_000),
)
.with_param_overrides(reasoning_profile("openai_responses_reasoning")),
BuiltinModel::chat(
"gpt-4.1",
"GPT-4.1",
vec![TextChat, Vision, FunctionCalling],
Some(1_047_576),
),
BuiltinModel::chat(
"gpt-4o",
"GPT-4o",
vec![TextChat, Vision, FunctionCalling],
Some(128_000),
)
.disabled(),
BuiltinModel::chat(
"gpt-4o-mini",
"GPT-4o Mini",
vec![TextChat, Vision, FunctionCalling],
Some(128_000),
)
.disabled(),
BuiltinModel::chat(
"o3",
"o3",
vec![TextChat, Reasoning, FunctionCalling],
Some(200_000),
)
.with_param_overrides(reasoning_profile("openai_responses_reasoning")),
BuiltinModel::chat(
"o4-mini",
"o4-mini",
vec![TextChat, Reasoning, FunctionCalling],
Some(200_000),
)
.with_param_overrides(reasoning_profile("openai_responses_reasoning")),
],
},
BuiltinProvider {
builtin_id: "gemini",
name: "Gemini",
provider_type: ProviderType::Gemini,
api_host: "https://generativelanguage.googleapis.com",
models: vec![
BuiltinModel::chat(
"gemini-3.1-pro-preview",
"Gemini 3.1 Pro Preview",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_048_576),
)
.with_param_overrides(reasoning_profile("gemini_thinking_level")),
BuiltinModel::chat(
"gemini-3.1-flash-lite-preview",
"Gemini 3.1 Flash-Lite Preview",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_048_576),
)
.with_param_overrides(reasoning_profile("gemini_thinking_level")),
BuiltinModel::chat(
"gemini-2.5-pro",
"Gemini 2.5 Pro",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_048_576),
)
.with_param_overrides(reasoning_profile("gemini_thinking_budget")),
BuiltinModel::chat(
"gemini-2.5-flash",
"Gemini 2.5 Flash",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_048_576),
)
.with_param_overrides(reasoning_profile("gemini_thinking_budget")),
BuiltinModel::chat(
"gemini-2.5-flash-lite",
"Gemini 2.5 Flash-Lite",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(1_048_576),
)
.with_param_overrides(reasoning_profile("gemini_thinking_budget")),
BuiltinModel::chat(
"gemini-2.0-flash",
"Gemini 2.0 Flash",
vec![TextChat, Vision, FunctionCalling],
Some(1_048_576),
)
.disabled(),
],
},
BuiltinProvider {
builtin_id: "anthropic",
name: "Claude",
provider_type: ProviderType::Anthropic,
api_host: "https://api.anthropic.com",
models: vec![
BuiltinModel::chat(
"claude-opus-4-7-20260127",
"Claude Opus 4.7",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(200_000),
)
.with_param_overrides(reasoning_profile("anthropic_adaptive")),
BuiltinModel::chat(
"claude-sonnet-4-6-20251117",
"Claude Sonnet 4.6",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(200_000),
)
.with_param_overrides(reasoning_profile("anthropic_adaptive")),
BuiltinModel::chat(
"claude-haiku-4-5-20251001",
"Claude Haiku 4.5",
vec![TextChat, Vision, FunctionCalling, Reasoning],
Some(200_000),
)
.with_param_overrides(reasoning_profile("anthropic_budget_tokens")),
],
},
BuiltinProvider {
builtin_id: "deepseek",
name: "DeepSeek",
provider_type: ProviderType::DeepSeek,
api_host: "https://api.deepseek.com",
models: vec![
BuiltinModel::chat(
"deepseek-v4-flash",
"DeepSeek v4 Flash",
vec![TextChat, Reasoning, FunctionCalling],
Some(1_000_000),
)
.with_param_overrides(reasoning_profile("openai_reasoning_effort")),
BuiltinModel::chat(
"deepseek-v4-pro",
"DeepSeek v4 Pro",
vec![TextChat, Reasoning, FunctionCalling],
Some(1_000_000),
)
.with_param_overrides(reasoning_profile("openai_reasoning_effort")),
BuiltinModel::chat(
"deepseek-chat",
"DeepSeek Chat",
vec![TextChat, FunctionCalling],
Some(64_000),
)
.disabled(),
BuiltinModel::chat(
"deepseek-reasoner",
"DeepSeek Reasoner",
vec![TextChat, Reasoning],
Some(64_000),
)
.with_param_overrides(reasoning_profile("openai_reasoning_effort"))
.disabled(),
],
},
BuiltinProvider {
builtin_id: "xai",
name: "xAI",
provider_type: ProviderType::XAI,
api_host: "https://api.x.ai",
models: vec![
BuiltinModel::chat(
"grok-4.3",
"Grok 4.3",
vec![TextChat, Vision, Reasoning, FunctionCalling],
None,
)
.with_param_overrides(reasoning_profile("openai_reasoning_effort")),
BuiltinModel::chat(
"grok-3",
"Grok 3",
vec![TextChat, Vision, FunctionCalling],
Some(131_072),
)
.disabled(),
BuiltinModel::chat(
"grok-3-mini",
"Grok 3 Mini",
vec![TextChat, Reasoning, FunctionCalling],
Some(131_072),
)
.with_param_overrides(reasoning_profile("none"))
.disabled(),
],
},
BuiltinProvider {
builtin_id: "glm",
name: "GLM",
provider_type: ProviderType::GLM,
api_host: "https://open.bigmodel.cn/api/paas",
models: vec![
BuiltinModel::chat(
"glm-5.1",
"GLM-5.1",
vec![TextChat, Vision, Reasoning, FunctionCalling],
Some(200_000),
)
.with_param_overrides(reasoning_profile("glm_thinking")),
BuiltinModel::chat(
"glm-5",
"GLM-5",
vec![TextChat, Vision, Reasoning, FunctionCalling],
Some(128_000),
)
.with_param_overrides(reasoning_profile("glm_thinking")),
BuiltinModel::chat(
"glm-4.6",
"GLM-4.6",
vec![TextChat, Vision, Reasoning, FunctionCalling],
Some(128_000),
)
.with_param_overrides(reasoning_profile("glm_thinking"))
.disabled(),
],
},
BuiltinProvider {
builtin_id: "siliconflow",
name: "SiliconFlow",
provider_type: ProviderType::SiliconFlow,
api_host: "https://api.siliconflow.cn",
models: vec![
BuiltinModel::chat(
"deepseek-ai/DeepSeek-V3.2-Exp",
"DeepSeek-V3.2-Exp",
vec![TextChat, FunctionCalling],
Some(64_000),
),
BuiltinModel::chat(
"deepseek-ai/DeepSeek-R1",
"DeepSeek-R1",
vec![TextChat, Reasoning],
Some(64_000),
)
.with_param_overrides(reasoning_profile("siliconflow_enable_thinking")),
BuiltinModel::chat(
"Qwen/Qwen3-235B-A22B",
"Qwen3-235B-A22B",
vec![TextChat, Reasoning, FunctionCalling],
Some(262_144),
)
.with_param_overrides(reasoning_profile("siliconflow_enable_thinking")),
BuiltinModel::chat(
"Qwen/Qwen3-Coder-480B-A35B-Instruct",
"Qwen3-Coder-480B-A35B-Instruct",
vec![TextChat, FunctionCalling],
Some(262_144),
),
],
},
BuiltinProvider {
builtin_id: "minimax",
name: "MiniMax",
provider_type: ProviderType::OpenAI,
api_host: "https://api.minimax.io",
models: vec![
BuiltinModel::chat(
"MiniMax-M2.7",
"MiniMax-M2.7",
vec![TextChat, FunctionCalling],
Some(250_000),
)
.with_param_overrides(minimax_m2_profile()),
BuiltinModel::chat(
"MiniMax-M2.5",
"MiniMax-M2.5",
vec![TextChat, FunctionCalling],
Some(250_000),
)
.with_param_overrides(minimax_m2_profile()),
BuiltinModel::chat(
"MiniMax-M1",
"MiniMax-M1",
vec![TextChat, Reasoning, FunctionCalling],
Some(1_000_000),
)
.disabled(),
],
},
BuiltinProvider {
builtin_id: "jina",
name: "Jina",
provider_type: ProviderType::Jina,
api_host: "https://api.jina.ai",
models: vec![
BuiltinModel::rerank("jina-reranker-v3", "Jina Reranker v3"),
BuiltinModel::rerank(
"jina-reranker-v2-base-multilingual",
"Jina Reranker v2 Base Multilingual",
)
.disabled(),
BuiltinModel::rerank("jina-colbert-v2", "Jina ColBERT v2").disabled(),
],
},
BuiltinProvider {
builtin_id: "cohere",
name: "Cohere",
provider_type: ProviderType::Cohere,
api_host: "https://api.cohere.com",
models: vec![
BuiltinModel::rerank("rerank-v4.0", "Rerank v4.0"),
BuiltinModel::rerank("rerank-v4.0-pro", "Rerank v4.0 Pro"),
BuiltinModel::rerank("rerank-v4.0-fast", "Rerank v4.0 Fast"),
BuiltinModel::rerank("rerank-v3.5", "Rerank v3.5").disabled(),
],
},
BuiltinProvider {
builtin_id: "voyage",
name: "Voyage",
provider_type: ProviderType::Voyage,
api_host: "https://api.voyageai.com",
models: vec![
BuiltinModel::rerank("rerank-2.5", "Rerank 2.5"),
BuiltinModel::rerank("rerank-2.5-lite", "Rerank 2.5 Lite"),
BuiltinModel::rerank("rerank-2", "Rerank 2").disabled(),
BuiltinModel::rerank("rerank-2-lite", "Rerank 2 Lite").disabled(),
],
},
]
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn builtin_rerank_providers_are_registered_with_rerank_models() {
let providers = get_builtin_providers();
for (builtin_id, provider_type, model_id) in [
("jina", ProviderType::Jina, "jina-reranker-v3"),
("cohere", ProviderType::Cohere, "rerank-v4.0-pro"),
("voyage", ProviderType::Voyage, "rerank-2.5"),
] {
let provider = providers
.iter()
.find(|provider| provider.builtin_id == builtin_id)
.expect("missing rerank provider");
assert_eq!(provider.provider_type, provider_type);
assert!(
provider
.models
.iter()
.any(|model| model.model_id == model_id
&& model.model_type.as_ref() == Some(&ModelType::Rerank))
);
}
}
#[test]
fn builtin_openai_compatible_providers_use_dedicated_types() {
let providers = get_builtin_providers();
for (builtin_id, provider_type) in [
("deepseek", ProviderType::DeepSeek),
("xai", ProviderType::XAI),
("glm", ProviderType::GLM),
("siliconflow", ProviderType::SiliconFlow),
] {
let provider = providers
.iter()
.find(|provider| provider.builtin_id == builtin_id)
.expect("missing builtin provider");
assert_eq!(provider.provider_type, provider_type);
}
}
#[test]
fn builtin_models_include_current_reasoning_params() {
let providers = get_builtin_providers();
let provider = providers
.iter()
.find(|provider| provider.builtin_id == "deepseek")
.expect("missing DeepSeek builtin provider");
let model = provider
.models
.iter()
.find(|model| model.model_id == "deepseek-v4-flash")
.expect("missing DeepSeek v4 Flash model");
assert!(model.capabilities.contains(&ModelCapability::Reasoning));
assert_eq!(model.max_tokens, Some(1_000_000));
assert_eq!(
model
.param_overrides
.as_ref()
.and_then(|params| params.reasoning_profile.as_deref()),
Some("openai_reasoning_effort")
);
let minimax = providers
.iter()
.find(|provider| provider.builtin_id == "minimax")
.expect("missing MiniMax builtin provider");
let m2 = minimax
.models
.iter()
.find(|model| model.model_id == "MiniMax-M2.7")
.expect("missing MiniMax-M2.7 model");
assert_eq!(
m2.param_overrides
.as_ref()
.and_then(|params| params.use_max_completion_tokens),
Some(true)
);
assert_eq!(
m2.param_overrides
.as_ref()
.and_then(|params| params.max_tokens),
Some(2048)
);
}
}
pub async fn create_test_pool() -> Result<DbHandle> {
create_pool("sqlite::memory:").await
}
@@ -0,0 +1,134 @@
use crate::error::{BoosApiError, Result};
use std::path::Path;
/// Extract plain text from a document file based on its MIME type.
pub fn extract_text(file_path: &Path, mime_type: &str) -> Result<String> {
match mime_type {
// Plain text files
"text/plain" | "text/markdown" | "text/csv" | "text/html" | "text/xml"
| "application/json" | "application/xml" => std::fs::read_to_string(file_path)
.map_err(|e| BoosApiError::Provider(format!("Failed to read file: {e}"))),
// PDF
"application/pdf" => extract_pdf(file_path),
// Legacy Word document
"application/msword" => extract_word_document(file_path),
// DOCX — basic XML extraction without external crate
"application/vnd.openxmlformats-officedocument.wordprocessingml.document" => {
extract_docx(file_path)
}
_ => {
// Try reading as plain text as fallback
std::fs::read_to_string(file_path).map_err(|e| {
BoosApiError::Provider(format!(
"Unsupported MIME type '{}', fallback read failed: {e}",
mime_type
))
})
}
}
}
/// Extract text from PDF using pdf-extract crate.
fn extract_pdf(file_path: &Path) -> Result<String> {
let bytes = std::fs::read(file_path)
.map_err(|e| BoosApiError::Provider(format!("Failed to read PDF file: {e}")))?;
pdf_extract::extract_text_from_mem(&bytes)
.map_err(|e| BoosApiError::Provider(format!("Failed to extract PDF text: {e}")))
}
/// Extract text from Word documents using litchi.
fn extract_word_document(file_path: &Path) -> Result<String> {
let document = litchi::Document::open(file_path)
.map_err(|e| BoosApiError::Provider(format!("Failed to open Word document: {e}")))?;
document
.text()
.map_err(|e| BoosApiError::Provider(format!("Failed to extract Word document text: {e}")))
}
/// Extract text from DOCX by reading the internal XML.
/// DOCX files are ZIP archives containing word/document.xml.
fn extract_docx(file_path: &Path) -> Result<String> {
let file = std::fs::File::open(file_path)
.map_err(|e| BoosApiError::Provider(format!("Failed to open DOCX file: {e}")))?;
let mut archive = zip::ZipArchive::new(file)
.map_err(|e| BoosApiError::Provider(format!("Failed to read DOCX as ZIP: {e}")))?;
let mut xml_content = String::new();
if let Ok(mut entry) = archive.by_name("word/document.xml") {
use std::io::Read;
entry
.read_to_string(&mut xml_content)
.map_err(|e| BoosApiError::Provider(format!("Failed to read document.xml: {e}")))?;
} else {
return Err(BoosApiError::Provider(
"DOCX: word/document.xml not found".into(),
));
}
// Simple XML text extraction: find all <w:t> tag contents
Ok(extract_text_from_xml(&xml_content))
}
/// Simple XML text extraction — pulls text from <w:t> and <w:t xml:space="preserve"> tags.
fn extract_text_from_xml(xml: &str) -> String {
let mut result = String::new();
let mut in_paragraph = false;
// Track <w:p> boundaries for paragraph breaks
for part in xml.split("<w:p") {
if in_paragraph && !result.is_empty() {
result.push('\n');
}
in_paragraph = true;
// Extract text from <w:t> tags within this paragraph
for segment in part.split("<w:t") {
if let Some(text_start) = segment.find('>') {
let after_tag = &segment[text_start + 1..];
if let Some(end) = after_tag.find("</w:t>") {
result.push_str(&after_tag[..end]);
}
}
}
}
result
}
/// Determine the MIME type from a file extension.
pub fn mime_from_extension(path: &Path) -> &'static str {
let extension = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
match extension.as_str() {
"txt" => "text/plain",
"md" | "markdown" => "text/markdown",
"csv" => "text/csv",
"html" | "htm" => "text/html",
"xml" => "text/xml",
"json" => "application/json",
"pdf" => "application/pdf",
"doc" => "application/msword",
"docx" => "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
_ => "text/plain",
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::Path;
#[test]
fn mime_from_extension_detects_legacy_word_documents() {
assert_eq!(mime_from_extension(Path::new("report.doc")), "application/msword");
}
}
@@ -0,0 +1,24 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "agent_sessions")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub cwd: Option<String>,
pub permission_mode: String,
pub runtime_status: String,
pub sdk_context_json: Option<String>,
pub sdk_context_backup_json: Option<String>,
pub total_tokens: i32,
pub total_cost_usd: f64,
pub created_at: String,
pub updated_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,22 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "artifacts")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub kind: String,
pub title: String,
#[sea_orm(column_type = "Text")]
pub content: String,
pub format: String,
pub pinned: i32,
pub updated_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,23 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "backup_manifests")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub version: String,
pub created_at: String,
pub encrypted: i32,
pub checksum: String,
#[sea_orm(column_type = "Text")]
pub object_counts_json: String,
pub source_app_version: String,
pub file_path: Option<String>,
pub file_size: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,17 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "backup_targets")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub kind: String,
#[sea_orm(column_type = "Text")]
pub config_json: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,17 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "builtin_model_deletions")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub builtin_id: String,
#[sea_orm(primary_key, auto_increment = false)]
pub model_id: String,
pub deleted_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,23 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "context_sources")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub message_id: Option<String>,
#[sea_orm(column_name = "type")]
pub source_type: String,
pub ref_id: String,
pub title: String,
pub enabled: i32,
#[sea_orm(column_type = "Text", nullable)]
pub summary: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,21 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "conversation_branches")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub parent_message_id: String,
pub branch_label: String,
pub branch_index: i32,
#[sea_orm(column_type = "Text", nullable)]
pub compared_message_ids_json: Option<String>,
pub created_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,28 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "conversation_categories")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
pub icon_type: Option<String>,
pub icon_value: Option<String>,
pub system_prompt: Option<String>,
pub default_provider_id: Option<String>,
pub default_model_id: Option<String>,
pub default_temperature: Option<f64>,
pub default_max_tokens: Option<i64>,
pub default_top_p: Option<f64>,
pub default_frequency_penalty: Option<f64>,
pub sort_order: i32,
pub is_collapsed: i32,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,34 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "conversation_summaries")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub summary_text: String,
pub compressed_until_message_id: Option<String>,
pub token_count: Option<i64>,
pub model_used: Option<String>,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::conversations::Entity",
from = "Column::ConversationId",
to = "super::conversations::Column::Id"
)]
Conversation,
}
impl Related<super::conversations::Entity> for Entity {
fn to() -> RelationDef {
Relation::Conversation.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,51 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "conversations")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub title: String,
pub model_id: String,
pub provider_id: String,
pub system_prompt: Option<String>,
pub temperature: Option<f64>,
pub max_tokens: Option<i64>,
pub top_p: Option<f64>,
pub frequency_penalty: Option<f64>,
pub search_enabled: i32,
pub search_provider_id: Option<String>,
pub thinking_budget: Option<i64>,
pub thinking_level: Option<String>,
pub enabled_mcp_server_ids: String,
pub enabled_knowledge_base_ids: String,
pub enabled_memory_namespace_ids: String,
pub message_count: i32,
pub created_at: i64,
pub updated_at: i64,
pub is_pinned: i32,
pub is_archived: i32,
pub workspace_snapshot_json: String,
pub active_branch_id: Option<String>,
pub active_artifact_id: Option<String>,
pub research_mode: i32,
pub context_compression: i32,
pub category_id: Option<String>,
pub parent_conversation_id: Option<String>,
pub mode: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::messages::Entity")]
Messages,
}
impl Related<super::messages::Entity> for Entity {
fn to() -> RelationDef {
Relation::Messages.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,20 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "desktop_state")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub window_key: String,
pub width: i32,
pub height: i32,
pub x: Option<i64>,
pub y: Option<i64>,
pub maximized: i32,
pub visible: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,31 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "drawing_generations")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub parent_generation_id: Option<String>,
pub provider_id: String,
pub key_id: String,
pub model_id: String,
pub api_kind: String,
pub action: String,
pub prompt: String,
pub parameters_json: String,
pub reference_file_ids_json: String,
pub source_image_ids_json: String,
pub mask_file_id: Option<String>,
pub status: String,
pub error_message: Option<String>,
pub response_id: Option<String>,
pub usage_json: Option<String>,
pub created_at: i64,
pub completed_at: Option<i64>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,22 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "drawing_images")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub generation_id: String,
pub stored_file_id: String,
pub storage_path: String,
pub mime_type: String,
pub width: Option<i32>,
pub height: Option<i32>,
pub revised_prompt: Option<String>,
pub created_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,18 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "gateway_diagnostics")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub category: String,
pub status: String,
pub message: String,
pub created_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,31 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "gateway_keys")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
#[sea_orm(unique)]
pub key_hash: String,
pub key_prefix: String,
pub enabled: i32,
pub created_at: i64,
pub last_used_at: Option<i64>,
pub encrypted_key: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::gateway_usage::Entity")]
GatewayUsage,
}
impl Related<super::gateway_usage::Entity> for Entity {
fn to() -> RelationDef {
Relation::GatewayUsage.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,26 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "gateway_request_logs")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub key_id: String,
pub key_name: String,
pub method: String,
pub path: String,
pub model: Option<String>,
pub provider_id: Option<String>,
pub status_code: i32,
pub duration_ms: i32,
pub request_tokens: i32,
pub response_tokens: i32,
pub error_message: Option<String>,
pub created_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,33 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "gateway_usage")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i64,
pub key_id: String,
pub provider_id: String,
pub model_id: Option<String>,
pub request_tokens: i32,
pub response_tokens: i32,
pub created_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::gateway_keys::Entity",
from = "Column::KeyId",
to = "super::gateway_keys::Column::Id"
)]
GatewayKeys,
}
impl Related<super::gateway_keys::Entity> for Entity {
fn to() -> RelationDef {
Relation::GatewayKeys.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,20 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "import_jobs")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub source_type: String,
pub status: String,
#[sea_orm(column_type = "Text", nullable)]
pub summary_json: Option<String>,
pub conflict_count: i32,
pub created_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,42 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "knowledge_bases")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
#[sea_orm(column_type = "Text", nullable)]
pub description: Option<String>,
pub embedding_provider: Option<String>,
pub enabled: i32,
pub icon_type: Option<String>,
pub icon_value: Option<String>,
pub sort_order: i32,
pub embedding_dimensions: Option<i32>,
pub retrieval_threshold: Option<f32>,
pub retrieval_top_k: Option<i32>,
pub rerank_provider: Option<String>,
pub rerank_candidate_k: Option<i32>,
pub chunk_size: Option<i32>,
pub chunk_overlap: Option<i32>,
#[sea_orm(column_type = "Text", nullable)]
pub separator: Option<String>,
pub index_concurrency: Option<i32>,
pub index_interval_ms: Option<i32>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::knowledge_documents::Entity")]
KnowledgeDocuments,
}
impl Related<super::knowledge_documents::Entity> for Entity {
fn to() -> RelationDef {
Relation::KnowledgeDocuments.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,36 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "knowledge_documents")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub knowledge_base_id: String,
pub title: String,
pub source_path: String,
pub mime_type: String,
pub size_bytes: i64,
pub indexing_status: String,
pub doc_type: String,
pub index_error: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::knowledge_bases::Entity",
from = "Column::KnowledgeBaseId",
to = "super::knowledge_bases::Column::Id",
on_delete = "Cascade"
)]
KnowledgeBase,
}
impl Related<super::knowledge_bases::Entity> for Entity {
fn to() -> RelationDef {
Relation::KnowledgeBase.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,37 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "mcp_servers")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
pub transport: String,
pub command: Option<String>,
pub args_json: Option<String>,
pub endpoint: Option<String>,
pub env_json: Option<String>,
pub enabled: i32,
pub permission_policy: String,
pub source: String,
pub discover_timeout_secs: Option<i32>,
pub execute_timeout_secs: Option<i32>,
pub headers_json: Option<String>,
pub icon_type: Option<String>,
pub icon_value: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::tool_descriptors::Entity")]
ToolDescriptors,
}
impl Related<super::tool_descriptors::Entity> for Entity {
fn to() -> RelationDef {
Relation::ToolDescriptors.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,36 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "memory_items")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub namespace_id: String,
pub title: String,
#[sea_orm(column_type = "Text")]
pub content: String,
pub source: String,
pub index_status: String,
pub index_error: Option<String>,
pub updated_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::memory_namespaces::Entity",
from = "Column::NamespaceId",
to = "super::memory_namespaces::Column::Id",
on_delete = "Cascade"
)]
MemoryNamespace,
}
impl Related<super::memory_namespaces::Entity> for Entity {
fn to() -> RelationDef {
Relation::MemoryNamespace.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,32 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "memory_namespaces")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
pub scope: String,
pub embedding_provider: Option<String>,
pub embedding_dimensions: Option<i32>,
pub retrieval_threshold: Option<f32>,
pub retrieval_top_k: Option<i32>,
pub icon_type: Option<String>,
pub icon_value: Option<String>,
pub sort_order: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::memory_items::Entity")]
MemoryItems,
}
impl Related<super::memory_items::Entity> for Entity {
fn to() -> RelationDef {
Relation::MemoryItems.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,48 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "messages")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub role: String,
pub content: String,
pub provider_id: Option<String>,
pub model_id: Option<String>,
pub token_count: Option<i64>,
pub prompt_tokens: Option<i64>,
pub completion_tokens: Option<i64>,
pub attachments: String,
pub thinking: Option<String>,
pub created_at: i64,
pub branch_id: Option<String>,
pub parent_message_id: Option<String>,
pub version_index: i32,
pub is_active: i32,
pub tool_calls_json: Option<String>,
pub tool_call_id: Option<String>,
#[sea_orm(default_value = "complete")]
pub status: String,
pub tokens_per_second: Option<f64>,
pub first_token_latency_ms: Option<i64>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::conversations::Entity",
from = "Column::ConversationId",
to = "super::conversations::Column::Id"
)]
Conversation,
}
impl Related<super::conversations::Entity> for Entity {
fn to() -> RelationDef {
Relation::Conversation.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,44 @@
//! SeaORM entity definitions for BoosApi database tables.
pub mod builtin_model_deletions;
pub mod conversation_categories;
pub mod conversation_summaries;
pub mod conversations;
pub mod desktop_state;
pub mod drawing_generations;
pub mod drawing_images;
pub mod gateway_diagnostics;
pub mod gateway_keys;
pub mod gateway_request_logs;
pub mod gateway_usage;
pub mod mcp_servers;
pub mod messages;
pub mod models;
pub mod program_policies;
pub mod provider_keys;
pub mod providers;
pub mod search_citations;
pub mod search_providers;
pub mod settings;
pub mod skill_states;
pub mod tool_descriptors;
pub mod tool_executions;
// Wave 2+ entities
pub mod artifacts;
pub mod backup_manifests;
pub mod backup_targets;
pub mod context_sources;
pub mod conversation_branches;
pub mod import_jobs;
pub mod knowledge_bases;
pub mod knowledge_documents;
pub mod memory_items;
pub mod memory_namespaces;
pub mod retrieval_hits;
pub mod stored_files;
pub mod agent_sessions;
pub use sea_orm;
@@ -0,0 +1,36 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "models")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub provider_id: String,
#[sea_orm(primary_key, auto_increment = false)]
pub model_id: String,
pub name: String,
pub group_name: Option<String>,
pub model_type: String,
pub capabilities: String,
pub max_tokens: Option<i64>,
pub enabled: i32,
pub param_overrides: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::providers::Entity",
from = "Column::ProviderId",
to = "super::providers::Column::Id"
)]
Provider,
}
impl Related<super::providers::Entity> for Entity {
fn to() -> RelationDef {
Relation::Provider.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,21 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "program_policies")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
#[sea_orm(unique)]
pub program_name: String,
pub allowed_provider_ids_json: String,
pub allowed_model_ids_json: String,
pub default_provider_id: Option<String>,
pub default_model_id: Option<String>,
pub rate_limit_per_minute: Option<i64>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,35 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "provider_keys")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub provider_id: String,
pub key_encrypted: String,
pub key_prefix: String,
pub enabled: i32,
pub last_validated_at: Option<i64>,
pub last_error: Option<String>,
pub rotation_index: i32,
pub created_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::providers::Entity",
from = "Column::ProviderId",
to = "super::providers::Column::Id"
)]
Provider,
}
impl Related<super::providers::Entity> for Entity {
fn to() -> RelationDef {
Relation::Provider.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,43 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "providers")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
pub provider_type: String,
pub api_host: String,
pub api_path: Option<String>,
pub enabled: i32,
pub proxy_config: Option<String>,
pub custom_headers: Option<String>,
pub icon: Option<String>,
pub builtin_id: Option<String>,
pub sort_order: i32,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::provider_keys::Entity")]
ProviderKeys,
#[sea_orm(has_many = "super::models::Entity")]
Models,
}
impl Related<super::provider_keys::Entity> for Entity {
fn to() -> RelationDef {
Relation::ProviderKeys.def()
}
}
impl Related<super::models::Entity> for Entity {
fn to() -> RelationDef {
Relation::Models.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,22 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "retrieval_hits")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub message_id: String,
pub knowledge_base_id: String,
pub document_id: String,
pub chunk_ref: String,
pub score: f64,
#[sea_orm(column_type = "Text")]
pub preview: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,21 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "search_citations")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub message_id: String,
pub title: String,
pub url: String,
pub snippet: Option<String>,
pub provider_id: String,
pub rank: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,24 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "search_providers")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub name: String,
pub provider_type: String,
pub endpoint: Option<String>,
pub api_key_ref: Option<String>,
pub enabled: i32,
pub region: Option<String>,
pub language: Option<String>,
pub safe_search: Option<i64>,
pub result_limit: i32,
pub timeout_ms: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,16 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "settings")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub key: String,
#[sea_orm(column_type = "Text")]
pub value: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,16 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "skill_states")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub name: String,
pub enabled: i32,
pub updated_at: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,21 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "stored_files")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub hash: String,
pub original_name: String,
pub mime_type: String,
pub size_bytes: i64,
pub storage_path: String,
pub conversation_id: Option<String>,
pub created_at: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,31 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "tool_descriptors")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub server_id: String,
pub name: String,
pub description: Option<String>,
pub input_schema_json: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::mcp_servers::Entity",
from = "Column::ServerId",
to = "super::mcp_servers::Column::Id"
)]
McpServers,
}
impl Related<super::mcp_servers::Entity> for Entity {
fn to() -> RelationDef {
Relation::McpServers.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,25 @@
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "tool_executions")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
pub conversation_id: String,
pub message_id: Option<String>,
pub server_id: String,
pub tool_name: String,
pub status: String,
pub input_preview: Option<String>,
pub output_preview: Option<String>,
pub error_message: Option<String>,
pub duration_ms: Option<i64>,
pub created_at: String,
pub approval_status: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
@@ -0,0 +1,39 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum BoosApiError {
#[error("Database error: {0}")]
Database(#[from] sea_orm::DbErr),
#[error("Provider error: {0}")]
Provider(String),
#[error("Gateway error: {0}")]
Gateway(String),
#[error("Crypto error: {0}")]
Crypto(String),
#[error("Not found: {0}")]
NotFound(String),
#[error("Validation error: {0}")]
Validation(String),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
}
impl serde::Serialize for BoosApiError {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
impl From<sea_orm::TransactionError<sea_orm::DbErr>> for BoosApiError {
fn from(err: sea_orm::TransactionError<sea_orm::DbErr>) -> Self {
match err {
sea_orm::TransactionError::Connection(e) => BoosApiError::Database(e),
sea_orm::TransactionError::Transaction(e) => BoosApiError::Database(e),
}
}
}
pub type Result<T> = std::result::Result<T, BoosApiError>;
@@ -0,0 +1,91 @@
use std::path::PathBuf;
use sha2::{Digest, Sha256};
use crate::error::Result;
pub struct FileStore {
base_dir: PathBuf,
}
pub struct SavedFile {
pub hash: String,
/// Relative path from the documents root (e.g. "images/abc123_photo.jpg")
pub storage_path: String,
pub size_bytes: i64,
}
impl FileStore {
/// Creates a FileStore rooted at `~/Documents/boosapi/`.
pub fn new() -> Self {
Self {
base_dir: crate::storage_paths::documents_root(),
}
}
/// Creates a FileStore with an explicit root directory (useful for testing).
pub fn with_root(root: PathBuf) -> Self {
Self { base_dir: root }
}
/// Save file bytes to disk. Returns hash and relative storage path.
/// Files are stored under `{base_dir}/{bucket}/{hash_prefix}_{sanitized_name}`
/// where bucket is determined by MIME type ("images" or "files").
pub fn save_file(
&self,
data: &[u8],
original_name: &str,
mime_type: &str,
) -> Result<SavedFile> {
let hash = Self::compute_hash(data);
let relative_path =
crate::storage_paths::build_relative_path(original_name, mime_type, &hash);
let abs_path = self.base_dir.join(&relative_path);
if let Some(parent) = abs_path.parent() {
std::fs::create_dir_all(parent)?;
}
// Deduplication: skip write if file already exists with same hash
if !abs_path.exists() {
std::fs::write(&abs_path, data)?;
}
Ok(SavedFile {
hash,
storage_path: relative_path,
size_bytes: data.len() as i64,
})
}
/// Read file bytes from a relative storage path.
pub fn read_file(&self, storage_path: &str) -> Result<Vec<u8>> {
let path = self.resolve_path(storage_path);
if !path.exists() {
return Err(crate::error::BoosApiError::NotFound(format!(
"File not found: {}",
storage_path
)));
}
Ok(std::fs::read(&path)?)
}
/// Delete a file from storage.
pub fn delete_file(&self, storage_path: &str) -> Result<()> {
let path = self.resolve_path(storage_path);
if path.exists() {
std::fs::remove_file(&path)?;
}
Ok(())
}
pub fn resolve_path(&self, storage_path: &str) -> PathBuf {
self.base_dir.join(storage_path)
}
fn compute_hash(data: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(data);
format!("{:x}", hasher.finalize())
}
}
+22
View File
@@ -0,0 +1,22 @@
pub mod builtin_tools;
pub mod crypto;
pub mod db;
pub mod document_parser;
pub mod entity;
pub mod error;
pub mod file_store;
pub mod mcp_client;
pub mod path_vars;
pub mod rag;
pub mod repo;
pub mod search;
pub mod s3_backup;
pub mod storage_inventory;
pub mod storage_migration;
pub mod storage_paths;
pub mod text_chunker;
pub mod token_counter;
pub mod types;
pub mod utils;
pub mod vector_store;
pub mod webdav;
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,339 @@
//! Path encoding/decoding for cross-device portability.
//!
//! Replaces absolute path prefixes with variables (`{{BOOSAPI_HOME}}`, `{{DOCUMENTS}}`,
//! `{{HOME}}`) when writing to the database, and resolves them back to platform-specific
//! absolute paths when reading. This allows backups and WebDAV syncs to be restored on
//! a different machine without hard-coded user paths breaking.
const VAR_BOOSAPI_HOME: &str = "{{BOOSAPI_HOME}}";
const VAR_DOCUMENTS: &str = "{{DOCUMENTS}}";
const VAR_HOME: &str = "{{HOME}}";
/// Encode an absolute path by replacing known prefixes with variables.
/// More-specific prefixes (BOOSAPI_HOME, DOCUMENTS) are tried before HOME.
pub fn encode_path(absolute_path: &str) -> String {
if absolute_path.is_empty() {
return absolute_path.to_string();
}
let home = match dirs::home_dir() {
Some(h) => h,
None => return absolute_path.to_string(),
};
let boosapi_home = home.join(".boosapi");
let documents_root = crate::storage_paths::documents_root();
// Try the most specific prefix first so that e.g. ~/.boosapi/… is not
// encoded as {{HOME}}/.boosapi/… when {{BOOSAPI_HOME}}/… is more precise.
if let Some(encoded) = try_encode(absolute_path, &boosapi_home.to_string_lossy(), VAR_BOOSAPI_HOME)
{
return encoded;
}
if let Some(encoded) =
try_encode(absolute_path, &documents_root.to_string_lossy(), VAR_DOCUMENTS)
{
return encoded;
}
if let Some(encoded) = try_encode(absolute_path, &home.to_string_lossy(), VAR_HOME) {
return encoded;
}
absolute_path.to_string()
}
/// Decode a path by replacing variables with the current system's actual paths.
pub fn decode_path(encoded_path: &str) -> String {
if encoded_path.is_empty() {
return encoded_path.to_string();
}
let home = match dirs::home_dir() {
Some(h) => h,
None => return encoded_path.to_string(),
};
let boosapi_home = home.join(".boosapi");
let documents_root = crate::storage_paths::documents_root();
if let Some(rest) = encoded_path.strip_prefix(VAR_BOOSAPI_HOME) {
return format!(
"{}{}",
boosapi_home.to_string_lossy(),
platform_sep(rest)
);
}
if let Some(rest) = encoded_path.strip_prefix(VAR_DOCUMENTS) {
return format!(
"{}{}",
documents_root.to_string_lossy(),
platform_sep(rest)
);
}
if let Some(rest) = encoded_path.strip_prefix(VAR_HOME) {
return format!("{}{}", home.to_string_lossy(), platform_sep(rest));
}
encoded_path.to_string()
}
/// Encode an `Option<String>` path.
pub fn encode_path_opt(path: &Option<String>) -> Option<String> {
path.as_ref().map(|p| encode_path(p))
}
/// Decode an `Option<String>` path.
pub fn decode_path_opt(path: &Option<String>) -> Option<String> {
path.as_ref().map(|p| decode_path(p))
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// Try to replace `prefix` at the start of `path` with `variable`.
fn try_encode(path: &str, prefix: &str, variable: &str) -> Option<String> {
let np = normalize_sep_fwd(path);
let npfx = normalize_sep_fwd(prefix)
.trim_end_matches('/')
.to_string();
if np == npfx {
return Some(variable.to_string());
}
if np.starts_with(&npfx) {
let rest = &np[npfx.len()..];
if rest.starts_with('/') {
return Some(format!("{}{}", variable, rest));
}
}
None
}
/// Normalize all path separators to forward slashes (for comparison / storage).
fn normalize_sep_fwd(s: &str) -> String {
s.replace('\\', "/")
}
/// Convert separators in `s` to the current platform's native separator.
#[cfg(windows)]
fn platform_sep(s: &str) -> String {
s.replace('/', "\\")
}
#[cfg(not(windows))]
fn platform_sep(s: &str) -> String {
s.replace('\\', "/")
}
// ---------------------------------------------------------------------------
// Database migration for hardcoded paths
// ---------------------------------------------------------------------------
/// Settings keys that store filesystem paths.
const PATH_SETTING_KEYS: &[&str] = &[
"backup_dir",
"gateway_ssl_cert_path",
"gateway_ssl_key_path",
"agent_workspace_root",
];
/// Migrate hardcoded absolute paths in settings to use dynamic variables.
/// Called once at startup. Only touches values that look like absolute paths
/// and do NOT already contain a `{{…}}` variable.
pub async fn migrate_hardcoded_paths(db: &sea_orm::DatabaseConnection) {
for &key in PATH_SETTING_KEYS {
match crate::repo::settings::get_setting(db, key).await {
Ok(Some(value)) if !value.is_empty() && !value.contains("{{") => {
let encoded = encode_path(&value);
if encoded != value {
if let Err(e) = crate::repo::settings::set_setting(db, key, &encoded).await {
tracing::warn!(
"path_vars: failed to migrate setting '{}': {}",
key,
e
);
} else {
tracing::info!(
"path_vars: migrated '{}': {} → {}",
key,
value,
encoded
);
}
}
}
_ => {}
}
}
// Also migrate backup_manifests.file_path entries
migrate_backup_manifest_paths(db).await;
}
/// Migrate hardcoded file_path values in backup_manifests table.
async fn migrate_backup_manifest_paths(db: &sea_orm::DatabaseConnection) {
use crate::entity::backup_manifests;
use sea_orm::*;
let manifests = match backup_manifests::Entity::find().all(db).await {
Ok(m) => m,
Err(_) => return,
};
for m in manifests {
if let Some(ref fp) = m.file_path {
if !fp.is_empty() && !fp.contains("{{") {
let encoded = encode_path(fp);
if encoded != *fp {
let mut am: backup_manifests::ActiveModel = m.into();
am.file_path = Set(Some(encoded));
if let Err(e) = am.update(db).await {
tracing::warn!(
"path_vars: failed to migrate backup manifest path: {}",
e
);
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn roundtrip_boosapi_home() {
let home = dirs::home_dir().unwrap();
let original = home
.join(".boosapi")
.join("ssl")
.join("cert.pem")
.to_string_lossy()
.to_string();
let encoded = encode_path(&original);
assert!(
encoded.starts_with(VAR_BOOSAPI_HOME),
"expected BOOSAPI_HOME prefix, got: {}",
encoded
);
assert_eq!(decode_path(&encoded), original);
}
#[test]
fn roundtrip_documents() {
let docs = dirs::document_dir().unwrap().join("boosapi");
let original = docs
.join("images")
.join("photo.jpg")
.to_string_lossy()
.to_string();
let encoded = encode_path(&original);
assert!(
encoded.starts_with(VAR_DOCUMENTS),
"expected DOCUMENTS prefix, got: {}",
encoded
);
assert_eq!(decode_path(&encoded), original);
}
#[test]
fn roundtrip_home() {
let home = dirs::home_dir().unwrap();
let original = home
.join("some")
.join("random")
.join("file.txt")
.to_string_lossy()
.to_string();
let encoded = encode_path(&original);
assert!(
encoded.starts_with(VAR_HOME),
"expected HOME prefix, got: {}",
encoded
);
assert_eq!(decode_path(&encoded), original);
}
#[test]
fn boosapi_home_takes_priority_over_home() {
let home = dirs::home_dir().unwrap();
let original = home
.join(".boosapi")
.join("backups")
.to_string_lossy()
.to_string();
let encoded = encode_path(&original);
assert!(
encoded.starts_with(VAR_BOOSAPI_HOME),
"BOOSAPI_HOME should take priority over HOME, got: {}",
encoded
);
}
#[test]
fn unrelated_path_unchanged() {
let path = "/opt/data/file.txt";
assert_eq!(encode_path(path), path);
assert_eq!(decode_path(path), path);
}
#[test]
fn empty_path_unchanged() {
assert_eq!(encode_path(""), "");
assert_eq!(decode_path(""), "");
}
#[test]
fn option_helpers_none() {
assert_eq!(encode_path_opt(&None), None);
assert_eq!(decode_path_opt(&None), None);
}
#[test]
fn option_helpers_some() {
let home = dirs::home_dir().unwrap();
let path = home
.join(".boosapi")
.join("backups")
.to_string_lossy()
.to_string();
let encoded = encode_path_opt(&Some(path.clone()));
assert!(encoded.as_ref().unwrap().starts_with(VAR_BOOSAPI_HOME));
let decoded = decode_path_opt(&encoded);
assert_eq!(decoded, Some(path));
}
#[test]
fn exact_prefix_without_trailing_component() {
let home = dirs::home_dir().unwrap();
let exact = home.join(".boosapi").to_string_lossy().to_string();
let encoded = encode_path(&exact);
assert_eq!(encoded, VAR_BOOSAPI_HOME);
assert_eq!(decode_path(&encoded), exact);
}
#[test]
fn already_encoded_path_is_decoded() {
let decoded = decode_path("{{BOOSAPI_HOME}}/ssl/cert.pem");
let home = dirs::home_dir().unwrap();
let expected = home
.join(".boosapi")
.join("ssl")
.join("cert.pem")
.to_string_lossy()
.to_string();
assert_eq!(decoded, expected);
}
#[test]
fn path_without_variable_left_unchanged_by_decode() {
let home = dirs::home_dir().unwrap();
let abs = home
.join(".boosapi")
.join("foo")
.to_string_lossy()
.to_string();
assert_eq!(decode_path(&abs), abs);
}
}
+650
View File
@@ -0,0 +1,650 @@
//! Unified RAG (Retrieval-Augmented Generation) abstraction layer.
//!
//! Provides a trait-based interface for different RAG sources (knowledge bases,
//! memory namespaces, etc.) to share indexing, searching, and context-collection
//! logic without code duplication.
use async_trait::async_trait;
use sea_orm::DatabaseConnection;
use crate::error::{BoosApiError, Result};
use crate::types::{
RagContextResult, RagRetrievedItem, RagSourceEmptyResult, RagSourceError, RagSourceResult,
};
use crate::vector_store::{EmbeddingRecord, VectorSearchResult, VectorStore};
use crate::{document_parser, text_chunker};
fn format_rag_failure_message(reason: impl AsRef<str>) -> String {
const PREFIX: &str = "检索失败";
let reason = reason.as_ref().trim();
if reason.is_empty() {
return PREFIX.to_string();
}
if reason.starts_with(PREFIX) {
return reason.to_string();
}
format!("{PREFIX}:{reason}")
}
// ── Trait ────────────────────────────────────────────────────────────────────
/// A source of RAG content that can be searched and indexed.
///
/// Each implementor describes how to look up its embedding provider and
/// what prefix / label to use for vector-store collections and conversation
/// context injection.
#[async_trait]
pub trait RAGSource: Send + Sync {
/// Collection prefix for vector-store table names (e.g. `"kb"`, `"mem"`).
fn collection_prefix(&self) -> &'static str;
/// Human-readable label inserted into conversation context
/// (e.g. `"Knowledge Base Reference"`, `"Memory Reference"`).
fn context_label(&self) -> &'static str;
/// Resolve the `"providerId::modelId"` embedding provider string
/// configured on the container identified by `container_id`.
async fn resolve_embedding_provider(
&self,
db: &DatabaseConnection,
container_id: &str,
) -> Result<String>;
}
// ── Built-in implementations ─────────────────────────────────────────────────
/// RAG source backed by a knowledge base (documents → parsed → chunked → embedded).
pub struct KnowledgeRAG;
#[async_trait]
impl RAGSource for KnowledgeRAG {
fn collection_prefix(&self) -> &'static str {
"kb"
}
fn context_label(&self) -> &'static str {
"Knowledge Base Reference"
}
async fn resolve_embedding_provider(
&self,
db: &DatabaseConnection,
container_id: &str,
) -> Result<String> {
let kb = crate::repo::knowledge::get_knowledge_base(db, container_id).await?;
kb.embedding_provider.ok_or_else(|| {
BoosApiError::Provider("Knowledge base has no embedding provider configured".to_string())
})
}
}
/// RAG source backed by a memory namespace (text items → directly embedded).
pub struct MemoryRAG;
#[async_trait]
impl RAGSource for MemoryRAG {
fn collection_prefix(&self) -> &'static str {
"mem"
}
fn context_label(&self) -> &'static str {
"Memory Reference"
}
async fn resolve_embedding_provider(
&self,
db: &DatabaseConnection,
container_id: &str,
) -> Result<String> {
let ns = crate::repo::memory::get_namespace(db, container_id).await?;
ns.embedding_provider.ok_or_else(|| {
BoosApiError::Provider(
"Memory namespace has no embedding provider configured".to_string(),
)
})
}
}
// ── Helpers ──────────────────────────────────────────────────────────────────
/// Build the sanitised collection ID for a RAG source.
pub fn collection_id(prefix: &str, container_id: &str) -> String {
format!("{}_{}", prefix, container_id)
}
// ── Unified search ───────────────────────────────────────────────────────────
/// Search a single RAG source for content relevant to `query`.
///
/// This is the generic replacement for the separate `search_knowledge` /
/// `search_memory` functions. The concrete `EmbedFn` is injected by the
/// caller (typically `crate::indexing::generate_embeddings`).
pub async fn search<S: RAGSource + ?Sized>(
source: &S,
db: &DatabaseConnection,
master_key: &[u8; 32],
vector_store: &VectorStore,
container_id: &str,
query: &str,
top_k: usize,
dimensions: Option<usize>,
embed_fn: impl AsyncEmbedFn,
) -> Result<Vec<VectorSearchResult>> {
let embedding_provider = source.resolve_embedding_provider(db, container_id).await?;
let cid = collection_id(source.collection_prefix(), container_id);
let embed_response = embed_fn
.generate(
db,
master_key,
&embedding_provider,
vec![query.to_string()],
dimensions,
)
.await?;
let query_embedding = embed_response
.embeddings
.into_iter()
.next()
.ok_or_else(|| BoosApiError::Provider("No query embedding returned".into()))?;
let results = vector_store.search(&cid, query_embedding, top_k).await?;
Ok(results)
}
// ── Unified indexing ─────────────────────────────────────────────────────────
/// Chunking strategy for indexing content.
pub enum ChunkStrategy {
/// Parse a file and chunk the resulting text.
ParseAndChunk {
source_path: String,
mime_type: String,
chunk_size: usize,
overlap: usize,
},
/// Embed the content directly as a single vector.
Direct,
}
/// Index content into a RAG source's vector collection.
///
/// Depending on the `ChunkStrategy`, the content is either:
/// - Parsed from a file, chunked, and batch-embedded (`ParseAndChunk`), or
/// - Embedded directly as a single item (`Direct`).
pub async fn index(
vector_store: &VectorStore,
collection_prefix: &str,
container_id: &str,
item_id: &str,
_content: &str,
embeddings: Vec<Vec<f32>>,
chunks: Vec<(String, String, i32)>, // (id, content, chunk_index)
) -> Result<()> {
if chunks.is_empty() || embeddings.is_empty() {
return Ok(());
}
if embeddings.len() != chunks.len() {
return Err(BoosApiError::Provider(format!(
"Embedding count mismatch: got {} embeddings for {} chunks",
embeddings.len(),
chunks.len()
)));
}
let cid = collection_id(collection_prefix, container_id);
let records: Vec<EmbeddingRecord> = chunks
.into_iter()
.zip(embeddings)
.map(|((id, text, chunk_index), embedding)| EmbeddingRecord {
id,
document_id: item_id.to_string(),
chunk_index,
content: text,
embedding,
})
.collect();
vector_store.upsert_embeddings(&cid, records).await
}
/// Prepare chunks from content using the given strategy.
///
/// Returns a list of `(chunk_id, chunk_content, chunk_index)` tuples.
pub fn prepare_chunks(
item_id: &str,
strategy: &ChunkStrategy,
) -> Result<Vec<(String, String, i32)>> {
match strategy {
ChunkStrategy::ParseAndChunk {
source_path,
mime_type,
chunk_size,
overlap,
} => {
let path = std::path::Path::new(source_path);
let text = document_parser::extract_text(path, mime_type)?;
if text.trim().is_empty() {
return Ok(vec![]);
}
let chunks = text_chunker::chunk_text(&text, *chunk_size, *overlap);
Ok(chunks
.into_iter()
.map(|c| (format!("{}_{}", item_id, c.index), c.content, c.index))
.collect())
}
ChunkStrategy::Direct => {
// Caller provides content directly; we don't read from strategy.
// The actual content is passed to `index()` separately.
// Return a placeholder that the caller fills in.
Ok(vec![])
}
}
}
/// Prepare a single direct chunk (for memory items).
pub fn prepare_direct_chunk(item_id: &str, content: &str) -> Vec<(String, String, i32)> {
if content.trim().is_empty() {
return vec![];
}
vec![(item_id.to_string(), content.to_string(), 0)]
}
pub fn vector_similarity(distance: f32) -> f32 {
1.0 / (1.0 + distance.max(0.0))
}
pub fn passes_retrieval_threshold(distance: f32, threshold: f32) -> bool {
if threshold <= 0.0 {
return true;
}
if threshold <= 1.0 {
vector_similarity(distance) >= threshold
} else {
distance <= threshold
}
}
fn filter_results_by_threshold(
raw_results: Vec<VectorSearchResult>,
threshold: f32,
) -> (Vec<VectorSearchResult>, Option<&'static str>) {
if raw_results.is_empty() {
return (raw_results, Some("no_candidates"));
}
if threshold <= 0.0 {
return (raw_results, None);
}
let results: Vec<_> = raw_results
.into_iter()
.filter(|r| passes_retrieval_threshold(r.score, threshold))
.collect();
if results.is_empty() {
(results, Some("threshold_filtered"))
} else {
(results, None)
}
}
// ── Context collection ───────────────────────────────────────────────────────
/// A typed RAG source reference for context collection.
pub struct RAGSourceRef {
pub source_type: RAGSourceType,
pub container_id: String,
}
/// The type of RAG source.
#[derive(PartialEq)]
pub enum RAGSourceType {
Knowledge,
Memory,
}
impl RAGSourceRef {
fn source(&self) -> Box<dyn RAGSource> {
match self.source_type {
RAGSourceType::Knowledge => Box::new(KnowledgeRAG),
RAGSourceType::Memory => Box::new(MemoryRAG),
}
}
fn source_type_str(&self) -> &'static str {
match self.source_type {
RAGSourceType::Knowledge => "knowledge",
RAGSourceType::Memory => "memory",
}
}
fn empty_result(&self, reason: &str) -> RagSourceEmptyResult {
RagSourceEmptyResult {
source_type: self.source_type_str().to_string(),
container_id: self.container_id.clone(),
reason: reason.to_string(),
}
}
}
/// Collect RAG context from all enabled sources for a conversation query.
///
/// Returns a `RagContextResult` containing both formatted context parts
/// (for injection into the system prompt) and structured results
/// (for frontend display). Errors for individual sources are logged and skipped.
pub async fn collect_rag_context(
db: &DatabaseConnection,
master_key: &[u8; 32],
vector_store: &VectorStore,
kb_ids: &[String],
mem_ids: &[String],
query: &str,
top_k: usize,
embed_fn: impl AsyncEmbedFn + Clone,
rerank_fn: impl AsyncRerankFn + Clone,
) -> RagContextResult {
let mut sources: Vec<RAGSourceRef> = Vec::new();
for id in kb_ids {
sources.push(RAGSourceRef {
source_type: RAGSourceType::Knowledge,
container_id: id.clone(),
});
}
for id in mem_ids {
sources.push(RAGSourceRef {
source_type: RAGSourceType::Memory,
container_id: id.clone(),
});
}
let mut context_parts = Vec::new();
let mut source_results = Vec::new();
let mut errors = Vec::new();
let mut empty_results = Vec::new();
for src_ref in &sources {
let source = src_ref.source();
// Resolve per-source search parameters (top_k, threshold, dimensions)
let (source_top_k, threshold, dims, final_top_k, rerank_provider) = match src_ref
.source_type
{
RAGSourceType::Memory => {
match crate::repo::memory::get_namespace(db, &src_ref.container_id).await {
Ok(ns) => {
let final_top_k = ns
.retrieval_top_k
.filter(|v| *v > 0)
.map(|v| v as usize)
.unwrap_or(top_k)
.max(1);
(
final_top_k,
ns.retrieval_threshold.unwrap_or(0.0),
ns.embedding_dimensions.map(|v| v as usize),
final_top_k,
None,
)
}
Err(_) => (top_k, 0.0, None, top_k, None),
}
}
RAGSourceType::Knowledge => {
match crate::repo::knowledge::get_knowledge_base(db, &src_ref.container_id).await {
Ok(kb) => {
let final_top_k = kb
.retrieval_top_k
.filter(|v| *v > 0)
.map(|v| v as usize)
.unwrap_or(top_k)
.max(1);
let rerank_provider = kb
.rerank_provider
.filter(|provider| !provider.trim().is_empty());
let source_top_k = if rerank_provider.is_some() {
let min_candidates = final_top_k.min(100) as i32;
kb.rerank_candidate_k
.unwrap_or(20)
.clamp(min_candidates, 100) as usize
} else {
final_top_k
};
(
source_top_k,
kb.retrieval_threshold.unwrap_or(0.0),
kb.embedding_dimensions
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0),
final_top_k,
rerank_provider,
)
}
Err(_) => (top_k, 0.0, None, top_k, None),
}
}
};
let result = search(
source.as_ref(),
db,
master_key,
vector_store,
&src_ref.container_id,
query,
source_top_k,
dims,
embed_fn.clone(),
)
.await;
match result {
Ok(raw_results) => {
let (mut results, empty_reason) =
filter_results_by_threshold(raw_results, threshold);
if let Some(reason) = empty_reason {
tracing::warn!(
"RAG search returned no usable results for {} {}: {}",
source.collection_prefix(),
src_ref.container_id,
reason,
);
empty_results.push(src_ref.empty_result(reason));
continue;
}
if let Some(rerank_provider) = rerank_provider {
match rerank_fn
.rerank(
db,
master_key,
&rerank_provider,
query,
&results,
final_top_k,
)
.await
{
Ok(reranked) if !reranked.is_empty() => {
results = reranked;
}
Ok(_) => {}
Err(e) => {
tracing::warn!(
"RAG rerank failed for {} {}: {}",
source.collection_prefix(),
src_ref.container_id,
e
);
errors.push(RagSourceError {
source_type: src_ref.source_type_str().to_string(),
container_id: src_ref.container_id.clone(),
message: format_rag_failure_message(format!("重排序失败:{e}")),
});
continue;
}
}
}
results.truncate(final_top_k);
let items: Vec<RagRetrievedItem> = results
.iter()
.map(|r| RagRetrievedItem {
content: r.content.clone(),
score: r.score,
rerank_score: r.rerank_score,
document_id: r.document_id.clone(),
id: r.id.clone(),
document_name: None,
})
.collect();
let snippets: Vec<String> = results.iter().map(|r| r.content.clone()).collect();
context_parts.push(format!(
"[{}]\n{}",
source.context_label(),
snippets.join("\n---\n")
));
source_results.push(RagSourceResult {
source_type: src_ref.source_type_str().to_string(),
container_id: src_ref.container_id.clone(),
items,
});
}
Err(e) => {
tracing::warn!(
"RAG search failed for {} {}: {}",
source.collection_prefix(),
src_ref.container_id,
e
);
errors.push(RagSourceError {
source_type: src_ref.source_type_str().to_string(),
container_id: src_ref.container_id.clone(),
message: format_rag_failure_message(format!("向量检索失败:{e}")),
});
}
}
}
// Batch-lookup document titles for knowledge sources
{
let kb_doc_ids: Vec<String> = source_results
.iter()
.filter(|s| s.source_type == "knowledge")
.flat_map(|s| s.items.iter().map(|it| it.document_id.clone()))
.collect::<std::collections::HashSet<_>>()
.into_iter()
.collect();
if !kb_doc_ids.is_empty() {
match crate::repo::knowledge::get_document_titles(db, &kb_doc_ids).await {
Ok(titles) => {
for src in source_results
.iter_mut()
.filter(|s| s.source_type == "knowledge")
{
for item in &mut src.items {
item.document_name = titles.get(&item.document_id).cloned();
}
}
}
Err(e) => {
tracing::warn!("Failed to lookup document titles: {e}");
}
}
}
}
RagContextResult {
context_parts,
source_results,
errors,
empty_results,
}
}
// ── Embed function trait ─────────────────────────────────────────────────────
/// Trait for embedding generation, allowing the RAG layer to be independent
/// of the concrete provider implementation in the `indexing` module.
#[async_trait]
pub trait AsyncEmbedFn: Send + Sync + Clone {
async fn generate(
&self,
db: &DatabaseConnection,
master_key: &[u8; 32],
embedding_provider: &str,
texts: Vec<String>,
dimensions: Option<usize>,
) -> Result<crate::types::EmbedResponse>;
}
#[async_trait]
pub trait AsyncRerankFn: Send + Sync + Clone {
async fn rerank(
&self,
db: &DatabaseConnection,
master_key: &[u8; 32],
rerank_provider: &str,
query: &str,
results: &[VectorSearchResult],
top_n: usize,
) -> Result<Vec<VectorSearchResult>>;
}
#[cfg(test)]
mod tests {
use super::*;
fn vector_result(id: &str, score: f32) -> VectorSearchResult {
VectorSearchResult {
id: id.to_string(),
document_id: format!("doc-{id}"),
chunk_index: 0,
content: format!("content {id}"),
score,
rerank_score: None,
has_embedding: true,
}
}
#[test]
fn retrieval_threshold_uses_displayed_similarity_for_zero_to_one_values() {
assert!(passes_retrieval_threshold(1.0, 0.5));
assert!(passes_retrieval_threshold(4.0, 0.1));
assert!(!passes_retrieval_threshold(4.0, 0.3));
}
#[test]
fn retrieval_threshold_keeps_distance_compatibility_for_values_above_one() {
assert!(passes_retrieval_threshold(1.5, 2.0));
assert!(!passes_retrieval_threshold(2.5, 2.0));
}
#[test]
fn retrieval_empty_reason_marks_threshold_filtered_candidates() {
let (results, empty_reason) =
filter_results_by_threshold(vec![vector_result("a", 4.0)], 0.49);
assert!(results.is_empty());
assert_eq!(empty_reason, Some("threshold_filtered"));
}
#[test]
fn retrieval_empty_reason_marks_no_candidates() {
let (results, empty_reason) = filter_results_by_threshold(Vec::new(), 0.1);
assert!(results.is_empty());
assert_eq!(empty_reason, Some("no_candidates"));
}
}
+29
View File
@@ -0,0 +1,29 @@
pub mod artifact;
pub mod backup;
pub mod chatgpt_import;
pub mod cherry_import;
pub mod kelivo_import;
pub mod cli_config;
pub mod codex_session_visibility;
pub mod context_source;
pub mod conversation;
pub mod conversation_branch;
pub mod conversation_category;
pub mod drawing;
pub mod gateway;
pub mod gateway_diagnostic;
pub mod gateway_key;
pub mod gateway_request_log;
pub mod knowledge;
pub mod mcp_server;
pub mod memory;
pub mod message;
pub mod program_policy;
pub mod provider;
pub mod provider_import;
pub mod search_provider;
pub mod settings;
pub mod skill;
pub mod stored_file;
pub mod tool_execution;
pub mod agent_session;
@@ -0,0 +1,250 @@
use crate::entity::agent_sessions;
use crate::error::{BoosApiError, Result};
use crate::types::AgentSession;
use crate::utils::gen_id;
use sea_orm::*;
use sea_orm::sea_query::Expr;
fn model_to_agent_session(model: agent_sessions::Model) -> AgentSession {
AgentSession {
id: model.id,
conversation_id: model.conversation_id,
cwd: model.cwd,
permission_mode: model.permission_mode,
runtime_status: model.runtime_status,
sdk_context_json: model.sdk_context_json,
sdk_context_backup_json: model.sdk_context_backup_json,
total_tokens: model.total_tokens,
total_cost_usd: model.total_cost_usd,
created_at: model.created_at,
updated_at: model.updated_at,
}
}
/// Upsert an agent session. If no session exists for the conversation, create one.
/// If one exists, update the provided fields.
pub async fn upsert_agent_session(
db: &DatabaseConnection,
conversation_id: &str,
cwd: Option<&str>,
permission_mode: Option<&str>,
) -> Result<AgentSession> {
let existing = agent_sessions::Entity::find()
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
if let Some(model) = existing {
let mut am: agent_sessions::ActiveModel = model.into();
if let Some(cwd) = cwd {
am.cwd = Set(Some(cwd.to_string()));
}
if let Some(pm) = permission_mode {
am.permission_mode = Set(pm.to_string());
}
am.updated_at = Set(now);
let updated = am.update(db).await?;
Ok(model_to_agent_session(updated))
} else {
let id = gen_id();
let model = agent_sessions::ActiveModel {
id: Set(id),
conversation_id: Set(conversation_id.to_string()),
cwd: Set(cwd.map(|s| s.to_string())),
permission_mode: Set(permission_mode.unwrap_or("default").to_string()),
runtime_status: Set("idle".to_string()),
sdk_context_json: Set(None),
sdk_context_backup_json: Set(None),
total_tokens: Set(0),
total_cost_usd: Set(0.0),
created_at: Set(now.clone()),
updated_at: Set(now),
};
let inserted = model.insert(db).await?;
Ok(model_to_agent_session(inserted))
}
}
/// Get agent session by conversation ID.
pub async fn get_agent_session_by_conversation_id(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Option<AgentSession>> {
let model = agent_sessions::Entity::find()
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
Ok(model.map(model_to_agent_session))
}
/// Update runtime status of an agent session.
pub async fn update_agent_session_status(
db: &DatabaseConnection,
id: &str,
runtime_status: &str,
) -> Result<()> {
let model = agent_sessions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {}", id)))?;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let mut am: agent_sessions::ActiveModel = model.into();
am.runtime_status = Set(runtime_status.to_string());
am.updated_at = Set(now);
am.update(db).await?;
Ok(())
}
/// Update the working directory of an agent session.
pub async fn update_agent_session_cwd(
db: &DatabaseConnection,
id: &str,
cwd: &str,
) -> Result<()> {
let model = agent_sessions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {}", id)))?;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let mut am: agent_sessions::ActiveModel = model.into();
am.cwd = Set(Some(cwd.to_string()));
am.updated_at = Set(now);
am.update(db).await?;
Ok(())
}
/// Update the permission mode of an agent session.
pub async fn update_agent_session_permission_mode(
db: &DatabaseConnection,
id: &str,
permission_mode: &str,
) -> Result<()> {
let model = agent_sessions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {}", id)))?;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let mut am: agent_sessions::ActiveModel = model.into();
am.permission_mode = Set(permission_mode.to_string());
am.updated_at = Set(now);
am.update(db).await?;
Ok(())
}
/// Reset all running/waiting_approval sessions to idle (for app startup recovery).
pub async fn reset_running_sessions(db: &DatabaseConnection) -> Result<u64> {
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let result = agent_sessions::Entity::update_many()
.col_expr(
agent_sessions::Column::RuntimeStatus,
Expr::value("idle"),
)
.col_expr(
agent_sessions::Column::UpdatedAt,
Expr::value(now),
)
.filter(
Condition::any()
.add(agent_sessions::Column::RuntimeStatus.eq("running"))
.add(agent_sessions::Column::RuntimeStatus.eq("waiting_approval")),
)
.exec(db)
.await?;
Ok(result.rows_affected)
}
/// Update sdk_context_json, total_tokens, and total_cost_usd after a successful query.
pub async fn update_agent_session_after_query(
db: &DatabaseConnection,
id: &str,
runtime_status: &str,
sdk_context_json: Option<&str>,
tokens_delta: i32,
cost_delta: f64,
) -> Result<()> {
let model = agent_sessions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {}", id)))?;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let mut am: agent_sessions::ActiveModel = model.clone().into();
am.runtime_status = Set(runtime_status.to_string());
if let Some(ctx) = sdk_context_json {
am.sdk_context_json = Set(Some(ctx.to_string()));
}
am.total_tokens = Set(model.total_tokens + tokens_delta);
am.total_cost_usd = Set(model.total_cost_usd + cost_delta);
am.updated_at = Set(now);
am.update(db).await?;
Ok(())
}
/// Clear the sdk_context_json for an agent session by conversation_id.
/// Called when conversation messages are cleared to prevent stale history.
pub async fn clear_sdk_context_by_conversation_id(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<()> {
let session = agent_sessions::Entity::find()
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
if let Some(model) = session {
let mut am: agent_sessions::ActiveModel = model.into();
am.sdk_context_json = Set(None);
am.sdk_context_backup_json = Set(None);
am.update(db).await?;
}
Ok(())
}
/// Backup current sdk_context_json and clear it.
/// Called when a context-clear marker is inserted.
pub async fn backup_and_clear_sdk_context_by_conversation_id(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<()> {
let session = agent_sessions::Entity::find()
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
if let Some(model) = session {
let mut am: agent_sessions::ActiveModel = model.clone().into();
am.sdk_context_backup_json = Set(model.sdk_context_json);
am.sdk_context_json = Set(None);
am.update(db).await?;
}
Ok(())
}
/// Restore sdk_context_json from backup.
/// Called when a context-clear marker is removed (undo).
pub async fn restore_sdk_context_from_backup_by_conversation_id(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<()> {
let session = agent_sessions::Entity::find()
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
if let Some(model) = session {
if model.sdk_context_backup_json.is_some() {
let mut am: agent_sessions::ActiveModel = model.clone().into();
am.sdk_context_json = Set(model.sdk_context_backup_json);
am.sdk_context_backup_json = Set(None);
am.update(db).await?;
}
}
Ok(())
}
@@ -0,0 +1,103 @@
use sea_orm::*;
use crate::entity::artifacts;
use crate::error::{BoosApiError, Result};
use crate::types::{Artifact, CreateArtifactInput, UpdateArtifactInput};
use crate::utils::gen_id;
fn now_datetime() -> String {
chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string()
}
fn model_to_artifact(m: artifacts::Model) -> Artifact {
Artifact {
id: m.id,
conversation_id: m.conversation_id,
kind: m.kind,
title: m.title,
content: m.content,
format: m.format,
pinned: m.pinned != 0,
updated_at: m.updated_at,
}
}
pub async fn list_artifacts(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<Artifact>> {
let models = artifacts::Entity::find()
.filter(artifacts::Column::ConversationId.eq(conversation_id))
.order_by_desc(artifacts::Column::UpdatedAt)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_artifact).collect())
}
pub async fn get_artifact(db: &DatabaseConnection, id: &str) -> Result<Artifact> {
let model = artifacts::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Artifact {}", id)))?;
Ok(model_to_artifact(model))
}
pub async fn create_artifact(
db: &DatabaseConnection,
input: &CreateArtifactInput,
) -> Result<Artifact> {
let id = gen_id();
let am = artifacts::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(input.conversation_id.clone()),
kind: Set(input.kind.clone()),
title: Set(input.title.clone()),
content: Set(input.content.clone()),
format: Set(input.format.clone()),
pinned: Set(0),
updated_at: Set(now_datetime()),
};
am.insert(db).await?;
get_artifact(db, &id).await
}
pub async fn update_artifact(
db: &DatabaseConnection,
id: &str,
input: &UpdateArtifactInput,
) -> Result<Artifact> {
let model = artifacts::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Artifact {}", id)))?;
let existing = model_to_artifact(model.clone());
let mut am: artifacts::ActiveModel = model.into();
am.title = Set(input.title.clone().unwrap_or(existing.title));
am.content = Set(input.content.clone().unwrap_or(existing.content));
am.format = Set(input.format.clone().unwrap_or(existing.format));
am.pinned = Set(if input.pinned.unwrap_or(existing.pinned) {
1
} else {
0
});
am.updated_at = Set(now_datetime());
am.update(db).await?;
get_artifact(db, id).await
}
pub async fn delete_artifact(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = artifacts::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("Artifact {}", id)));
}
Ok(())
}
@@ -0,0 +1,267 @@
use sea_orm::*;
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
use crate::entity::backup_manifests;
use crate::error::{BoosApiError, Result};
use crate::types::BackupManifest;
use crate::utils::gen_id;
fn model_to_manifest(m: backup_manifests::Model) -> BackupManifest {
BackupManifest {
id: m.id,
version: m.version,
created_at: m.created_at,
encrypted: m.encrypted != 0,
checksum: m.checksum,
object_counts_json: m.object_counts_json,
source_app_version: m.source_app_version,
file_path: m.file_path.as_ref().map(|p| crate::path_vars::decode_path(p)),
file_size: m.file_size,
}
}
/// Get the backup directory, using the configured path or defaulting to the BoosApi home backups dir.
pub fn resolve_backup_dir(backup_dir_setting: Option<&str>, app_data_dir: &Path) -> PathBuf {
if let Some(dir) = backup_dir_setting {
if !dir.is_empty() {
return PathBuf::from(dir);
}
}
app_data_dir.join("backups")
}
/// Ensure the backup directory exists
pub fn ensure_backup_dir(dir: &Path) -> Result<()> {
std::fs::create_dir_all(dir)
.map_err(|e| BoosApiError::Gateway(format!("Failed to create backup directory: {}", e)))
}
/// Create a real backup file (SQLite copy or JSON export)
pub async fn create_backup(
db: &DatabaseConnection,
format: &str,
backup_dir: &Path,
) -> Result<BackupManifest> {
ensure_backup_dir(backup_dir)?;
let id = gen_id();
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
let extension = match format {
"sqlite" => "db",
_ => "json",
};
let filename = format!("boosapi-backup-{}.{}", timestamp, extension);
let file_path = backup_dir.join(&filename);
match format {
"sqlite" => {
create_sqlite_backup(db, &file_path).await?;
}
_ => {
create_json_backup(db, &file_path).await?;
}
}
let file_size = std::fs::metadata(&file_path)
.map(|m| m.len() as i64)
.unwrap_or(0);
let checksum = compute_file_checksum(&file_path)?;
// Count objects for manifest
let object_counts = count_objects(db).await?;
let am = backup_manifests::ActiveModel {
id: Set(id.clone()),
version: Set(format.to_string()),
encrypted: Set(0),
checksum: Set(checksum),
object_counts_json: Set(object_counts),
source_app_version: Set(env!("CARGO_PKG_VERSION").to_string()),
file_path: Set(Some(crate::path_vars::encode_path(&file_path.to_string_lossy()))),
file_size: Set(file_size),
..Default::default()
};
am.insert(db).await?;
get_backup(db, &id).await
}
/// Create a SQLite backup using VACUUM INTO
async fn create_sqlite_backup(db: &DatabaseConnection, dest: &Path) -> Result<()> {
let dest_str = dest.to_string_lossy().to_string();
// Remove existing file if present (VACUUM INTO fails otherwise)
if dest.exists() {
std::fs::remove_file(dest).map_err(|e| {
BoosApiError::Gateway(format!("Failed to remove existing backup file: {}", e))
})?;
}
db.execute(Statement::from_string(
sea_orm::DatabaseBackend::Sqlite,
format!("VACUUM INTO '{}'", dest_str.replace('\'', "''")),
))
.await
.map_err(|e| BoosApiError::Gateway(format!("VACUUM INTO failed: {}", e)))?;
Ok(())
}
/// Create a JSON backup by exporting all important tables
async fn create_json_backup(db: &DatabaseConnection, dest: &Path) -> Result<()> {
use crate::entity::*;
let conversations = conversations::Entity::find().all(db).await?;
let messages = messages::Entity::find().all(db).await?;
let providers = providers::Entity::find().all(db).await?;
let provider_keys = provider_keys::Entity::find().all(db).await?;
let models = models::Entity::find().all(db).await?;
let settings = settings::Entity::find().all(db).await?;
let gateway_keys = gateway_keys::Entity::find().all(db).await?;
let data = serde_json::json!({
"version": env!("CARGO_PKG_VERSION"),
"exported_at": chrono::Utc::now().to_rfc3339(),
"tables": {
"conversations": conversations,
"messages": messages,
"providers": providers,
"provider_keys": provider_keys,
"models": models,
"settings": settings,
"gateway_keys": gateway_keys,
}
});
let json_str = serde_json::to_string_pretty(&data)
.map_err(|e| BoosApiError::Gateway(format!("JSON serialization failed: {}", e)))?;
std::fs::write(dest, json_str)
.map_err(|e| BoosApiError::Gateway(format!("Failed to write backup file: {}", e)))?;
Ok(())
}
fn compute_file_checksum(path: &Path) -> Result<String> {
let data = std::fs::read(path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read file for checksum: {}", e)))?;
let hash = Sha256::digest(&data);
Ok(format!("{:x}", hash))
}
async fn count_objects(db: &DatabaseConnection) -> Result<String> {
use crate::entity::*;
let conv_count = conversations::Entity::find().count(db).await.unwrap_or(0);
let msg_count = messages::Entity::find().count(db).await.unwrap_or(0);
let provider_count = providers::Entity::find().count(db).await.unwrap_or(0);
let counts = serde_json::json!({
"conversations": conv_count,
"messages": msg_count,
"providers": provider_count,
});
Ok(counts.to_string())
}
pub async fn list_backups(db: &DatabaseConnection) -> Result<Vec<BackupManifest>> {
let models = backup_manifests::Entity::find()
.order_by_desc(backup_manifests::Column::CreatedAt)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_manifest).collect())
}
pub async fn get_backup(db: &DatabaseConnection, id: &str) -> Result<BackupManifest> {
let model = backup_manifests::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("BackupManifest {}", id)))?;
Ok(model_to_manifest(model))
}
pub async fn delete_backup(db: &DatabaseConnection, id: &str) -> Result<()> {
let manifest = get_backup(db, id).await?;
// Delete the file from disk if it exists
if let Some(ref path) = manifest.file_path {
let p = Path::new(path);
if p.exists() {
std::fs::remove_file(p).ok();
}
}
let result = backup_manifests::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("BackupManifest {}", id)));
}
Ok(())
}
pub async fn batch_delete_backups(db: &DatabaseConnection, ids: &[String]) -> Result<()> {
for id in ids {
delete_backup(db, id).await?;
}
Ok(())
}
/// Restore from a SQLite backup by replacing the current database file
pub async fn restore_sqlite_backup(backup_path: &str, current_db_path: &str) -> Result<()> {
let src = Path::new(backup_path);
if !src.exists() {
return Err(BoosApiError::NotFound(format!(
"Backup file not found: {}",
backup_path
)));
}
std::fs::copy(src, current_db_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to restore backup: {}", e)))?;
Ok(())
}
/// Clean up old backups exceeding max_count (keeps most recent)
pub async fn cleanup_old_backups(db: &DatabaseConnection, max_count: u32) -> Result<u32> {
let all = list_backups(db).await?;
if all.len() <= max_count as usize {
return Ok(0);
}
let to_delete = &all[max_count as usize..];
let mut deleted = 0u32;
for backup in to_delete {
delete_backup(db, &backup.id).await?;
deleted += 1;
}
Ok(deleted)
}
#[cfg(test)]
mod tests {
use super::resolve_backup_dir;
use std::path::PathBuf;
#[test]
fn resolve_backup_dir_defaults_to_boosapi_backups_subdir() {
let boosapi_home = PathBuf::from("/Users/test/.boosapi");
assert_eq!(
resolve_backup_dir(None, &boosapi_home),
boosapi_home.join("backups")
);
assert_eq!(
resolve_backup_dir(Some(""), &boosapi_home),
boosapi_home.join("backups")
);
}
#[test]
fn resolve_backup_dir_honors_explicit_absolute_override() {
let boosapi_home = PathBuf::from("/Users/test/.boosapi");
let override_dir = PathBuf::from("/Volumes/external/boosapi-backups");
assert_eq!(
resolve_backup_dir(Some(override_dir.to_str().unwrap()), &boosapi_home),
override_dir
);
}
}
@@ -0,0 +1,780 @@
use sea_orm::*;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::Read;
use std::path::Path;
use crate::entity::{conversations, import_jobs, messages};
use crate::error::{BoosApiError, Result};
use crate::repo::settings::get_settings;
use crate::utils::{gen_id, now_ts};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ChatGptImportWarning {
pub code: String,
pub message: String,
pub source_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct ChatGptImportSummary {
pub conversation_count: u32,
pub message_count: u32,
pub skipped_empty_conversation_count: u32,
pub duplicate_conversation_count: u32,
pub warnings: Vec<ChatGptImportWarning>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct ChatGptImportResult {
pub imported_conversation_count: u32,
pub imported_message_count: u32,
pub skipped_duplicate_conversation_count: u32,
pub warnings: Vec<ChatGptImportWarning>,
}
#[derive(Debug, Clone, Deserialize, Default)]
struct ChatGptConversation {
#[serde(default)]
id: Option<String>,
#[serde(default)]
title: Option<String>,
#[serde(default)]
create_time: Option<Value>,
#[serde(default)]
update_time: Option<Value>,
#[serde(default)]
current_node: Option<String>,
#[serde(default)]
mapping: HashMap<String, ChatGptNode>,
}
#[derive(Debug, Clone, Deserialize, Default)]
struct ChatGptNode {
#[serde(default)]
id: Option<String>,
#[serde(default)]
message: Option<ChatGptMessage>,
#[serde(default)]
parent: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Default)]
struct ChatGptMessage {
#[serde(default)]
id: Option<String>,
#[serde(default)]
author: Option<ChatGptAuthor>,
#[serde(default)]
create_time: Option<Value>,
#[serde(default)]
content: Option<ChatGptContent>,
#[serde(default)]
status: Option<String>,
#[serde(default)]
metadata: HashMap<String, Value>,
}
#[derive(Debug, Clone, Deserialize, Default)]
struct ChatGptAuthor {
#[serde(default)]
role: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Default)]
struct ChatGptContent {
#[serde(default)]
content_type: Option<String>,
#[serde(default)]
parts: Vec<Value>,
#[serde(default)]
text: Option<String>,
}
struct MaterializedConversation {
id: String,
title: String,
created_at: i64,
updated_at: i64,
messages: Vec<MaterializedMessage>,
warnings: Vec<ChatGptImportWarning>,
}
struct MaterializedMessage {
id: String,
role: String,
content: String,
created_at: i64,
status: String,
}
pub async fn scan_chatgpt_import_from_path(
db: &DatabaseConnection,
path: &Path,
) -> Result<ChatGptImportSummary> {
let conversations = parse_chatgpt_export(path)?;
let mut summary = ChatGptImportSummary::default();
for (index, conversation) in conversations.iter().enumerate() {
let materialized = materialize_conversation(conversation, index);
summary.warnings.extend(materialized.warnings);
if materialized.messages.is_empty() {
summary.skipped_empty_conversation_count += 1;
continue;
}
summary.conversation_count += 1;
summary.message_count += materialized.messages.len() as u32;
if conversations::Entity::find_by_id(&materialized.id)
.one(db)
.await?
.is_some()
{
summary.duplicate_conversation_count += 1;
}
}
Ok(summary)
}
pub async fn import_chatgpt_export_from_path(
db: &DatabaseConnection,
path: &Path,
) -> Result<ChatGptImportResult> {
let conversations = parse_chatgpt_export(path)?;
let settings = get_settings(db).await.unwrap_or_default();
let fallback_provider_id = settings
.default_provider_id
.clone()
.unwrap_or_else(|| "chatgpt-export".to_string());
let fallback_model_id = settings
.default_model_id
.clone()
.unwrap_or_else(|| "unknown-model".to_string());
let txn = db.begin().await?;
let mut result = ChatGptImportResult::default();
for (index, conversation) in conversations.iter().enumerate() {
let materialized = materialize_conversation(conversation, index);
result.warnings.extend(materialized.warnings);
if materialized.messages.is_empty() {
continue;
}
if conversations::Entity::find_by_id(&materialized.id)
.one(&txn)
.await?
.is_some()
{
result.skipped_duplicate_conversation_count += 1;
continue;
}
conversations::ActiveModel {
id: Set(materialized.id.clone()),
title: Set(materialized.title),
model_id: Set(fallback_model_id.clone()),
provider_id: Set(fallback_provider_id.clone()),
system_prompt: Set(None),
temperature: Set(None),
max_tokens: Set(None),
top_p: Set(None),
frequency_penalty: Set(None),
search_enabled: Set(0),
search_provider_id: Set(None),
thinking_budget: Set(None),
thinking_level: Set(None),
enabled_mcp_server_ids: Set("[]".to_string()),
enabled_knowledge_base_ids: Set("[]".to_string()),
enabled_memory_namespace_ids: Set("[]".to_string()),
message_count: Set(materialized.messages.len() as i32),
created_at: Set(materialized.created_at),
updated_at: Set(materialized.updated_at),
is_pinned: Set(0),
is_archived: Set(0),
workspace_snapshot_json: Set("{}".to_string()),
active_branch_id: Set(None),
active_artifact_id: Set(None),
research_mode: Set(0),
context_compression: Set(0),
category_id: Set(None),
parent_conversation_id: Set(None),
mode: Set("chat".to_string()),
}
.insert(&txn)
.await?;
result.imported_conversation_count += 1;
let mut previous_message_id: Option<String> = None;
for message in materialized.messages {
messages::ActiveModel {
id: Set(message.id.clone()),
conversation_id: Set(materialized.id.clone()),
role: Set(message.role),
content: Set(message.content),
provider_id: Set(Some(fallback_provider_id.clone())),
model_id: Set(Some(fallback_model_id.clone())),
token_count: Set(None),
prompt_tokens: Set(None),
completion_tokens: Set(None),
attachments: Set("[]".to_string()),
thinking: Set(None),
created_at: Set(message.created_at),
branch_id: Set(None),
parent_message_id: Set(previous_message_id.clone()),
version_index: Set(0),
is_active: Set(1),
tool_calls_json: Set(None),
tool_call_id: Set(None),
status: Set(message.status),
tokens_per_second: Set(None),
first_token_latency_ms: Set(None),
}
.insert(&txn)
.await?;
previous_message_id = Some(message.id);
result.imported_message_count += 1;
}
}
import_jobs::ActiveModel {
id: Set(gen_id()),
source_type: Set("chatgpt_export".to_string()),
status: Set("success".to_string()),
summary_json: Set(Some(serde_json::to_string(&result).unwrap_or_default())),
conflict_count: Set(result.skipped_duplicate_conversation_count as i32),
created_at: Set(chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string()),
}
.insert(&txn)
.await?;
txn.commit().await?;
Ok(result)
}
fn parse_chatgpt_export(path: &Path) -> Result<Vec<ChatGptConversation>> {
let bytes = if path
.extension()
.and_then(|value| value.to_str())
.is_some_and(|ext| ext.eq_ignore_ascii_case("zip"))
{
read_conversations_json_from_zip(path)?
} else {
std::fs::read(path)?
};
let value = serde_json::from_slice::<Value>(&bytes).map_err(|e| {
BoosApiError::Validation(format!("Invalid ChatGPT conversations.json: {e}"))
})?;
let conversations_value = value.get("conversations").cloned().unwrap_or(value);
serde_json::from_value::<Vec<ChatGptConversation>>(conversations_value).map_err(|e| {
BoosApiError::Validation(format!("Invalid ChatGPT conversations.json structure: {e}"))
})
}
fn read_conversations_json_from_zip(path: &Path) -> Result<Vec<u8>> {
let file = File::open(path)?;
let mut archive = zip::ZipArchive::new(file)
.map_err(|e| BoosApiError::Validation(format!("Invalid ChatGPT zip export: {e}")))?;
let mut entry = archive
.by_name("conversations.json")
.map_err(|_| BoosApiError::Validation("ChatGPT export is missing conversations.json".into()))?;
let mut data = Vec::new();
entry.read_to_end(&mut data)?;
Ok(data)
}
fn materialize_conversation(
conversation: &ChatGptConversation,
index: usize,
) -> MaterializedConversation {
let source_id = non_empty(conversation.id.as_deref())
.map(str::to_string)
.unwrap_or_else(|| format!("conversation-{index}"));
let id = prefixed_id("chatgpt", &source_id);
let created_at = parse_chatgpt_ts(conversation.create_time.as_ref()).unwrap_or_else(now_ts);
let updated_at = parse_chatgpt_ts(conversation.update_time.as_ref()).unwrap_or(created_at);
let title = non_empty(conversation.title.as_deref())
.unwrap_or("ChatGPT Chat")
.to_string();
let mut warnings = Vec::new();
let mut seen = HashSet::new();
let messages = ordered_node_ids(conversation)
.into_iter()
.filter_map(|node_id| {
let node = conversation.mapping.get(&node_id)?;
let message = node.message.as_ref()?;
let source_message_id = non_empty(message.id.as_deref())
.or_else(|| non_empty(node.id.as_deref()))
.unwrap_or(node_id.as_str());
if !seen.insert(source_message_id.to_string()) {
return None;
}
materialize_message(message, source_message_id, &mut warnings)
})
.collect();
MaterializedConversation {
id,
title,
created_at,
updated_at,
messages,
warnings,
}
}
fn ordered_node_ids(conversation: &ChatGptConversation) -> Vec<String> {
if let Some(current_node) = non_empty(conversation.current_node.as_deref()) {
let mut ids = Vec::new();
let mut seen = HashSet::new();
let mut cursor = Some(current_node.to_string());
while let Some(node_id) = cursor {
if !seen.insert(node_id.clone()) {
break;
}
let Some(node) = conversation.mapping.get(&node_id) else {
break;
};
ids.push(node_id);
cursor = node.parent.clone().and_then(|parent| non_empty(Some(&parent)).map(str::to_string));
}
ids.reverse();
if ids
.iter()
.any(|id| conversation.mapping.get(id).and_then(|node| node.message.as_ref()).is_some())
{
return ids;
}
}
let mut nodes = conversation
.mapping
.iter()
.filter(|(_, node)| node.message.is_some())
.map(|(id, node)| {
(
id.clone(),
node.message
.as_ref()
.and_then(|message| parse_chatgpt_ts(message.create_time.as_ref()))
.unwrap_or(0),
)
})
.collect::<Vec<_>>();
nodes.sort_by(|left, right| left.1.cmp(&right.1).then_with(|| left.0.cmp(&right.0)));
nodes.into_iter().map(|(id, _)| id).collect()
}
fn materialize_message(
message: &ChatGptMessage,
source_id: &str,
warnings: &mut Vec<ChatGptImportWarning>,
) -> Option<MaterializedMessage> {
if metadata_bool(&message.metadata, "is_visually_hidden_from_conversation")
|| metadata_bool(&message.metadata, "is_user_system_message")
{
return None;
}
let role = map_role(
message
.author
.as_ref()
.and_then(|author| author.role.as_deref()),
);
let content = content_to_markdown(message.content.as_ref(), source_id, warnings);
if content.trim().is_empty() {
return None;
}
if role == "system" && content.trim().is_empty() {
return None;
}
Some(MaterializedMessage {
id: prefixed_id("chatgpt", source_id),
role,
content,
created_at: parse_chatgpt_ts(message.create_time.as_ref()).unwrap_or_else(now_ts),
status: map_status(message.status.as_deref()),
})
}
fn content_to_markdown(
content: Option<&ChatGptContent>,
source_id: &str,
warnings: &mut Vec<ChatGptImportWarning>,
) -> String {
let Some(content) = content else {
return String::new();
};
if let Some(text) = non_empty(content.text.as_deref()) {
return text.to_string();
}
if content.parts.is_empty() {
if let Some(content_type) = non_empty(content.content_type.as_deref()) {
if content_type != "text" {
warnings.push(warning(
"unsupported_content_type",
format!("Unsupported ChatGPT content type '{content_type}' had no text parts."),
Some(source_id.to_string()),
));
}
}
return String::new();
}
content
.parts
.iter()
.filter_map(|part| part_to_markdown(part, source_id, warnings))
.collect::<Vec<_>>()
.join("\n\n")
}
fn part_to_markdown(
part: &Value,
source_id: &str,
warnings: &mut Vec<ChatGptImportWarning>,
) -> Option<String> {
match part {
Value::String(text) => non_empty(Some(text)).map(str::to_string),
Value::Object(map) => {
let part_type = map
.get("content_type")
.and_then(Value::as_str)
.or_else(|| map.get("type").and_then(Value::as_str))
.unwrap_or("object");
if matches!(part_type, "text" | "input_text" | "output_text") {
if let Some(text) = map.get("text").and_then(Value::as_str).and_then(|text| non_empty(Some(text))) {
return Some(text.to_string());
}
}
warnings.push(warning(
"unsupported_content_part",
format!("Unsupported ChatGPT content part '{part_type}' was preserved as Markdown."),
Some(source_id.to_string()),
));
Some(format!(
"### Unsupported ChatGPT content part: {part_type}\n```json\n{}\n```",
serde_json::to_string(part).unwrap_or_else(|_| part.to_string())
))
}
Value::Null => None,
other => {
warnings.push(warning(
"unsupported_content_part",
"Unsupported ChatGPT content part was preserved as Markdown.",
Some(source_id.to_string()),
));
Some(other.to_string())
}
}
}
fn parse_chatgpt_ts(value: Option<&Value>) -> Option<i64> {
let value = value?;
let seconds = match value {
Value::Number(number) => number.as_f64()?,
Value::String(text) => text
.parse::<f64>()
.ok()
.or_else(|| chrono::DateTime::parse_from_rfc3339(text).ok().map(|value| value.timestamp() as f64))?,
_ => return None,
};
Some(if seconds > 10_000_000_000.0 {
(seconds / 1000.0) as i64
} else {
seconds as i64
})
}
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
fn prefixed_id(prefix: &str, source_id: &str) -> String {
format!("{prefix}-{source_id}")
}
fn map_role(role: Option<&str>) -> String {
match role {
Some("system") => "system",
Some("user") => "user",
Some("tool") => "tool",
_ => "assistant",
}
.to_string()
}
fn map_status(status: Option<&str>) -> String {
match status.unwrap_or("") {
value if value.contains("error") || value.contains("failed") => "error",
_ => "complete",
}
.to_string()
}
fn metadata_bool(metadata: &HashMap<String, Value>, key: &str) -> bool {
metadata.get(key).and_then(Value::as_bool).unwrap_or(false)
}
fn warning(
code: impl Into<String>,
message: impl Into<String>,
source_id: Option<String>,
) -> ChatGptImportWarning {
ChatGptImportWarning {
code: code.into(),
message: message.into(),
source_id,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::create_test_pool;
use crate::entity::{conversations, messages};
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder};
use serde_json::json;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use tempfile::tempdir;
use zip::write::SimpleFileOptions;
fn write_chatgpt_zip(path: &Path, data: serde_json::Value) {
let file = File::create(path).unwrap();
let mut zip = zip::ZipWriter::new(file);
zip.start_file("conversations.json", SimpleFileOptions::default())
.unwrap();
zip.write_all(serde_json::to_string(&data).unwrap().as_bytes())
.unwrap();
zip.finish().unwrap();
}
fn sample_export_json() -> serde_json::Value {
json!([
{
"id": "conv-1",
"title": "Mainline chat",
"create_time": 1780000000.0,
"update_time": 1780000200.0,
"current_node": "msg-assistant-2",
"mapping": {
"root": {
"id": "root",
"message": null,
"parent": null,
"children": ["msg-system"]
},
"msg-system": {
"id": "msg-system",
"message": {
"id": "msg-system",
"author": {"role": "system"},
"create_time": 1780000001.0,
"content": {"content_type": "text", "parts": [""]},
"status": "finished_successfully",
"metadata": {"is_visually_hidden_from_conversation": true}
},
"parent": "root",
"children": ["msg-user-1"]
},
"msg-user-1": {
"id": "msg-user-1",
"message": {
"id": "msg-user-1",
"author": {"role": "user"},
"create_time": 1780000010.0,
"content": {"content_type": "text", "parts": ["hello"]},
"status": "finished_successfully",
"metadata": {}
},
"parent": "msg-system",
"children": ["msg-assistant-1", "branch-assistant"]
},
"branch-assistant": {
"id": "branch-assistant",
"message": {
"id": "branch-assistant",
"author": {"role": "assistant"},
"create_time": 1780000015.0,
"content": {"content_type": "text", "parts": ["old branch"]},
"status": "finished_successfully",
"metadata": {}
},
"parent": "msg-user-1",
"children": []
},
"msg-assistant-1": {
"id": "msg-assistant-1",
"message": {
"id": "msg-assistant-1",
"author": {"role": "assistant"},
"create_time": 1780000020.0,
"content": {
"content_type": "multimodal_text",
"parts": [
{"content_type": "text", "text": "hi"},
{"content_type": "image_asset_pointer", "asset_pointer": "file-service://image-1"}
]
},
"status": "finished_successfully",
"metadata": {"model_slug": "gpt-4o"}
},
"parent": "msg-user-1",
"children": ["msg-tool-1"]
},
"msg-tool-1": {
"id": "msg-tool-1",
"message": {
"id": "msg-tool-1",
"author": {"role": "tool"},
"create_time": 1780000025.0,
"content": {"content_type": "text", "parts": ["tool output"]},
"status": "finished_successfully",
"metadata": {}
},
"parent": "msg-assistant-1",
"children": ["msg-assistant-2"]
},
"msg-assistant-2": {
"id": "msg-assistant-2",
"message": {
"id": "msg-assistant-2",
"author": {"role": "assistant"},
"create_time": 1780000030.0,
"content": {"content_type": "text", "parts": ["final"]},
"status": "finished_successfully",
"metadata": {"model_slug": "gpt-4o"}
},
"parent": "msg-tool-1",
"children": []
}
}
},
{
"id": "conv-empty",
"title": "",
"create_time": null,
"update_time": null,
"current_node": null,
"mapping": {}
}
])
}
#[tokio::test]
async fn scan_chatgpt_zip_summarizes_importable_conversations() {
let dir = tempdir().unwrap();
let zip_path = dir.path().join("chatgpt.zip");
write_chatgpt_zip(&zip_path, sample_export_json());
let db = create_test_pool().await.unwrap();
let summary = scan_chatgpt_import_from_path(&db.conn, &zip_path)
.await
.unwrap();
assert_eq!(summary.conversation_count, 1);
assert_eq!(summary.message_count, 4);
assert_eq!(summary.skipped_empty_conversation_count, 1);
assert_eq!(summary.duplicate_conversation_count, 0);
assert!(summary
.warnings
.iter()
.any(|warning| warning.code == "unsupported_content_part"));
}
#[tokio::test]
async fn import_chatgpt_json_materializes_mainline_messages_and_skips_duplicates() {
let dir = tempdir().unwrap();
let json_path = dir.path().join("conversations.json");
std::fs::write(&json_path, serde_json::to_vec(&sample_export_json()).unwrap()).unwrap();
let db = create_test_pool().await.unwrap();
let result = import_chatgpt_export_from_path(&db.conn, &json_path)
.await
.unwrap();
assert_eq!(result.imported_conversation_count, 1);
assert_eq!(result.imported_message_count, 4);
assert_eq!(result.skipped_duplicate_conversation_count, 0);
let conversation = conversations::Entity::find_by_id("chatgpt-conv-1")
.one(&db.conn)
.await
.unwrap()
.unwrap();
assert_eq!(conversation.title, "Mainline chat");
assert_eq!(conversation.provider_id, "chatgpt-export");
assert_eq!(conversation.model_id, "unknown-model");
assert_eq!(conversation.message_count, 4);
assert_eq!(conversation.created_at, 1780000000);
assert_eq!(conversation.updated_at, 1780000200);
let imported_messages = messages::Entity::find()
.filter(messages::Column::ConversationId.eq("chatgpt-conv-1"))
.order_by_asc(messages::Column::CreatedAt)
.all(&db.conn)
.await
.unwrap();
let ids = imported_messages
.iter()
.map(|message| message.id.as_str())
.collect::<Vec<_>>();
assert_eq!(
ids,
vec![
"chatgpt-msg-user-1",
"chatgpt-msg-assistant-1",
"chatgpt-msg-tool-1",
"chatgpt-msg-assistant-2",
]
);
assert_eq!(imported_messages[0].role, "user");
assert_eq!(imported_messages[0].content, "hello");
assert_eq!(imported_messages[1].role, "assistant");
assert!(imported_messages[1].content.contains("hi"));
assert!(imported_messages[1]
.content
.contains("Unsupported ChatGPT content part"));
assert_eq!(
imported_messages[1].parent_message_id.as_deref(),
Some("chatgpt-msg-user-1")
);
assert_eq!(
imported_messages[2].parent_message_id.as_deref(),
Some("chatgpt-msg-assistant-1")
);
assert_eq!(
imported_messages[3].parent_message_id.as_deref(),
Some("chatgpt-msg-tool-1")
);
assert_eq!(imported_messages[1].model_id.as_deref(), Some("unknown-model"));
let duplicate = import_chatgpt_export_from_path(&db.conn, &json_path)
.await
.unwrap();
assert_eq!(duplicate.imported_conversation_count, 0);
assert_eq!(duplicate.skipped_duplicate_conversation_count, 1);
}
#[tokio::test]
async fn scan_chatgpt_zip_rejects_missing_conversations_json() {
let dir = tempdir().unwrap();
let zip_path = dir.path().join("bad.zip");
zip::ZipWriter::new(File::create(&zip_path).unwrap())
.finish()
.unwrap();
let db = create_test_pool().await.unwrap();
let err = scan_chatgpt_import_from_path(&db.conn, &zip_path)
.await
.unwrap_err();
assert!(err.to_string().contains("conversations.json"));
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,48 @@
use std::fs;
use std::path::{Path, PathBuf};
use chrono::Utc;
use serde_json::json;
use super::{
json_error, Result, SessionChange, GLOBAL_STATE_BACKUP_FILE, GLOBAL_STATE_FILE, STATE_DB_FILE,
};
pub(super) fn create_backup(codex_home: &Path, changes: &[SessionChange]) -> Result<PathBuf> {
let backup_dir = codex_home
.join("backups_state")
.join("boosapi-session-visibility")
.join(Utc::now().format("%Y%m%d-%H%M%S-%3f").to_string());
fs::create_dir_all(&backup_dir)?;
for change in changes {
copy_if_exists(&change.path, &backup_dir.join(&change.relative_path))?;
}
for file_name in [STATE_DB_FILE, GLOBAL_STATE_FILE, GLOBAL_STATE_BACKUP_FILE] {
copy_if_exists(&codex_home.join(file_name), &backup_dir.join(file_name))?;
}
write_backup_metadata(&backup_dir, changes.len())?;
Ok(backup_dir)
}
fn copy_if_exists(src: &Path, dest: &Path) -> Result<()> {
if !src.exists() {
return Ok(());
}
if let Some(parent) = dest.parent() {
fs::create_dir_all(parent)?;
}
fs::copy(src, dest)?;
Ok(())
}
fn write_backup_metadata(backup_dir: &Path, changed_files: usize) -> Result<()> {
let content = serde_json::to_string_pretty(&json!({
"version": 1,
"namespace": "boosapi-session-visibility",
"createdAt": Utc::now().to_rfc3339(),
"changedSessionFiles": changed_files
}))
.map_err(json_error)?;
fs::write(backup_dir.join("metadata.json"), content)?;
Ok(())
}
@@ -0,0 +1,150 @@
mod backup;
mod repair;
mod session_files;
mod sqlite_state;
mod sqlite_status;
mod status;
mod workspace_roots;
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use std::time::SystemTime;
use crate::error::{BoosApiError, Result};
use crate::types::CodexSessionVisibilityStatus;
pub use repair::{
repair_codex_session_visibility, repair_codex_session_visibility_at,
repair_codex_session_visibility_with_backup, repair_codex_session_visibility_with_backup_at,
};
const DEFAULT_PROVIDER: &str = "openai";
const SESSION_DIRS: [&str; 2] = ["sessions", "archived_sessions"];
const GLOBAL_STATE_FILE: &str = ".codex-global-state.json";
const GLOBAL_STATE_BACKUP_FILE: &str = ".codex-global-state.json.bak";
const STATE_DB_FILE: &str = "state_5.sqlite";
#[derive(Debug, Clone)]
pub(super) struct SessionChange {
pub path: PathBuf,
pub relative_path: PathBuf,
pub original_first_line: String,
pub updated_first_line: String,
pub separator: String,
pub offset: u64,
pub original_mtime: Option<SystemTime>,
}
#[derive(Debug, Default)]
pub(super) struct SessionScan {
pub changes: Vec<SessionChange>,
pub user_event_thread_ids: BTreeSet<String>,
pub thread_cwd_by_id: BTreeMap<String, String>,
pub encrypted_content_by_provider: BTreeMap<String, usize>,
}
#[derive(Debug, Default)]
pub(super) struct SqliteUpdateResult {
pub present: bool,
pub provider_rows: usize,
pub user_event_rows: usize,
pub cwd_rows: usize,
}
#[derive(Debug, Default)]
pub(super) struct SqliteStatus {
pub present: bool,
pub rows: usize,
pub mismatched_rows: usize,
pub user_event_rows: usize,
pub cwd_rows: usize,
pub provider_counts: BTreeMap<String, usize>,
}
impl SqliteUpdateResult {
pub(super) fn total_rows(&self) -> usize {
self.provider_rows + self.user_event_rows + self.cwd_rows
}
}
pub async fn get_codex_session_visibility_status() -> Result<CodexSessionVisibilityStatus> {
let home = default_codex_home()?;
get_codex_session_visibility_status_at(&home).await
}
pub async fn get_codex_session_visibility_status_at(
codex_home: &Path,
) -> Result<CodexSessionVisibilityStatus> {
status::get_status(codex_home)
}
fn default_codex_home() -> Result<PathBuf> {
if let Some(value) = std::env::var_os("CODEX_HOME") {
return Ok(PathBuf::from(value));
}
let home = std::env::var_os("HOME")
.or_else(|| std::env::var_os("USERPROFILE"))
.map(PathBuf::from)
.ok_or_else(|| BoosApiError::NotFound("Could not determine home directory".into()))?;
Ok(home.join(".codex"))
}
fn ensure_codex_home(codex_home: &Path) -> Result<()> {
if codex_home.exists() {
Ok(())
} else {
Err(BoosApiError::NotFound(format!(
"Codex home not found: {}",
codex_home.display()
)))
}
}
fn read_target_provider(codex_home: &Path) -> Result<String> {
let config_path = codex_home.join("config.toml");
let content = match std::fs::read_to_string(&config_path) {
Ok(content) => content,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Ok(DEFAULT_PROVIDER.into());
}
Err(error) => return Err(BoosApiError::Io(error)),
};
let doc = content.parse::<toml_edit::DocumentMut>().map_err(|error| {
BoosApiError::Gateway(format!("Failed to parse Codex config.toml: {}", error))
})?;
Ok(doc
.get("model_provider")
.and_then(|value| value.as_str())
.filter(|value| !value.trim().is_empty())
.unwrap_or(DEFAULT_PROVIDER)
.to_string())
}
fn encrypted_content_warning(
counts: &BTreeMap<String, usize>,
target_provider: &str,
) -> Option<String> {
let risky: Vec<String> = counts
.iter()
.filter(|(provider, count)| provider.as_str() != target_provider && **count > 0)
.map(|(provider, _)| provider.clone())
.collect();
if risky.is_empty() {
None
} else {
Some(format!(
"{} rollout file(s) contain encrypted_content from provider(s) {}. Visibility metadata can be synchronized to {}, but continuing or compacting those histories may fail with invalid_encrypted_content.",
counts.values().sum::<usize>(),
risky.join(", "),
target_provider
))
}
}
pub(super) fn state_db_path(codex_home: &Path) -> PathBuf {
codex_home.join(STATE_DB_FILE)
}
pub(super) fn json_error(error: serde_json::Error) -> BoosApiError {
BoosApiError::Gateway(format!("Failed to process Codex JSON metadata: {}", error))
}
@@ -0,0 +1,192 @@
use std::path::{Path, PathBuf};
use crate::error::Result;
use crate::types::CodexSessionVisibilityRepairResult;
use super::{
backup, default_codex_home, encrypted_content_warning, ensure_codex_home, read_target_provider,
session_files, sqlite_state, state_db_path, workspace_roots, SessionChange, SessionScan,
SqliteUpdateResult, GLOBAL_STATE_FILE,
};
pub async fn repair_codex_session_visibility() -> Result<CodexSessionVisibilityRepairResult> {
let home = default_codex_home()?;
repair_codex_session_visibility_with_backup_at(&home, true).await
}
pub async fn repair_codex_session_visibility_at(
codex_home: &Path,
) -> Result<CodexSessionVisibilityRepairResult> {
repair_codex_session_visibility_with_backup_at(codex_home, true).await
}
pub async fn repair_codex_session_visibility_with_backup(
create_backup: bool,
) -> Result<CodexSessionVisibilityRepairResult> {
let home = default_codex_home()?;
repair_codex_session_visibility_with_backup_at(&home, create_backup).await
}
pub async fn repair_codex_session_visibility_with_backup_at(
codex_home: &Path,
create_backup: bool,
) -> Result<CodexSessionVisibilityRepairResult> {
ensure_codex_home(codex_home)?;
let target_provider = read_target_provider(codex_home)?;
let scan = session_files::collect_session_changes(codex_home, &target_provider)?;
sqlite_state::assert_sqlite_writable(codex_home)?;
let backup_dir = if create_backup {
Some(backup::create_backup(codex_home, &scan.changes)?)
} else {
None
};
let mut sqlite_result = SqliteUpdateResult::default();
let mut applied_changes: Vec<SessionChange> = Vec::new();
let global_state_backup = backup_dir.as_ref().map(|path| path.join(GLOBAL_STATE_FILE));
let global_state_original = if create_backup {
None
} else {
workspace_roots::read_global_state_text(codex_home)?
};
let mut global_state_updated = false;
match apply_repair(
codex_home,
&target_provider,
&scan,
&mut sqlite_result,
&mut applied_changes,
&mut global_state_updated,
) {
Ok(updated_workspace_roots) => Ok(build_result(
target_provider,
scan,
sqlite_result,
updated_workspace_roots,
backup_dir,
)),
Err(error) => {
session_files::restore_applied_changes(&applied_changes);
restore_global_state(
codex_home,
global_state_updated,
global_state_backup.as_ref(),
global_state_original.as_deref(),
);
Err(error)
}
}
}
fn restore_global_state(
codex_home: &Path,
updated: bool,
backup_path: Option<&PathBuf>,
original: Option<&str>,
) {
if !updated {
return;
}
if let Some(backup_path) = backup_path {
workspace_roots::restore_global_state(codex_home, backup_path);
} else {
workspace_roots::restore_global_state_text(codex_home, original);
}
}
fn apply_repair(
codex_home: &Path,
target_provider: &str,
scan: &SessionScan,
sqlite_result: &mut SqliteUpdateResult,
applied_changes: &mut Vec<SessionChange>,
global_state_updated: &mut bool,
) -> Result<usize> {
if state_db_path(codex_home).exists() {
return apply_repair_with_sqlite(
codex_home,
target_provider,
scan,
sqlite_result,
applied_changes,
global_state_updated,
);
}
session_files::apply_session_changes(&scan.changes, applied_changes)?;
let updated_roots = workspace_roots::sync_workspace_roots(codex_home, &scan.thread_cwd_by_id)?;
*global_state_updated = updated_roots > 0;
*sqlite_result = SqliteUpdateResult::default();
Ok(updated_roots)
}
fn apply_repair_with_sqlite(
codex_home: &Path,
target_provider: &str,
scan: &SessionScan,
sqlite_result: &mut SqliteUpdateResult,
applied_changes: &mut Vec<SessionChange>,
global_state_updated: &mut bool,
) -> Result<usize> {
let db = sqlite_state::SqliteDb::open(&state_db_path(codex_home))?;
db.exec("BEGIN IMMEDIATE")?;
let result = apply_repair_with_open_tx(
codex_home,
target_provider,
scan,
&db,
applied_changes,
global_state_updated,
);
match result {
Ok((sqlite_stats, updated_roots)) => {
db.exec("COMMIT")?;
*sqlite_result = sqlite_stats;
Ok(updated_roots)
}
Err(error) => {
let _ = db.exec("ROLLBACK");
Err(error)
}
}
}
fn apply_repair_with_open_tx(
codex_home: &Path,
target_provider: &str,
scan: &SessionScan,
db: &sqlite_state::SqliteDb,
applied_changes: &mut Vec<SessionChange>,
global_state_updated: &mut bool,
) -> Result<(SqliteUpdateResult, usize)> {
let sqlite_stats = sqlite_state::update_sqlite_provider(db, target_provider, scan)?;
session_files::apply_session_changes(&scan.changes, applied_changes)?;
let updated_roots = workspace_roots::sync_workspace_roots(codex_home, &scan.thread_cwd_by_id)?;
*global_state_updated = updated_roots > 0;
Ok((sqlite_stats, updated_roots))
}
fn build_result(
target_provider: String,
scan: SessionScan,
sqlite: SqliteUpdateResult,
updated_workspace_roots: usize,
backup_dir: Option<PathBuf>,
) -> CodexSessionVisibilityRepairResult {
CodexSessionVisibilityRepairResult {
target_provider: target_provider.clone(),
changed_session_files: scan.changes.len(),
skipped_locked_session_files: 0,
sqlite_rows_updated: sqlite.total_rows(),
sqlite_provider_rows_updated: sqlite.provider_rows,
sqlite_user_event_rows_updated: sqlite.user_event_rows,
sqlite_cwd_rows_updated: sqlite.cwd_rows,
sqlite_present: sqlite.present,
updated_workspace_roots,
backup_dir: backup_dir.map(|path| path.display().to_string()),
encrypted_content_warning: encrypted_content_warning(
&scan.encrypted_content_by_provider,
&target_provider,
),
}
}
@@ -0,0 +1,300 @@
use std::fs::{self, File, FileTimes};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use super::{json_error, workspace_roots, Result, SessionChange, SessionScan, SESSION_DIRS};
use serde_json::Value;
#[derive(Debug, Clone)]
struct FirstLineRecord {
first_line: String,
separator: String,
offset: u64,
}
pub(super) fn collect_session_changes(
codex_home: &Path,
target_provider: &str,
) -> Result<SessionScan> {
let mut scan = SessionScan::default();
for dir_name in SESSION_DIRS {
let root = codex_home.join(dir_name);
if root.exists() {
for path in list_rollout_files(&root)? {
collect_rollout_change(codex_home, &path, target_provider, &mut scan)?;
}
}
}
Ok(scan)
}
fn collect_rollout_change(
codex_home: &Path,
path: &Path,
target_provider: &str,
scan: &mut SessionScan,
) -> Result<()> {
let record = read_first_line_record(path)?;
let mut meta = parse_session_meta(&record.first_line)?;
if meta.is_null() {
return Ok(());
}
let payload = meta.get_mut("payload").and_then(Value::as_object_mut);
let Some(payload) = payload else {
return Ok(());
};
let current_provider = payload
.get("model_provider")
.and_then(Value::as_str)
.unwrap_or("(missing)")
.to_string();
record_rollout_facts(path, &record, payload, &current_provider, scan)?;
if current_provider != target_provider {
payload.insert(
"model_provider".into(),
Value::String(target_provider.into()),
);
scan.changes
.push(build_change(codex_home, path, record, meta)?);
}
Ok(())
}
fn build_change(
codex_home: &Path,
path: &Path,
record: FirstLineRecord,
meta: Value,
) -> Result<SessionChange> {
Ok(SessionChange {
path: path.to_path_buf(),
relative_path: path.strip_prefix(codex_home).unwrap_or(path).to_path_buf(),
original_first_line: record.first_line,
updated_first_line: serde_json::to_string(&meta).map_err(json_error)?,
separator: record.separator,
offset: record.offset,
original_mtime: fs::metadata(path)
.ok()
.and_then(|meta| meta.modified().ok()),
})
}
fn record_rollout_facts(
path: &Path,
record: &FirstLineRecord,
payload: &serde_json::Map<String, Value>,
provider: &str,
scan: &mut SessionScan,
) -> Result<()> {
let thread_id = payload
.get("id")
.and_then(Value::as_str)
.unwrap_or_default();
if !thread_id.is_empty() {
record_thread_cwd(payload, thread_id, scan);
if file_has_user_event(path, &record.first_line, record.offset)? {
scan.user_event_thread_ids.insert(thread_id.into());
}
}
if file_has_encrypted_content(path, &record.first_line, record.offset)? {
*scan
.encrypted_content_by_provider
.entry(provider.to_string())
.or_default() += 1;
}
Ok(())
}
fn record_thread_cwd(
payload: &serde_json::Map<String, Value>,
thread_id: &str,
scan: &mut SessionScan,
) {
let Some(cwd) = payload.get("cwd").and_then(Value::as_str) else {
return;
};
if !cwd.trim().is_empty() {
scan.thread_cwd_by_id.insert(
thread_id.into(),
workspace_roots::to_desktop_workspace_path(cwd),
);
}
}
fn list_rollout_files(root: &Path) -> Result<Vec<PathBuf>> {
let mut files = Vec::new();
list_rollout_files_into(root, &mut files)?;
files.sort();
Ok(files)
}
fn list_rollout_files_into(root: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
for entry in fs::read_dir(root)? {
let entry = entry?;
let path = entry.path();
let file_type = entry.file_type()?;
if file_type.is_dir() {
list_rollout_files_into(&path, files)?;
} else if file_type.is_file() && is_rollout_file(&path) {
files.push(path);
}
}
Ok(())
}
fn is_rollout_file(path: &Path) -> bool {
path.file_name()
.and_then(|name| name.to_str())
.map(|name| name.starts_with("rollout-") && name.ends_with(".jsonl"))
.unwrap_or(false)
}
fn read_first_line_record(path: &Path) -> Result<FirstLineRecord> {
let mut file = File::open(path)?;
let mut bytes = Vec::new();
file.read_to_end(&mut bytes)?;
let newline = bytes.iter().position(|byte| *byte == b'\n');
let (line_bytes, separator, offset) = match newline {
Some(index) if index > 0 && bytes[index - 1] == b'\r' => {
(&bytes[..index - 1], "\r\n".to_string(), index as u64 + 1)
}
Some(index) => (&bytes[..index], "\n".to_string(), index as u64 + 1),
None => (&bytes[..], String::new(), bytes.len() as u64),
};
Ok(FirstLineRecord {
first_line: String::from_utf8_lossy(line_bytes).to_string(),
separator,
offset,
})
}
fn parse_session_meta(first_line: &str) -> Result<Value> {
if first_line.trim().is_empty() {
return Ok(Value::Null);
}
let parsed: Value = serde_json::from_str(first_line).map_err(json_error)?;
if parsed.get("type").and_then(Value::as_str) == Some("session_meta")
&& parsed.get("payload").is_some()
{
Ok(parsed)
} else {
Ok(Value::Null)
}
}
fn file_has_encrypted_content(path: &Path, first_line: &str, offset: u64) -> Result<bool> {
Ok(first_line.contains("encrypted_content")
|| file_tail_contains(path, offset, "encrypted_content")?)
}
fn file_tail_contains(path: &Path, offset: u64, needle: &str) -> Result<bool> {
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(offset))?;
let mut content = String::new();
file.read_to_string(&mut content)?;
Ok(content.contains(needle))
}
fn file_has_user_event(path: &Path, first_line: &str, offset: u64) -> Result<bool> {
if line_has_user_event(first_line) {
return Ok(true);
}
let mut file = File::open(path)?;
file.seek(SeekFrom::Start(offset))?;
let mut content = String::new();
file.read_to_string(&mut content)?;
Ok(content.lines().any(line_has_user_event))
}
fn line_has_user_event(line: &str) -> bool {
let Ok(value) = serde_json::from_str::<Value>(line) else {
return false;
};
let event_msg = value.get("type").and_then(Value::as_str) == Some("event_msg")
&& value
.get("payload")
.and_then(|payload| payload.get("type"))
.and_then(Value::as_str)
== Some("user_message");
event_msg
|| ["payload", "item", "msg"]
.iter()
.any(|key| user_message_at(&value, key))
}
fn user_message_at(value: &Value, key: &str) -> bool {
let Some(message) = value.get(key) else {
return false;
};
message.get("type").and_then(Value::as_str) == Some("message")
&& message.get("role").and_then(Value::as_str) == Some("user")
}
pub(super) fn apply_session_changes(
changes: &[SessionChange],
applied: &mut Vec<SessionChange>,
) -> Result<usize> {
let mut changed = 0;
for change in changes {
if rewrite_collected_first_line(change)? {
changed += 1;
applied.push(change.clone());
}
}
Ok(changed)
}
fn rewrite_collected_first_line(change: &SessionChange) -> Result<bool> {
let current = read_first_line_record(&change.path)?;
if current.first_line != change.original_first_line || current.offset != change.offset {
return Ok(false);
}
let tmp_path = change.path.with_extension(format!(
"boosapi-session-visibility-{}.tmp",
chrono::Utc::now().timestamp_millis()
));
rewrite_first_line_to_tmp(change, &tmp_path)?;
fs::rename(&tmp_path, &change.path)?;
restore_mtime(&change.path, change.original_mtime);
Ok(true)
}
fn rewrite_first_line_to_tmp(change: &SessionChange, tmp_path: &Path) -> Result<()> {
let mut source = File::open(&change.path)?;
source.seek(SeekFrom::Start(change.offset))?;
let mut tmp = File::create(tmp_path)?;
tmp.write_all(change.updated_first_line.as_bytes())?;
tmp.write_all(change.separator.as_bytes())?;
std::io::copy(&mut source, &mut tmp)?;
tmp.flush()?;
Ok(())
}
pub(super) fn restore_applied_changes(changes: &[SessionChange]) {
for change in changes {
let _ = rewrite_first_line(&change.path, &change.original_first_line, &change.separator);
restore_mtime(&change.path, change.original_mtime);
}
}
fn rewrite_first_line(path: &Path, first_line: &str, separator: &str) -> Result<()> {
let current = read_first_line_record(path)?;
let tmp_path = path.with_extension("boosapi-session-visibility-restore.tmp");
let mut source = File::open(path)?;
source.seek(SeekFrom::Start(current.offset))?;
let mut tmp = File::create(&tmp_path)?;
tmp.write_all(first_line.as_bytes())?;
tmp.write_all(separator.as_bytes())?;
std::io::copy(&mut source, &mut tmp)?;
tmp.flush()?;
fs::rename(tmp_path, path)?;
Ok(())
}
fn restore_mtime(path: &Path, modified: Option<std::time::SystemTime>) {
let Some(modified) = modified else { return };
let Ok(file) = File::options().write(true).open(path) else {
return;
};
let _ = file.set_times(FileTimes::new().set_modified(modified));
}
@@ -0,0 +1,241 @@
use std::collections::BTreeSet;
use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_int};
use std::path::Path;
use std::ptr;
use libsqlite3_sys as sqlite;
use super::{state_db_path, BoosApiError, Result, SessionScan, SqliteUpdateResult};
pub(super) fn assert_sqlite_writable(codex_home: &Path) -> Result<()> {
let db_path = state_db_path(codex_home);
if !db_path.exists() {
return Ok(());
}
let db = SqliteDb::open(&db_path)?;
db.exec("BEGIN IMMEDIATE")?;
db.exec("ROLLBACK")
}
pub(super) fn update_sqlite_provider(
db: &SqliteDb,
target_provider: &str,
scan: &SessionScan,
) -> Result<SqliteUpdateResult> {
let columns = db.table_columns("threads")?;
let provider_rows = db.execute_with_params(
"UPDATE threads SET model_provider = ? WHERE COALESCE(model_provider, '') <> ?",
&[target_provider, target_provider],
)?;
let user_event_rows = update_user_event_rows(db, &columns, scan)?;
let cwd_rows = update_cwd_rows(db, &columns, scan)?;
Ok(SqliteUpdateResult {
present: true,
provider_rows,
user_event_rows,
cwd_rows,
})
}
fn update_user_event_rows(
db: &SqliteDb,
columns: &BTreeSet<String>,
scan: &SessionScan,
) -> Result<usize> {
if !columns.contains("has_user_event") {
return Ok(0);
}
let mut rows = 0;
for thread_id in &scan.user_event_thread_ids {
rows += db.execute_with_params(
"UPDATE threads SET has_user_event = 1 WHERE id = ? AND COALESCE(has_user_event, 0) <> 1",
&[thread_id],
)?;
}
Ok(rows)
}
fn update_cwd_rows(db: &SqliteDb, columns: &BTreeSet<String>, scan: &SessionScan) -> Result<usize> {
if !columns.contains("cwd") {
return Ok(0);
}
let mut rows = 0;
for (thread_id, cwd) in &scan.thread_cwd_by_id {
rows += db.execute_with_params(
"UPDATE threads SET cwd = ? WHERE id = ? AND COALESCE(cwd, '') <> ?",
&[cwd, thread_id, cwd],
)?;
}
Ok(rows)
}
pub(super) struct SqliteDb {
conn: *mut sqlite::sqlite3,
}
impl SqliteDb {
pub(super) fn open(path: &Path) -> Result<Self> {
let path = CString::new(path.to_string_lossy().as_bytes())
.map_err(|_| BoosApiError::Gateway("Invalid state_5.sqlite path".into()))?;
let mut conn = ptr::null_mut();
let flags = sqlite::SQLITE_OPEN_READWRITE | sqlite::SQLITE_OPEN_URI;
let rc = unsafe { sqlite::sqlite3_open_v2(path.as_ptr(), &mut conn, flags, ptr::null()) };
if rc == sqlite::SQLITE_OK {
return Ok(Self { conn });
}
let error = sqlite_error(conn, "open state_5.sqlite");
if !conn.is_null() {
unsafe { sqlite::sqlite3_close(conn) };
}
Err(error)
}
pub(super) fn exec(&self, sql: &str) -> Result<()> {
let sql = CString::new(sql).map_err(|_| BoosApiError::Gateway("Invalid SQL".into()))?;
let mut error: *mut c_char = ptr::null_mut();
let rc = unsafe {
sqlite::sqlite3_exec(self.conn, sql.as_ptr(), None, ptr::null_mut(), &mut error)
};
if rc == sqlite::SQLITE_OK {
return Ok(());
}
Err(BoosApiError::Gateway(format!(
"state_5.sqlite: {}",
sqlite_exec_error(self.conn, error)
)))
}
pub(super) fn table_columns(&self, table: &str) -> Result<BTreeSet<String>> {
let sql = format!("PRAGMA table_info({})", quote_sql_identifier(table));
let mut stmt = self.prepare(&sql)?;
let mut columns = BTreeSet::new();
while stmt.step_row()? {
if let Some(name) = stmt.column_text(1) {
columns.insert(name);
}
}
Ok(columns)
}
fn execute_with_params(&self, sql: &str, params: &[&str]) -> Result<usize> {
let mut stmt = self.prepare(sql)?;
for (index, value) in params.iter().enumerate() {
stmt.bind_text(index as c_int + 1, value)?;
}
stmt.step_done()?;
Ok(unsafe { sqlite::sqlite3_changes(self.conn) as usize })
}
pub(super) fn prepare(&self, sql: &str) -> Result<SqliteStatement<'_>> {
let sql = CString::new(sql).map_err(|_| BoosApiError::Gateway("Invalid SQL".into()))?;
let mut stmt = ptr::null_mut();
let rc = unsafe {
sqlite::sqlite3_prepare_v2(self.conn, sql.as_ptr(), -1, &mut stmt, ptr::null_mut())
};
if rc == sqlite::SQLITE_OK {
Ok(SqliteStatement { db: self, stmt })
} else {
Err(sqlite_error(self.conn, "prepare state_5.sqlite statement"))
}
}
}
impl Drop for SqliteDb {
fn drop(&mut self) {
if !self.conn.is_null() {
unsafe { sqlite::sqlite3_close(self.conn) };
}
}
}
pub(super) struct SqliteStatement<'a> {
db: &'a SqliteDb,
stmt: *mut sqlite::sqlite3_stmt,
}
impl SqliteStatement<'_> {
pub(super) fn bind_text(&mut self, index: c_int, value: &str) -> Result<()> {
let value = CString::new(value)
.map_err(|_| BoosApiError::Gateway("Invalid SQLite bind value".into()))?;
let rc = unsafe {
sqlite::sqlite3_bind_text(
self.stmt,
index,
value.as_ptr(),
-1,
sqlite::SQLITE_TRANSIENT(),
)
};
if rc == sqlite::SQLITE_OK {
Ok(())
} else {
Err(sqlite_error(self.db.conn, "bind state_5.sqlite parameter"))
}
}
fn step_done(&mut self) -> Result<()> {
match unsafe { sqlite::sqlite3_step(self.stmt) } {
sqlite::SQLITE_DONE => Ok(()),
_ => Err(sqlite_error(
self.db.conn,
"execute state_5.sqlite statement",
)),
}
}
pub(super) fn step_row(&mut self) -> Result<bool> {
match unsafe { sqlite::sqlite3_step(self.stmt) } {
sqlite::SQLITE_ROW => Ok(true),
sqlite::SQLITE_DONE => Ok(false),
_ => Err(sqlite_error(self.db.conn, "query state_5.sqlite")),
}
}
pub(super) fn column_text(&self, index: c_int) -> Option<String> {
let text = unsafe { sqlite::sqlite3_column_text(self.stmt, index) };
if text.is_null() {
return None;
}
let cstr = unsafe { CStr::from_ptr(text as *const c_char) };
Some(cstr.to_string_lossy().to_string())
}
pub(super) fn column_i64(&self, index: c_int) -> i64 {
unsafe { sqlite::sqlite3_column_int64(self.stmt, index) }
}
}
impl Drop for SqliteStatement<'_> {
fn drop(&mut self) {
if !self.stmt.is_null() {
unsafe { sqlite::sqlite3_finalize(self.stmt) };
}
}
}
fn quote_sql_identifier(value: &str) -> String {
format!("\"{}\"", value.replace('"', "\"\""))
}
fn sqlite_error(conn: *mut sqlite::sqlite3, action: &str) -> BoosApiError {
let message = if conn.is_null() {
"unknown SQLite error".into()
} else {
let ptr = unsafe { sqlite::sqlite3_errmsg(conn) };
unsafe { CStr::from_ptr(ptr).to_string_lossy().to_string() }
};
BoosApiError::Gateway(format!("Failed to {}: {}", action, message))
}
fn sqlite_exec_error(conn: *mut sqlite::sqlite3, error: *mut c_char) -> String {
if error.is_null() {
return match sqlite_error(conn, "execute") {
BoosApiError::Gateway(message) => message,
other => other.to_string(),
};
}
let message = unsafe { CStr::from_ptr(error).to_string_lossy().to_string() };
unsafe { sqlite::sqlite3_free(error as *mut _) };
message
}
@@ -0,0 +1,98 @@
use std::collections::BTreeMap;
use std::path::Path;
use super::{sqlite_state::SqliteDb, state_db_path, Result, SessionScan, SqliteStatus};
pub(super) fn read_sqlite_status(
codex_home: &Path,
target_provider: &str,
scan: &SessionScan,
) -> Result<SqliteStatus> {
let db_path = state_db_path(codex_home);
if !db_path.exists() {
return Ok(SqliteStatus::default());
}
let db = SqliteDb::open(&db_path)?;
let provider_counts = read_provider_counts(&db)?;
let rows = provider_counts.values().sum();
let mismatched_rows = provider_counts
.iter()
.filter(|(provider, _)| provider.as_str() != target_provider)
.map(|(_, count)| *count)
.sum();
let columns = db.table_columns("threads")?;
Ok(SqliteStatus {
present: true,
rows,
mismatched_rows,
user_event_rows: count_user_event_rows(&db, &columns, scan)?,
cwd_rows: count_cwd_rows(&db, &columns, scan)?,
provider_counts,
})
}
fn read_provider_counts(db: &SqliteDb) -> Result<BTreeMap<String, usize>> {
let mut stmt = db.prepare(
"SELECT COALESCE(model_provider, '') AS provider, COUNT(*) \
FROM threads GROUP BY COALESCE(model_provider, '')",
)?;
let mut counts = BTreeMap::new();
while stmt.step_row()? {
let provider = stmt
.column_text(0)
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "(missing)".into());
counts.insert(provider, stmt.column_i64(1).max(0) as usize);
}
Ok(counts)
}
fn count_user_event_rows(
db: &SqliteDb,
columns: &std::collections::BTreeSet<String>,
scan: &SessionScan,
) -> Result<usize> {
if !columns.contains("has_user_event") {
return Ok(0);
}
let mut rows = 0;
for thread_id in &scan.user_event_thread_ids {
rows += count_with_params(
db,
"SELECT COUNT(*) FROM threads WHERE id = ? AND COALESCE(has_user_event, 0) <> 1",
&[thread_id],
)?;
}
Ok(rows)
}
fn count_cwd_rows(
db: &SqliteDb,
columns: &std::collections::BTreeSet<String>,
scan: &SessionScan,
) -> Result<usize> {
if !columns.contains("cwd") {
return Ok(0);
}
let mut rows = 0;
for (thread_id, cwd) in &scan.thread_cwd_by_id {
rows += count_with_params(
db,
"SELECT COUNT(*) FROM threads WHERE id = ? AND COALESCE(cwd, '') <> ?",
&[thread_id, cwd],
)?;
}
Ok(rows)
}
fn count_with_params(db: &SqliteDb, sql: &str, params: &[&str]) -> Result<usize> {
let mut stmt = db.prepare(sql)?;
for (index, value) in params.iter().enumerate() {
stmt.bind_text(index as i32 + 1, value)?;
}
if stmt.step_row()? {
Ok(stmt.column_i64(0).max(0) as usize)
} else {
Ok(0)
}
}
@@ -0,0 +1,181 @@
use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use serde_json::Value;
use super::{
encrypted_content_warning, ensure_codex_home, json_error, read_target_provider, session_files,
sqlite_status, workspace_roots, Result, SESSION_DIRS,
};
use crate::types::{CodexSessionVisibilityStatus, CodexSessionVisibilityStatusRow};
#[derive(Debug, Default)]
struct SessionProviderCounts {
counts: BTreeMap<(String, String), usize>,
total: usize,
}
pub(super) fn get_status(codex_home: &Path) -> Result<CodexSessionVisibilityStatus> {
ensure_codex_home(codex_home)?;
let target_provider = read_target_provider(codex_home)?;
let scan = session_files::collect_session_changes(codex_home, &target_provider)?;
let session_counts = collect_session_provider_counts(codex_home)?;
let sqlite = sqlite_status::read_sqlite_status(codex_home, &target_provider, &scan)?;
let workspace_roots =
workspace_roots::preview_workspace_roots(codex_home, &scan.thread_cwd_by_id)?;
let mut rows = build_session_rows(&session_counts.counts, &target_provider);
rows.extend(build_provider_rows(
"state_5.sqlite",
&sqlite.provider_counts,
&target_provider,
));
if workspace_roots > 0 {
rows.push(CodexSessionVisibilityStatusRow {
scope: ".codex-global-state.json".into(),
provider: None,
count: workspace_roots,
mismatched_count: workspace_roots,
status: "needs_repair".into(),
});
}
Ok(CodexSessionVisibilityStatus {
target_provider: target_provider.clone(),
codex_home: codex_home.display().to_string(),
total_session_files: session_counts.total,
mismatched_session_files: scan.changes.len(),
sqlite_present: sqlite.present,
sqlite_rows: sqlite.rows,
sqlite_mismatched_rows: sqlite.mismatched_rows,
sqlite_user_event_rows_needing_repair: sqlite.user_event_rows,
sqlite_cwd_rows_needing_repair: sqlite.cwd_rows,
workspace_roots_needing_update: workspace_roots,
status_rows: rows,
encrypted_content_warning: encrypted_content_warning(
&scan.encrypted_content_by_provider,
&target_provider,
),
})
}
fn build_session_rows(
counts: &BTreeMap<(String, String), usize>,
target_provider: &str,
) -> Vec<CodexSessionVisibilityStatusRow> {
counts
.iter()
.map(|((scope, provider), count)| provider_row(scope, provider, *count, target_provider))
.collect()
}
fn build_provider_rows(
scope: &str,
counts: &BTreeMap<String, usize>,
target_provider: &str,
) -> Vec<CodexSessionVisibilityStatusRow> {
counts
.iter()
.map(|(provider, count)| provider_row(scope, provider, *count, target_provider))
.collect()
}
fn provider_row(
scope: &str,
provider: &str,
count: usize,
target_provider: &str,
) -> CodexSessionVisibilityStatusRow {
let mismatched_count = if provider == target_provider {
0
} else {
count
};
CodexSessionVisibilityStatusRow {
scope: scope.into(),
provider: Some(provider.into()),
count,
mismatched_count,
status: if mismatched_count == 0 {
"ok"
} else {
"needs_repair"
}
.into(),
}
}
fn collect_session_provider_counts(codex_home: &Path) -> Result<SessionProviderCounts> {
let mut result = SessionProviderCounts::default();
for dir_name in SESSION_DIRS {
let root = codex_home.join(dir_name);
if root.exists() {
for path in list_rollout_files(&root)? {
collect_file_provider(dir_name, &path, &mut result)?;
}
}
}
Ok(result)
}
fn collect_file_provider(
scope: &str,
path: &Path,
counts: &mut SessionProviderCounts,
) -> Result<()> {
let Some(first_line) = read_first_line(path)? else {
return Ok(());
};
let parsed: Value = serde_json::from_str(&first_line).map_err(json_error)?;
if parsed.get("type").and_then(Value::as_str) != Some("session_meta") {
return Ok(());
}
let provider = parsed
.get("payload")
.and_then(|payload| payload.get("model_provider"))
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.unwrap_or("(missing)")
.to_string();
*counts.counts.entry((scope.into(), provider)).or_default() += 1;
counts.total += 1;
Ok(())
}
fn list_rollout_files(root: &Path) -> Result<Vec<PathBuf>> {
let mut files = Vec::new();
list_rollout_files_into(root, &mut files)?;
files.sort();
Ok(files)
}
fn list_rollout_files_into(root: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
for entry in fs::read_dir(root)? {
let entry = entry?;
let path = entry.path();
let file_type = entry.file_type()?;
if file_type.is_dir() {
list_rollout_files_into(&path, files)?;
} else if file_type.is_file() && is_rollout_file(&path) {
files.push(path);
}
}
Ok(())
}
fn is_rollout_file(path: &Path) -> bool {
path.file_name()
.and_then(|name| name.to_str())
.map(|name| name.starts_with("rollout-") && name.ends_with(".jsonl"))
.unwrap_or(false)
}
fn read_first_line(path: &Path) -> Result<Option<String>> {
let mut reader = BufReader::new(File::open(path)?);
let mut first_line = String::new();
if reader.read_line(&mut first_line)? == 0 {
return Ok(None);
}
Ok(Some(first_line.trim_end_matches(['\r', '\n']).to_string()))
}
@@ -0,0 +1,231 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::Path;
use serde_json::{json, Value};
use super::{json_error, BoosApiError, Result, GLOBAL_STATE_BACKUP_FILE, GLOBAL_STATE_FILE};
pub(super) fn sync_workspace_roots(
codex_home: &Path,
cwd_by_id: &BTreeMap<String, String>,
) -> Result<usize> {
let path = codex_home.join(GLOBAL_STATE_FILE);
let content = match fs::read_to_string(&path) {
Ok(content) => content,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(0),
Err(error) => return Err(BoosApiError::Io(error)),
};
let mut state: Value = serde_json::from_str(&content).map_err(json_error)?;
let original_state = state.clone();
let known_cwds = dedupe_paths(cwd_by_id.values().cloned().collect());
let changed = rewrite_workspace_root_state(&mut state, &known_cwds);
let backup_missing = !codex_home.join(GLOBAL_STATE_BACKUP_FILE).exists();
if changed == 0 && state == original_state && !backup_missing {
return Ok(0);
}
let next = format!(
"{}\n",
serde_json::to_string_pretty(&state).map_err(json_error)?
);
fs::write(&path, &next)?;
fs::write(codex_home.join(GLOBAL_STATE_BACKUP_FILE), next)?;
Ok(changed)
}
pub(super) fn preview_workspace_roots(
codex_home: &Path,
cwd_by_id: &BTreeMap<String, String>,
) -> Result<usize> {
let Some(content) = read_global_state_text(codex_home)? else {
return Ok(0);
};
let mut state: Value = serde_json::from_str(&content).map_err(json_error)?;
let original_state = state.clone();
let known_cwds = dedupe_paths(cwd_by_id.values().cloned().collect());
let changed = rewrite_workspace_root_state(&mut state, &known_cwds);
if state == original_state {
Ok(0)
} else {
Ok(changed.max(1))
}
}
pub(super) fn read_global_state_text(codex_home: &Path) -> Result<Option<String>> {
match fs::read_to_string(codex_home.join(GLOBAL_STATE_FILE)) {
Ok(content) => Ok(Some(content)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(error) => Err(BoosApiError::Io(error)),
}
}
pub(super) fn restore_global_state_text(codex_home: &Path, content: Option<&str>) {
let path = codex_home.join(GLOBAL_STATE_FILE);
match content {
Some(content) => {
let _ = fs::write(path, content);
}
None => {
let _ = fs::remove_file(path);
}
}
}
fn rewrite_workspace_root_state(state: &mut Value, known_cwds: &[String]) -> usize {
let saved = to_path_array(state.get("electron-saved-workspace-roots"));
let project_order = to_path_array(state.get("project-order"));
let active = to_path_array(state.get("active-workspace-roots"));
let next_saved = build_next_saved_roots(&project_order, &saved, &active, known_cwds);
let next_order = build_next_project_order(&project_order, &saved, &next_saved, known_cwds);
let next_active = dedupe_paths(resolve_paths(active, known_cwds));
let changed = count_array_changes(&saved, &next_saved);
state["electron-saved-workspace-roots"] = json!(next_saved);
state["project-order"] = json!(next_order);
rewrite_active_roots(state, next_active);
rewrite_object_keys(state, "electron-workspace-root-labels", known_cwds);
rewrite_open_target_preferences(state, known_cwds);
changed
}
fn build_next_saved_roots(
project_order: &[String],
saved: &[String],
active: &[String],
known_cwds: &[String],
) -> Vec<String> {
let mut combined = Vec::new();
if !project_order.is_empty() {
combined.extend_from_slice(project_order);
}
combined.extend_from_slice(saved);
combined.extend_from_slice(active);
dedupe_paths(resolve_paths(combined, known_cwds))
}
fn build_next_project_order(
project_order: &[String],
saved: &[String],
next_saved: &[String],
known_cwds: &[String],
) -> Vec<String> {
if project_order.is_empty() {
return next_saved.to_vec();
}
let mut combined = Vec::new();
combined.extend_from_slice(project_order);
combined.extend_from_slice(saved);
dedupe_paths(resolve_paths(combined, known_cwds))
}
fn rewrite_active_roots(state: &mut Value, next_active: Vec<String>) {
let original_is_array = state
.get("active-workspace-roots")
.map(Value::is_array)
.unwrap_or(false);
if original_is_array {
state["active-workspace-roots"] = json!(next_active);
} else if let Some(first) = next_active.first() {
state["active-workspace-roots"] = json!(first);
}
}
fn rewrite_object_keys(state: &mut Value, key: &str, known_cwds: &[String]) {
let Some(object) = state.get(key).and_then(Value::as_object) else {
return;
};
let mut next = serde_json::Map::new();
for (path, value) in object {
next.insert(resolve_stored_path(path, known_cwds), value.clone());
}
state[key] = Value::Object(next);
}
fn rewrite_open_target_preferences(state: &mut Value, known_cwds: &[String]) {
let Some(preferences) = state.get_mut("open-in-target-preferences") else {
return;
};
let Some(per_path) = preferences.get_mut("perPath") else {
return;
};
let Some(object) = per_path.as_object() else {
return;
};
let mut next = serde_json::Map::new();
for (path, value) in object {
next.insert(resolve_stored_path(path, known_cwds), value.clone());
}
*per_path = Value::Object(next);
}
fn to_path_array(value: Option<&Value>) -> Vec<String> {
match value {
Some(Value::Array(values)) => values
.iter()
.filter_map(Value::as_str)
.filter(|value| !value.trim().is_empty())
.map(str::to_string)
.collect(),
Some(Value::String(value)) if !value.trim().is_empty() => vec![value.to_string()],
_ => Vec::new(),
}
}
fn resolve_paths(paths: Vec<String>, known_cwds: &[String]) -> Vec<String> {
paths
.into_iter()
.map(|path| resolve_stored_path(&path, known_cwds))
.collect()
}
fn resolve_stored_path(value: &str, known_cwds: &[String]) -> String {
let comparable = normalize_comparable_path(value);
known_cwds
.iter()
.find(|cwd| normalize_comparable_path(cwd) == comparable)
.cloned()
.unwrap_or_else(|| to_desktop_workspace_path(value))
}
fn dedupe_paths(paths: Vec<String>) -> Vec<String> {
let mut seen = BTreeSet::new();
let mut result = Vec::new();
for path in paths {
let comparable = normalize_comparable_path(&path);
if !comparable.is_empty() && seen.insert(comparable) {
result.push(path);
}
}
result
}
fn count_array_changes(previous: &[String], next: &[String]) -> usize {
let count = previous.len().max(next.len());
(0..count)
.filter(|index| previous.get(*index) != next.get(*index))
.count()
}
fn normalize_comparable_path(value: &str) -> String {
let mut normalized = to_desktop_workspace_path(value).trim().replace('/', "\\");
while normalized.ends_with('\\') && normalized.len() > 3 {
normalized.pop();
}
normalized.to_lowercase()
}
pub(super) fn to_desktop_workspace_path(value: &str) -> String {
let trimmed = value.trim();
if let Some(rest) = trimmed.strip_prefix(r"\\?\UNC\") {
return format!(r"\\{}", rest).replace('/', "\\");
}
if let Some(rest) = trimmed.strip_prefix(r"\\?\") {
return rest.replace('/', "\\");
}
trimmed.trim_end_matches(['/', '\\']).to_string()
}
pub(super) fn restore_global_state(codex_home: &Path, backup_path: &Path) {
if backup_path.exists() {
let _ = fs::copy(backup_path, codex_home.join(GLOBAL_STATE_FILE));
}
}
@@ -0,0 +1,85 @@
use sea_orm::*;
use crate::entity::context_sources;
use crate::error::{BoosApiError, Result};
use crate::types::{ContextSource, CreateContextSourceInput};
use crate::utils::gen_id;
fn model_to_context_source(m: context_sources::Model) -> ContextSource {
ContextSource {
id: m.id,
conversation_id: m.conversation_id,
message_id: m.message_id,
source_type: m.source_type,
ref_id: m.ref_id,
title: m.title,
enabled: m.enabled != 0,
summary: m.summary,
}
}
pub async fn list_context_sources(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<ContextSource>> {
let models = context_sources::Entity::find()
.filter(context_sources::Column::ConversationId.eq(conversation_id))
.all(db)
.await?;
Ok(models.into_iter().map(model_to_context_source).collect())
}
pub async fn get_context_source(db: &DatabaseConnection, id: &str) -> Result<ContextSource> {
let model = context_sources::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ContextSource {}", id)))?;
Ok(model_to_context_source(model))
}
pub async fn add_context_source(
db: &DatabaseConnection,
input: &CreateContextSourceInput,
) -> Result<ContextSource> {
let id = gen_id();
let am = context_sources::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(input.conversation_id.clone()),
message_id: Set(input.message_id.clone()),
source_type: Set(input.source_type.clone()),
ref_id: Set(input.ref_id.clone()),
title: Set(input.title.clone()),
enabled: Set(1),
summary: Set(input.summary.clone()),
};
am.insert(db).await?;
get_context_source(db, &id).await
}
pub async fn remove_context_source(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = context_sources::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("ContextSource {}", id)));
}
Ok(())
}
pub async fn toggle_context_source(db: &DatabaseConnection, id: &str) -> Result<ContextSource> {
let model = context_sources::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ContextSource {}", id)))?;
let new_enabled = if model.enabled != 0 { 0 } else { 1 };
let mut am: context_sources::ActiveModel = model.into();
am.enabled = Set(new_enabled);
am.update(db).await?;
get_context_source(db, id).await
}
@@ -0,0 +1,635 @@
use sea_orm::*;
use serde_json;
use crate::entity::{conversation_summaries, conversations, messages};
use crate::error::{BoosApiError, Result};
use crate::types::{
Conversation, ConversationSearchResult, ConversationSummary, UpdateConversationInput,
};
use crate::utils::{gen_id, now_ts};
fn conversation_from_entity(m: conversations::Model) -> Conversation {
Conversation {
id: m.id,
title: m.title,
model_id: m.model_id,
provider_id: m.provider_id,
system_prompt: m.system_prompt,
temperature: m.temperature.map(|v| v as f32),
max_tokens: m.max_tokens.map(|v| v as u32),
top_p: m.top_p.map(|v| v as f32),
frequency_penalty: m.frequency_penalty.map(|v| v as f32),
search_enabled: m.search_enabled != 0,
search_provider_id: m.search_provider_id,
thinking_budget: m.thinking_budget,
thinking_level: m.thinking_level,
enabled_mcp_server_ids: parse_string_list(&m.enabled_mcp_server_ids),
enabled_knowledge_base_ids: parse_string_list(&m.enabled_knowledge_base_ids),
enabled_memory_namespace_ids: parse_string_list(&m.enabled_memory_namespace_ids),
message_count: m.message_count as u32,
is_pinned: m.is_pinned != 0,
is_archived: m.is_archived != 0,
context_compression: m.context_compression != 0,
category_id: m.category_id,
parent_conversation_id: m.parent_conversation_id,
mode: m.mode,
created_at: m.created_at,
updated_at: m.updated_at,
}
}
fn parse_string_list(raw: &str) -> Vec<String> {
serde_json::from_str(raw)
.expect("conversation preference JSON is invalid; database contents are corrupted")
}
fn stringify_string_list(values: &[String]) -> String {
serde_json::to_string(values).expect("failed to serialize conversation preference JSON")
}
pub async fn list_conversations(db: &DatabaseConnection) -> Result<Vec<Conversation>> {
let rows = conversations::Entity::find()
.filter(conversations::Column::IsArchived.eq(0))
.order_by_desc(conversations::Column::IsPinned)
.order_by_desc(conversations::Column::UpdatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(conversation_from_entity).collect())
}
pub async fn list_archived_conversations(db: &DatabaseConnection) -> Result<Vec<Conversation>> {
let rows = conversations::Entity::find()
.filter(conversations::Column::IsArchived.ne(0))
.order_by_desc(conversations::Column::UpdatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(conversation_from_entity).collect())
}
pub async fn get_conversation(db: &DatabaseConnection, id: &str) -> Result<Conversation> {
let row = conversations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Conversation {}", id)))?;
Ok(conversation_from_entity(row))
}
pub async fn create_conversation(
db: &DatabaseConnection,
title: &str,
model_id: &str,
provider_id: &str,
system_prompt: Option<&str>,
) -> Result<Conversation> {
let id = gen_id();
let now = now_ts();
conversations::ActiveModel {
id: Set(id.clone()),
title: Set(title.to_string()),
model_id: Set(model_id.to_string()),
provider_id: Set(provider_id.to_string()),
system_prompt: Set(system_prompt.map(|s| s.to_string())),
message_count: Set(0),
is_pinned: Set(0),
created_at: Set(now),
updated_at: Set(now),
..Default::default()
}
.insert(db)
.await?;
get_conversation(db, &id).await
}
pub async fn update_conversation(
db: &DatabaseConnection,
id: &str,
input: UpdateConversationInput,
) -> Result<Conversation> {
let row = conversations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Conversation {}", id)))?;
let now = now_ts();
let existing = conversation_from_entity(row.clone());
let title = input.title.unwrap_or(existing.title);
let provider_id = input.provider_id.unwrap_or(existing.provider_id);
let model_id = input.model_id.unwrap_or(existing.model_id);
let is_pinned = input.is_pinned.unwrap_or(existing.is_pinned);
let is_archived = input.is_archived.unwrap_or(existing.is_archived);
let mut am: conversations::ActiveModel = row.into();
am.title = Set(title);
am.provider_id = Set(provider_id);
am.model_id = Set(model_id);
am.is_pinned = Set(if is_pinned { 1 } else { 0 });
am.is_archived = Set(if is_archived { 1 } else { 0 });
if let Some(ref sp) = input.system_prompt {
am.system_prompt = Set(if sp.is_empty() {
None
} else {
Some(sp.clone())
});
}
if let Some(temperature) = input.temperature {
am.temperature = Set(temperature);
}
if let Some(max_tokens) = input.max_tokens {
am.max_tokens = Set(max_tokens);
}
if let Some(top_p) = input.top_p {
am.top_p = Set(top_p);
}
if let Some(frequency_penalty) = input.frequency_penalty {
am.frequency_penalty = Set(frequency_penalty);
}
if let Some(search_enabled) = input.search_enabled {
am.search_enabled = Set(if search_enabled { 1 } else { 0 });
}
if let Some(search_provider_id) = input.search_provider_id {
am.search_provider_id = Set(search_provider_id);
}
if let Some(thinking_budget) = input.thinking_budget {
am.thinking_budget = Set(thinking_budget);
}
if let Some(thinking_level) = input.thinking_level {
am.thinking_level = Set(thinking_level);
}
if let Some(enabled_mcp_server_ids) = input.enabled_mcp_server_ids {
am.enabled_mcp_server_ids = Set(stringify_string_list(&enabled_mcp_server_ids));
}
if let Some(enabled_knowledge_base_ids) = input.enabled_knowledge_base_ids {
am.enabled_knowledge_base_ids = Set(stringify_string_list(&enabled_knowledge_base_ids));
}
if let Some(enabled_memory_namespace_ids) = input.enabled_memory_namespace_ids {
am.enabled_memory_namespace_ids = Set(stringify_string_list(&enabled_memory_namespace_ids));
}
if let Some(context_compression) = input.context_compression {
am.context_compression = Set(if context_compression { 1 } else { 0 });
}
if let Some(category_id) = input.category_id {
am.category_id = Set(category_id);
}
if let Some(parent_conversation_id) = input.parent_conversation_id {
am.parent_conversation_id = Set(parent_conversation_id);
}
if let Some(mode) = input.mode {
am.mode = Set(mode);
}
am.updated_at = Set(now);
am.update(db).await?;
get_conversation(db, id).await
}
pub async fn update_conversation_title(
db: &DatabaseConnection,
id: &str,
title: &str,
) -> Result<()> {
if let Some(row) = conversations::Entity::find_by_id(id).one(db).await? {
let mut am: conversations::ActiveModel = row.into();
am.title = Set(title.to_string());
am.updated_at = Set(now_ts());
am.update(db).await?;
}
Ok(())
}
pub async fn toggle_pin(db: &DatabaseConnection, id: &str) -> Result<Conversation> {
let row = conversations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Conversation {}", id)))?;
let new_pinned = if row.is_pinned != 0 { 0 } else { 1 };
let now = now_ts();
let mut am: conversations::ActiveModel = row.into();
am.is_pinned = Set(new_pinned);
am.updated_at = Set(now);
am.update(db).await?;
get_conversation(db, id).await
}
pub async fn toggle_archive(db: &DatabaseConnection, id: &str) -> Result<Conversation> {
let row = conversations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Conversation {}", id)))?;
let new_archived = if row.is_archived != 0 { 0 } else { 1 };
let now = now_ts();
let mut am: conversations::ActiveModel = row.into();
am.is_archived = Set(new_archived);
am.updated_at = Set(now);
am.update(db).await?;
get_conversation(db, id).await
}
pub async fn delete_conversation(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = conversations::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("Conversation {}", id)));
}
Ok(())
}
/// Branch a conversation: copy settings + messages up to `until_message_id`.
/// If `as_child` is true, the new conversation is nested under the source (or its parent).
pub async fn branch_conversation(
db: &DatabaseConnection,
conversation_id: &str,
until_message_id: &str,
as_child: bool,
custom_title: Option<&str>,
) -> Result<Conversation> {
// 1. Load source conversation
let source = conversations::Entity::find_by_id(conversation_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Conversation {}", conversation_id)))?;
// 2. Load all active messages ordered by created_at
let all_msgs = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::IsActive.eq(1))
.order_by_asc(messages::Column::CreatedAt)
.all(db)
.await?;
// 3. Build the branch candidate list. Normal branches target an active
// message. Multi-model cards may target an inactive assistant version, in
// which case the selected version replaces its active sibling for the
// branch snapshot.
let candidate_msgs: Vec<messages::Model> = if let Some(target_idx) = all_msgs
.iter()
.position(|m| m.id == until_message_id)
{
all_msgs[..=target_idx].to_vec()
} else {
let target = messages::Entity::find_by_id(until_message_id)
.one(db)
.await?
.ok_or_else(|| {
BoosApiError::NotFound(format!("Message {} in conversation", until_message_id))
})?;
if target.conversation_id != conversation_id {
return Err(BoosApiError::NotFound(format!(
"Message {} in conversation {}",
until_message_id, conversation_id
)));
}
let parent_message_id = target.parent_message_id.clone().ok_or_else(|| {
BoosApiError::NotFound(format!("Message {} in conversation", until_message_id))
})?;
if target.role != "assistant" {
return Err(BoosApiError::NotFound(format!(
"Message {} in conversation",
until_message_id
)));
}
let parent_idx = all_msgs
.iter()
.position(|m| m.id == parent_message_id)
.ok_or_else(|| {
BoosApiError::NotFound(format!("Message {} in conversation", until_message_id))
})?;
let mut selected_branch = all_msgs[..=parent_idx].to_vec();
selected_branch.push(target);
selected_branch
};
// 4. Find last context-clear marker to determine effective start
let start_idx = candidate_msgs
.iter()
.rposition(|m| {
m.role == "system"
&& (m.content == "<!-- context-clear -->"
|| m.content == "<!-- context-compressed -->")
})
.map(|idx| idx + 1) // skip the marker itself
.unwrap_or(0);
let effective_msgs = &candidate_msgs[start_idx..];
// 5. Create new conversation with copied settings
let new_id = gen_id();
let now = now_ts();
let branch_title = custom_title
.map(|t| t.to_string())
.unwrap_or_else(|| source.title.clone());
// Determine parent_conversation_id
let parent_id = if as_child {
// If source already has a parent, new branch is a sibling (same parent)
// Otherwise, source becomes the parent
Some(
source
.parent_conversation_id
.clone()
.unwrap_or_else(|| source.id.clone()),
)
} else {
None
};
conversations::ActiveModel {
id: Set(new_id.clone()),
title: Set(branch_title),
model_id: Set(source.model_id.clone()),
provider_id: Set(source.provider_id.clone()),
system_prompt: Set(source.system_prompt.clone()),
temperature: Set(source.temperature),
max_tokens: Set(source.max_tokens),
top_p: Set(source.top_p),
frequency_penalty: Set(source.frequency_penalty),
search_enabled: Set(source.search_enabled),
search_provider_id: Set(source.search_provider_id.clone()),
thinking_budget: Set(source.thinking_budget),
thinking_level: Set(source.thinking_level.clone()),
enabled_mcp_server_ids: Set(source.enabled_mcp_server_ids.clone()),
enabled_knowledge_base_ids: Set(source.enabled_knowledge_base_ids.clone()),
enabled_memory_namespace_ids: Set(source.enabled_memory_namespace_ids.clone()),
message_count: Set(effective_msgs.len() as i32),
is_pinned: Set(0),
is_archived: Set(0),
context_compression: Set(source.context_compression),
category_id: Set(source.category_id.clone()),
parent_conversation_id: Set(parent_id),
research_mode: Set(source.research_mode),
created_at: Set(now),
updated_at: Set(now),
..Default::default()
}
.insert(db)
.await?;
// 6. Copy messages — assign new IDs and remap parent_message_id references
let mut id_map = std::collections::HashMap::new();
for msg in effective_msgs {
let new_msg_id = gen_id();
id_map.insert(msg.id.clone(), new_msg_id.clone());
let new_parent = msg
.parent_message_id
.as_ref()
.and_then(|pid| id_map.get(pid))
.cloned();
messages::ActiveModel {
id: Set(new_msg_id),
conversation_id: Set(new_id.clone()),
role: Set(msg.role.clone()),
content: Set(msg.content.clone()),
provider_id: Set(msg.provider_id.clone()),
model_id: Set(msg.model_id.clone()),
token_count: Set(msg.token_count),
prompt_tokens: Set(msg.prompt_tokens),
completion_tokens: Set(msg.completion_tokens),
attachments: Set(msg.attachments.clone()),
thinking: Set(msg.thinking.clone()),
created_at: Set(msg.created_at),
parent_message_id: Set(new_parent),
version_index: Set(msg.version_index),
is_active: Set(1),
tool_calls_json: Set(msg.tool_calls_json.clone()),
tool_call_id: Set(msg.tool_call_id.clone()),
status: Set(msg.status.clone()),
tokens_per_second: Set(msg.tokens_per_second),
first_token_latency_ms: Set(msg.first_token_latency_ms),
..Default::default()
}
.insert(db)
.await?;
}
get_conversation(db, &new_id).await
}
pub async fn search_conversations(
db: &DatabaseConnection,
query: &str,
) -> Result<Vec<ConversationSearchResult>> {
#[derive(Debug, FromQueryResult)]
struct FtsRow {
conversation_id: String,
preview: String,
}
let fts_rows = FtsRow::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT m.conversation_id, snippet(messages_fts, 0, '', '', '...', 32) as preview \
FROM messages_fts \
JOIN messages m ON m.rowid = messages_fts.rowid \
WHERE messages_fts MATCH ? \
GROUP BY m.conversation_id \
ORDER BY rank",
[query.into()],
))
.all(db)
.await?;
let mut results = Vec::with_capacity(fts_rows.len());
for fts in fts_rows {
if let Ok(conv) = get_conversation(db, &fts.conversation_id).await {
results.push(ConversationSearchResult {
conversation: conv,
matched_message_preview: Some(fts.preview),
});
}
}
Ok(results)
}
pub async fn increment_message_count(db: &DatabaseConnection, conversation_id: &str) -> Result<()> {
db.execute(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"UPDATE conversations SET message_count = message_count + 1, updated_at = ? WHERE id = ?",
[now_ts().into(), conversation_id.into()],
))
.await?;
Ok(())
}
pub async fn decrement_message_count(db: &DatabaseConnection, conversation_id: &str) -> Result<()> {
db.execute(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"UPDATE conversations SET message_count = MAX(0, message_count - 1), updated_at = ? WHERE id = ?",
[now_ts().into(), conversation_id.into()],
))
.await?;
Ok(())
}
// ── Conversation summaries ──────────────────────────────────────────────
fn summary_from_entity(m: conversation_summaries::Model) -> ConversationSummary {
ConversationSummary {
id: m.id,
conversation_id: m.conversation_id,
summary_text: m.summary_text,
compressed_until_message_id: m.compressed_until_message_id,
token_count: m.token_count.map(|v| v as u32),
model_used: m.model_used,
created_at: m.created_at,
updated_at: m.updated_at,
}
}
pub async fn get_summary(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Option<ConversationSummary>> {
let row = conversation_summaries::Entity::find()
.filter(conversation_summaries::Column::ConversationId.eq(conversation_id))
.order_by_desc(conversation_summaries::Column::UpdatedAt)
.one(db)
.await?;
Ok(row.map(summary_from_entity))
}
pub async fn upsert_summary(
db: &DatabaseConnection,
conversation_id: &str,
summary_text: &str,
compressed_until_message_id: Option<&str>,
token_count: Option<u32>,
model_used: Option<&str>,
) -> Result<ConversationSummary> {
let now = now_ts();
let existing = conversation_summaries::Entity::find()
.filter(conversation_summaries::Column::ConversationId.eq(conversation_id))
.one(db)
.await?;
match existing {
Some(row) => {
let mut am: conversation_summaries::ActiveModel = row.into();
am.summary_text = Set(summary_text.to_string());
am.compressed_until_message_id =
Set(compressed_until_message_id.map(|s| s.to_string()));
am.token_count = Set(token_count.map(|v| v as i64));
am.model_used = Set(model_used.map(|s| s.to_string()));
am.updated_at = Set(now);
am.update(db).await?;
}
None => {
let id = gen_id();
conversation_summaries::ActiveModel {
id: Set(id),
conversation_id: Set(conversation_id.to_string()),
summary_text: Set(summary_text.to_string()),
compressed_until_message_id: Set(
compressed_until_message_id.map(|s| s.to_string()),
),
token_count: Set(token_count.map(|v| v as i64)),
model_used: Set(model_used.map(|s| s.to_string())),
created_at: Set(now),
updated_at: Set(now),
}
.insert(db)
.await?;
}
}
get_summary(db, conversation_id).await?.ok_or_else(|| {
BoosApiError::Database(sea_orm::DbErr::Custom(
"Failed to read back upserted summary".into(),
))
})
}
pub async fn delete_summary(db: &DatabaseConnection, conversation_id: &str) -> Result<()> {
conversation_summaries::Entity::delete_many()
.filter(conversation_summaries::Column::ConversationId.eq(conversation_id))
.exec(db)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::create_test_pool;
use crate::repo::message;
use crate::types::MessageRole;
#[tokio::test]
async fn branch_conversation_from_inactive_assistant_version_uses_selected_version() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let conv = create_conversation(db, "Branch Source", "model-a", "provider-a", None)
.await
.unwrap();
let user = message::create_message(
db,
&conv.id,
MessageRole::User,
"Compare answers",
&[],
None,
0,
)
.await
.unwrap();
let active = message::create_message(
db,
&conv.id,
MessageRole::Assistant,
"Active answer",
&[],
Some(&user.id),
0,
)
.await
.unwrap();
let inactive = message::create_message(
db,
&conv.id,
MessageRole::Assistant,
"Inactive answer",
&[],
Some(&user.id),
1,
)
.await
.unwrap();
message::set_active_version(db, &conv.id, &user.id, &active.id)
.await
.unwrap();
let branched = branch_conversation(
db,
&conv.id,
&inactive.id,
false,
Some("Branched from inactive"),
)
.await
.unwrap();
let branched_messages = message::list_messages(db, &branched.id).await.unwrap();
assert_eq!(branched_messages.len(), 2);
assert_eq!(branched_messages[0].content, "Compare answers");
assert_eq!(branched_messages[1].content, "Inactive answer");
assert!(branched_messages[1].is_active);
}
}
@@ -0,0 +1,70 @@
use sea_orm::*;
use crate::entity::conversation_branches;
use crate::error::{BoosApiError, Result};
use crate::types::ConversationBranch;
use crate::utils::gen_id;
fn model_to_branch(m: conversation_branches::Model) -> ConversationBranch {
ConversationBranch {
id: m.id,
conversation_id: m.conversation_id,
parent_message_id: m.parent_message_id,
branch_label: m.branch_label,
branch_index: m.branch_index,
compared_message_ids_json: m.compared_message_ids_json,
created_at: m.created_at,
}
}
pub async fn list_branches(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<ConversationBranch>> {
let models = conversation_branches::Entity::find()
.filter(conversation_branches::Column::ConversationId.eq(conversation_id))
.order_by_asc(conversation_branches::Column::BranchIndex)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_branch).collect())
}
pub async fn get_branch(db: &DatabaseConnection, id: &str) -> Result<ConversationBranch> {
let model = conversation_branches::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ConversationBranch {}", id)))?;
Ok(model_to_branch(model))
}
pub async fn create_branch(
db: &DatabaseConnection,
conversation_id: &str,
parent_message_id: &str,
label: &str,
) -> Result<ConversationBranch> {
let id = gen_id();
let count = conversation_branches::Entity::find()
.filter(conversation_branches::Column::ConversationId.eq(conversation_id))
.count(db)
.await? as i32;
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
let am = conversation_branches::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(conversation_id.to_string()),
parent_message_id: Set(parent_message_id.to_string()),
branch_label: Set(label.to_string()),
branch_index: Set(count),
compared_message_ids_json: Set(None),
created_at: Set(now),
};
am.insert(db).await?;
get_branch(db, &id).await
}
@@ -0,0 +1,184 @@
use sea_orm::*;
use sea_orm::prelude::Expr;
use crate::entity::conversation_categories;
use crate::error::{BoosApiError, Result};
use crate::types::{
ConversationCategory, CreateConversationCategoryInput, UpdateConversationCategoryInput,
};
use crate::utils::{gen_id, now_ts};
fn category_from_entity(m: conversation_categories::Model) -> ConversationCategory {
ConversationCategory {
id: m.id,
name: m.name,
icon_type: m.icon_type,
icon_value: m.icon_value,
system_prompt: m.system_prompt,
default_provider_id: m.default_provider_id,
default_model_id: m.default_model_id,
default_temperature: m.default_temperature,
default_max_tokens: m.default_max_tokens,
default_top_p: m.default_top_p,
default_frequency_penalty: m.default_frequency_penalty,
sort_order: m.sort_order,
is_collapsed: m.is_collapsed != 0,
created_at: m.created_at,
updated_at: m.updated_at,
}
}
pub async fn list_conversation_categories(
db: &DatabaseConnection,
) -> Result<Vec<ConversationCategory>> {
let rows = conversation_categories::Entity::find()
.order_by_asc(conversation_categories::Column::SortOrder)
.all(db)
.await?;
Ok(rows.into_iter().map(category_from_entity).collect())
}
pub async fn create_conversation_category(
db: &DatabaseConnection,
input: CreateConversationCategoryInput,
) -> Result<ConversationCategory> {
let id = gen_id();
let now = now_ts();
let max_order = conversation_categories::Entity::find()
.select_only()
.column_as(conversation_categories::Column::SortOrder.max(), "m")
.into_tuple::<Option<i32>>()
.one(db)
.await?
.flatten();
let sort_order = max_order.unwrap_or(-1) + 1;
let am = conversation_categories::ActiveModel {
id: Set(id),
name: Set(input.name),
icon_type: Set(input.icon_type),
icon_value: Set(input.icon_value),
system_prompt: Set(input.system_prompt),
default_provider_id: Set(input.default_provider_id),
default_model_id: Set(input.default_model_id),
default_temperature: Set(input.default_temperature),
default_max_tokens: Set(input.default_max_tokens),
default_top_p: Set(input.default_top_p),
default_frequency_penalty: Set(input.default_frequency_penalty),
sort_order: Set(sort_order),
is_collapsed: Set(1),
created_at: Set(now),
updated_at: Set(now),
};
let row = am.insert(db).await?;
Ok(category_from_entity(row))
}
pub async fn update_conversation_category(
db: &DatabaseConnection,
id: &str,
input: UpdateConversationCategoryInput,
) -> Result<ConversationCategory> {
let row = conversation_categories::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ConversationCategory {}", id)))?;
let mut am: conversation_categories::ActiveModel = row.into();
if let Some(name) = input.name {
am.name = Set(name);
}
if let Some(icon_type) = input.icon_type {
am.icon_type = Set(icon_type);
}
if let Some(icon_value) = input.icon_value {
am.icon_value = Set(icon_value);
}
if let Some(system_prompt) = input.system_prompt {
am.system_prompt = Set(system_prompt);
}
if let Some(default_provider_id) = input.default_provider_id {
am.default_provider_id = Set(default_provider_id);
}
if let Some(default_model_id) = input.default_model_id {
am.default_model_id = Set(default_model_id);
}
if let Some(default_temperature) = input.default_temperature {
am.default_temperature = Set(default_temperature);
}
if let Some(default_max_tokens) = input.default_max_tokens {
am.default_max_tokens = Set(default_max_tokens);
}
if let Some(default_top_p) = input.default_top_p {
am.default_top_p = Set(default_top_p);
}
if let Some(default_frequency_penalty) = input.default_frequency_penalty {
am.default_frequency_penalty = Set(default_frequency_penalty);
}
am.updated_at = Set(now_ts());
am.update(db).await?;
let updated = conversation_categories::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ConversationCategory {}", id)))?;
Ok(category_from_entity(updated))
}
pub async fn delete_conversation_category(db: &DatabaseConnection, id: &str) -> Result<()> {
// Unset category_id on conversations that belong to this category
use crate::entity::conversations;
conversations::Entity::update_many()
.col_expr(
conversations::Column::CategoryId,
Expr::value(Option::<String>::None),
)
.filter(conversations::Column::CategoryId.eq(id))
.exec(db)
.await?;
conversation_categories::Entity::delete_by_id(id)
.exec(db)
.await?;
Ok(())
}
pub async fn reorder_conversation_categories(
db: &DatabaseConnection,
category_ids: &[String],
) -> Result<()> {
let now = now_ts();
for (i, id) in category_ids.iter().enumerate() {
conversation_categories::Entity::update_many()
.col_expr(
conversation_categories::Column::SortOrder,
Expr::value(i as i32),
)
.col_expr(
conversation_categories::Column::UpdatedAt,
Expr::value(now),
)
.filter(conversation_categories::Column::Id.eq(id))
.exec(db)
.await?;
}
Ok(())
}
pub async fn set_conversation_category_collapsed(
db: &DatabaseConnection,
id: &str,
collapsed: bool,
) -> Result<()> {
let row = conversation_categories::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ConversationCategory {}", id)))?;
let mut am: conversation_categories::ActiveModel = row.into();
am.is_collapsed = Set(if collapsed { 1 } else { 0 });
am.updated_at = Set(now_ts());
am.update(db).await?;
Ok(())
}
@@ -0,0 +1,443 @@
use sea_orm::*;
use serde::{Deserialize, Serialize};
use crate::entity::{drawing_generations, drawing_images};
use crate::error::{BoosApiError, Result};
use crate::repo::stored_file::StoredFile;
use crate::utils::{gen_id, now_ts};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DrawingImage {
pub id: String,
pub generation_id: String,
pub stored_file_id: String,
pub storage_path: String,
pub mime_type: String,
pub width: Option<i32>,
pub height: Option<i32>,
pub revised_prompt: Option<String>,
pub created_at: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DrawingGeneration {
pub id: String,
pub parent_generation_id: Option<String>,
pub provider_id: String,
pub key_id: String,
pub model_id: String,
pub api_kind: String,
pub action: String,
pub prompt: String,
pub parameters_json: String,
pub reference_file_ids_json: String,
pub source_image_ids_json: String,
pub mask_file_id: Option<String>,
pub status: String,
pub error_message: Option<String>,
pub response_id: Option<String>,
pub usage_json: Option<String>,
pub created_at: i64,
pub completed_at: Option<i64>,
pub images: Vec<DrawingImage>,
#[serde(default)]
pub reference_files: Vec<StoredFile>,
#[serde(default)]
pub source_images: Vec<DrawingImage>,
pub mask_file: Option<StoredFile>,
}
#[derive(Debug, Clone)]
pub struct NewDrawingGeneration {
pub parent_generation_id: Option<String>,
pub provider_id: String,
pub key_id: String,
pub model_id: String,
pub action: String,
pub prompt: String,
pub parameters_json: String,
pub reference_file_ids_json: String,
pub source_image_ids_json: String,
pub mask_file_id: Option<String>,
}
#[derive(Debug, Clone)]
pub struct NewDrawingImage {
pub generation_id: String,
pub stored_file_id: String,
pub storage_path: String,
pub mime_type: String,
pub width: Option<i32>,
pub height: Option<i32>,
pub revised_prompt: Option<String>,
}
fn image_from_entity(model: drawing_images::Model) -> DrawingImage {
DrawingImage {
id: model.id,
generation_id: model.generation_id,
stored_file_id: model.stored_file_id,
storage_path: model.storage_path,
mime_type: model.mime_type,
width: model.width,
height: model.height,
revised_prompt: model.revised_prompt,
created_at: model.created_at,
}
}
fn generation_from_entity(
model: drawing_generations::Model,
images: Vec<DrawingImage>,
reference_files: Vec<StoredFile>,
source_images: Vec<DrawingImage>,
mask_file: Option<StoredFile>,
) -> DrawingGeneration {
DrawingGeneration {
id: model.id,
parent_generation_id: model.parent_generation_id,
provider_id: model.provider_id,
key_id: model.key_id,
model_id: model.model_id,
api_kind: model.api_kind,
action: model.action,
prompt: model.prompt,
parameters_json: model.parameters_json,
reference_file_ids_json: model.reference_file_ids_json,
source_image_ids_json: model.source_image_ids_json,
mask_file_id: model.mask_file_id,
status: model.status,
error_message: model.error_message,
response_id: model.response_id,
usage_json: model.usage_json,
created_at: model.created_at,
completed_at: model.completed_at,
images,
reference_files,
source_images,
mask_file,
}
}
fn parse_id_list(raw: &str) -> Vec<String> {
serde_json::from_str::<Vec<String>>(raw).unwrap_or_default()
}
async fn list_reference_files(
db: &DatabaseConnection,
ids_json: &str,
) -> Result<Vec<StoredFile>> {
let mut files = Vec::new();
for id in parse_id_list(ids_json) {
if let Ok(file) = crate::repo::stored_file::get_stored_file(db, &id).await {
files.push(file);
}
}
Ok(files)
}
async fn list_source_images(
db: &DatabaseConnection,
ids_json: &str,
) -> Result<Vec<DrawingImage>> {
let mut images = Vec::new();
for id in parse_id_list(ids_json) {
if let Ok(image) = get_image(db, &id).await {
images.push(image);
}
}
Ok(images)
}
async fn get_mask_file(
db: &DatabaseConnection,
mask_file_id: Option<&str>,
) -> Result<Option<StoredFile>> {
let Some(mask_file_id) = mask_file_id else {
return Ok(None);
};
Ok(crate::repo::stored_file::get_stored_file(db, mask_file_id)
.await
.ok())
}
async fn hydrate_generation(
db: &DatabaseConnection,
row: drawing_generations::Model,
) -> Result<DrawingGeneration> {
let id = row.id.clone();
let images = list_images_for_generation(db, &id).await?;
let reference_files = list_reference_files(db, &row.reference_file_ids_json).await?;
let source_images = list_source_images(db, &row.source_image_ids_json).await?;
let mask_file = get_mask_file(db, row.mask_file_id.as_deref()).await?;
Ok(generation_from_entity(
row,
images,
reference_files,
source_images,
mask_file,
))
}
pub async fn create_generation(
db: &DatabaseConnection,
input: NewDrawingGeneration,
) -> Result<DrawingGeneration> {
let id = gen_id();
let now = now_ts();
drawing_generations::ActiveModel {
id: Set(id.clone()),
parent_generation_id: Set(input.parent_generation_id),
provider_id: Set(input.provider_id),
key_id: Set(input.key_id),
model_id: Set(input.model_id),
api_kind: Set("image_api".to_string()),
action: Set(input.action),
prompt: Set(input.prompt),
parameters_json: Set(input.parameters_json),
reference_file_ids_json: Set(input.reference_file_ids_json),
source_image_ids_json: Set(input.source_image_ids_json),
mask_file_id: Set(input.mask_file_id),
status: Set("running".to_string()),
error_message: Set(None),
response_id: Set(None),
usage_json: Set(None),
created_at: Set(now),
completed_at: Set(None),
}
.insert(db)
.await?;
get_generation(db, &id).await
}
pub async fn add_image(db: &DatabaseConnection, input: NewDrawingImage) -> Result<DrawingImage> {
let id = gen_id();
let now = now_ts();
drawing_images::ActiveModel {
id: Set(id.clone()),
generation_id: Set(input.generation_id),
stored_file_id: Set(input.stored_file_id),
storage_path: Set(input.storage_path),
mime_type: Set(input.mime_type),
width: Set(input.width),
height: Set(input.height),
revised_prompt: Set(input.revised_prompt),
created_at: Set(now),
}
.insert(db)
.await?;
get_image(db, &id).await
}
pub async fn mark_generation_succeeded(
db: &DatabaseConnection,
id: &str,
response_id: Option<String>,
usage_json: Option<String>,
) -> Result<()> {
let row = drawing_generations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("DrawingGeneration {}", id)))?;
let mut am: drawing_generations::ActiveModel = row.into();
am.status = Set("succeeded".to_string());
am.error_message = Set(None);
am.response_id = Set(response_id);
am.usage_json = Set(usage_json);
am.completed_at = Set(Some(now_ts()));
am.update(db).await?;
Ok(())
}
pub async fn mark_generation_failed(
db: &DatabaseConnection,
id: &str,
error_message: String,
) -> Result<()> {
let row = drawing_generations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("DrawingGeneration {}", id)))?;
let mut am: drawing_generations::ActiveModel = row.into();
am.status = Set("failed".to_string());
am.error_message = Set(Some(error_message));
am.completed_at = Set(Some(now_ts()));
am.update(db).await?;
Ok(())
}
pub async fn get_image(db: &DatabaseConnection, id: &str) -> Result<DrawingImage> {
let row = drawing_images::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("DrawingImage {}", id)))?;
Ok(image_from_entity(row))
}
pub async fn get_generation(db: &DatabaseConnection, id: &str) -> Result<DrawingGeneration> {
let row = drawing_generations::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("DrawingGeneration {}", id)))?;
hydrate_generation(db, row).await
}
pub async fn list_images_for_generation(
db: &DatabaseConnection,
generation_id: &str,
) -> Result<Vec<DrawingImage>> {
let rows = drawing_images::Entity::find()
.filter(drawing_images::Column::GenerationId.eq(generation_id))
.order_by_asc(drawing_images::Column::CreatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(image_from_entity).collect())
}
pub async fn list_generations(
db: &DatabaseConnection,
limit: u64,
cursor: Option<i64>,
) -> Result<Vec<DrawingGeneration>> {
let mut query = drawing_generations::Entity::find()
.order_by_desc(drawing_generations::Column::CreatedAt)
.limit(limit.min(100));
if let Some(cursor) = cursor {
query = query.filter(drawing_generations::Column::CreatedAt.lt(cursor));
}
let rows = query.all(db).await?;
let mut result = Vec::with_capacity(rows.len());
for row in rows {
result.push(hydrate_generation(db, row).await?);
}
Ok(result)
}
pub async fn delete_generation(db: &DatabaseConnection, id: &str) -> Result<()> {
drawing_images::Entity::delete_many()
.filter(drawing_images::Column::GenerationId.eq(id))
.exec(db)
.await?;
let result = drawing_generations::Entity::delete_by_id(id)
.exec(db)
.await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("DrawingGeneration {}", id)));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::create_test_pool;
use crate::repo::stored_file;
#[tokio::test]
async fn hydrates_reference_source_and_mask_files() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let source_file = stored_file::create_stored_file(
db,
&gen_id(),
"source-hash",
"source.png",
"image/png",
1024,
"images/source.png",
None,
)
.await
.unwrap();
let ref_file = stored_file::create_stored_file(
db,
&gen_id(),
"ref-hash",
"ref.png",
"image/png",
2048,
"images/ref.png",
None,
)
.await
.unwrap();
let mask_file = stored_file::create_stored_file(
db,
&gen_id(),
"mask-hash",
"mask.png",
"image/png",
512,
"images/mask.png",
None,
)
.await
.unwrap();
let source_generation = create_generation(
db,
NewDrawingGeneration {
parent_generation_id: None,
provider_id: "provider-1".into(),
key_id: "key-1".into(),
model_id: "gpt-image-2".into(),
action: "generate".into(),
prompt: "source".into(),
parameters_json: "{}".into(),
reference_file_ids_json: "[]".into(),
source_image_ids_json: "[]".into(),
mask_file_id: None,
},
)
.await
.unwrap();
let source_image = add_image(
db,
NewDrawingImage {
generation_id: source_generation.id.clone(),
stored_file_id: source_file.id.clone(),
storage_path: source_file.storage_path.clone(),
mime_type: source_file.mime_type.clone(),
width: Some(1024),
height: Some(1024),
revised_prompt: None,
},
)
.await
.unwrap();
let edit_generation = create_generation(
db,
NewDrawingGeneration {
parent_generation_id: Some(source_generation.id),
provider_id: "provider-1".into(),
key_id: "key-1".into(),
model_id: "gpt-image-2".into(),
action: "mask_edit".into(),
prompt: "edit".into(),
parameters_json: "{}".into(),
reference_file_ids_json: serde_json::to_string(&vec![ref_file.id.clone()]).unwrap(),
source_image_ids_json: serde_json::to_string(&vec![source_image.id.clone()])
.unwrap(),
mask_file_id: Some(mask_file.id.clone()),
},
)
.await
.unwrap();
let fetched = get_generation(db, &edit_generation.id).await.unwrap();
assert_eq!(fetched.reference_files.len(), 1);
assert_eq!(fetched.reference_files[0].id, ref_file.id);
assert_eq!(fetched.source_images.len(), 1);
assert_eq!(fetched.source_images[0].id, source_image.id);
assert_eq!(
fetched.mask_file.as_ref().map(|file| file.id.as_str()),
Some(mask_file.id.as_str())
);
}
}
@@ -0,0 +1,354 @@
use sea_orm::*;
use crate::crypto;
use crate::entity::{gateway_keys, gateway_usage};
use crate::error::{BoosApiError, Result};
use crate::types::*;
use crate::utils::now_ts;
fn key_from_entity(m: gateway_keys::Model) -> GatewayKey {
GatewayKey {
id: m.id,
name: m.name,
key_hash: m.key_hash,
key_prefix: m.key_prefix,
enabled: m.enabled != 0,
created_at: m.created_at,
last_used_at: m.last_used_at,
has_encrypted_key: m.encrypted_key.is_some(),
}
}
// --- Gateway Key CRUD ---
pub async fn list_gateway_keys(db: &DatabaseConnection) -> Result<Vec<GatewayKey>> {
let rows = gateway_keys::Entity::find()
.order_by_desc(gateway_keys::Column::CreatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(key_from_entity).collect())
}
pub async fn create_gateway_key(
db: &DatabaseConnection,
name: &str,
master_key: Option<&[u8; 32]>,
) -> Result<CreateGatewayKeyResult> {
crate::repo::gateway_key::create_gateway_key(db, name, master_key).await
}
pub async fn delete_gateway_key(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = gateway_keys::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("GatewayKey {}", id)));
}
Ok(())
}
pub async fn toggle_gateway_key(db: &DatabaseConnection, id: &str, enabled: bool) -> Result<()> {
let row = gateway_keys::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("GatewayKey {}", id)))?;
let mut am: gateway_keys::ActiveModel = row.into();
am.enabled = Set(if enabled { 1 } else { 0 });
am.update(db).await?;
Ok(())
}
/// Verify an incoming API key against stored hashes. Returns the matching key if found.
pub async fn verify_key(db: &DatabaseConnection, plain_key: &str) -> Result<GatewayKey> {
let key_hash = crypto::sha256_hash(plain_key);
let row = gateway_keys::Entity::find()
.filter(gateway_keys::Column::KeyHash.eq(&key_hash))
.filter(gateway_keys::Column::Enabled.eq(1))
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound("Invalid or disabled gateway key".to_string()))?;
Ok(key_from_entity(row))
}
pub async fn update_last_used(db: &DatabaseConnection, id: &str) -> Result<()> {
if let Some(row) = gateway_keys::Entity::find_by_id(id).one(db).await? {
let mut am: gateway_keys::ActiveModel = row.into();
am.last_used_at = Set(Some(now_ts()));
am.update(db).await?;
}
Ok(())
}
// --- Gateway Usage ---
pub async fn record_usage(
db: &DatabaseConnection,
key_id: &str,
provider_id: &str,
model_id: Option<&str>,
request_tokens: u64,
response_tokens: u64,
) -> Result<()> {
gateway_usage::ActiveModel {
key_id: Set(key_id.to_string()),
provider_id: Set(provider_id.to_string()),
model_id: Set(model_id.map(|s| s.to_string())),
request_tokens: Set(request_tokens as i32),
response_tokens: Set(response_tokens as i32),
created_at: Set(now_ts()),
..Default::default()
}
.insert(db)
.await?;
Ok(())
}
pub async fn get_gateway_metrics(db: &DatabaseConnection) -> Result<GatewayMetrics> {
let today_start = chrono::Utc::now()
.date_naive()
.and_hms_opt(0, 0, 0)
.unwrap()
.and_utc()
.timestamp();
#[derive(Debug, FromQueryResult)]
struct MetricsRow {
total_requests: i64,
total_tokens: i64,
request_tokens: i64,
response_tokens: i64,
}
let all = MetricsRow::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT COUNT(*) as total_requests, \
COALESCE(SUM(request_tokens + response_tokens), 0) as total_tokens, \
COALESCE(SUM(request_tokens), 0) as request_tokens, \
COALESCE(SUM(response_tokens), 0) as response_tokens \
FROM gateway_usage",
))
.one(db)
.await?
.unwrap_or(MetricsRow {
total_requests: 0,
total_tokens: 0,
request_tokens: 0,
response_tokens: 0,
});
let today = MetricsRow::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT COUNT(*) as total_requests, \
COALESCE(SUM(request_tokens + response_tokens), 0) as total_tokens, \
COALESCE(SUM(request_tokens), 0) as request_tokens, \
COALESCE(SUM(response_tokens), 0) as response_tokens \
FROM gateway_usage WHERE created_at >= ?",
[today_start.into()],
))
.one(db)
.await?
.unwrap_or(MetricsRow {
total_requests: 0,
total_tokens: 0,
request_tokens: 0,
response_tokens: 0,
});
Ok(GatewayMetrics {
total_requests: all.total_requests as u64,
total_tokens: all.total_tokens as u64,
total_request_tokens: all.request_tokens as u64,
total_response_tokens: all.response_tokens as u64,
active_connections: 0,
today_requests: today.total_requests as u64,
today_tokens: today.total_tokens as u64,
today_request_tokens: today.request_tokens as u64,
today_response_tokens: today.response_tokens as u64,
})
}
pub async fn get_usage_by_key(db: &DatabaseConnection) -> Result<Vec<UsageByKey>> {
#[derive(Debug, FromQueryResult)]
struct Row {
key_id: String,
key_name: String,
request_count: i64,
token_count: i64,
request_tokens: i64,
response_tokens: i64,
}
let rows = Row::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT gu.key_id, gk.name as key_name, \
COUNT(*) as request_count, \
COALESCE(SUM(gu.request_tokens + gu.response_tokens), 0) as token_count, \
COALESCE(SUM(gu.request_tokens), 0) as request_tokens, \
COALESCE(SUM(gu.response_tokens), 0) as response_tokens \
FROM gateway_usage gu \
JOIN gateway_keys gk ON gk.id = gu.key_id \
GROUP BY gu.key_id \
ORDER BY token_count DESC",
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByKey {
key_id: r.key_id,
key_name: r.key_name,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
request_tokens: r.request_tokens as u64,
response_tokens: r.response_tokens as u64,
})
.collect())
}
pub async fn get_usage_by_provider(db: &DatabaseConnection) -> Result<Vec<UsageByProvider>> {
#[derive(Debug, FromQueryResult)]
struct Row {
provider_id: String,
provider_name: String,
request_count: i64,
token_count: i64,
request_tokens: i64,
response_tokens: i64,
}
let rows = Row::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT gu.provider_id, COALESCE(p.name, gu.provider_id) as provider_name, \
COUNT(*) as request_count, \
COALESCE(SUM(gu.request_tokens + gu.response_tokens), 0) as token_count, \
COALESCE(SUM(gu.request_tokens), 0) as request_tokens, \
COALESCE(SUM(gu.response_tokens), 0) as response_tokens \
FROM gateway_usage gu \
LEFT JOIN providers p ON p.id = gu.provider_id \
GROUP BY gu.provider_id \
ORDER BY token_count DESC",
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByProvider {
provider_id: r.provider_id,
provider_name: r.provider_name,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
request_tokens: r.request_tokens as u64,
response_tokens: r.response_tokens as u64,
})
.collect())
}
pub async fn get_usage_by_day(db: &DatabaseConnection, days: u32) -> Result<Vec<UsageByDay>> {
#[derive(Debug, FromQueryResult)]
struct Row {
date: String,
request_count: i64,
token_count: i64,
request_tokens: i64,
response_tokens: i64,
}
let since = now_ts() - (days as i64 * 86400);
let rows = Row::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT date(created_at, 'unixepoch') as date, \
COUNT(*) as request_count, \
COALESCE(SUM(request_tokens + response_tokens), 0) as token_count, \
COALESCE(SUM(request_tokens), 0) as request_tokens, \
COALESCE(SUM(response_tokens), 0) as response_tokens \
FROM gateway_usage \
WHERE created_at >= ? \
GROUP BY date \
ORDER BY date ASC",
[since.into()],
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByDay {
date: r.date,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
request_tokens: r.request_tokens as u64,
response_tokens: r.response_tokens as u64,
})
.collect())
}
pub async fn get_connected_programs(db: &DatabaseConnection) -> Result<Vec<ConnectedProgram>> {
let today_start = chrono::Utc::now()
.date_naive()
.and_hms_opt(0, 0, 0)
.unwrap()
.and_utc()
.timestamp();
let active_threshold = now_ts() - 300;
#[derive(Debug, FromQueryResult)]
struct Row {
key_id: String,
key_name: String,
key_prefix: String,
today_requests: i64,
today_tokens: i64,
today_request_tokens: i64,
today_response_tokens: i64,
last_active_at: Option<i64>,
}
let rows = Row::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT gk.id as key_id, gk.name as key_name, gk.key_prefix, \
COALESCE(t.cnt, 0) as today_requests, \
COALESCE(t.tokens, 0) as today_tokens, \
COALESCE(t.request_tokens, 0) as today_request_tokens, \
COALESCE(t.response_tokens, 0) as today_response_tokens, \
gk.last_used_at as last_active_at \
FROM gateway_keys gk \
LEFT JOIN ( \
SELECT key_id, COUNT(*) as cnt, \
SUM(request_tokens + response_tokens) as tokens, \
SUM(request_tokens) as request_tokens, \
SUM(response_tokens) as response_tokens \
FROM gateway_usage WHERE created_at >= ? \
GROUP BY key_id \
) t ON t.key_id = gk.id \
WHERE gk.enabled = 1 \
ORDER BY gk.created_at DESC",
[today_start.into()],
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| ConnectedProgram {
key_id: r.key_id,
key_name: r.key_name,
key_prefix: r.key_prefix,
today_requests: r.today_requests as u64,
today_tokens: r.today_tokens as u64,
today_request_tokens: r.today_request_tokens as u64,
today_response_tokens: r.today_response_tokens as u64,
last_active_at: r.last_active_at,
is_active: r
.last_active_at
.map(|t| t >= active_threshold)
.unwrap_or(false),
})
.collect())
}
@@ -0,0 +1,39 @@
use crate::error::Result;
use crate::types::GatewayDiagnostic;
use crate::utils::gen_id;
use sea_orm::DatabaseConnection;
pub async fn get_diagnostics(_db: &DatabaseConnection) -> Result<Vec<GatewayDiagnostic>> {
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
Ok(vec![
GatewayDiagnostic {
id: gen_id(),
category: "port".to_string(),
status: "ok".to_string(),
message: "Gateway port 8080 is available".to_string(),
created_at: now.clone(),
},
GatewayDiagnostic {
id: gen_id(),
category: "auth".to_string(),
status: "ok".to_string(),
message: "API key authentication is configured".to_string(),
created_at: now.clone(),
},
GatewayDiagnostic {
id: gen_id(),
category: "proxy".to_string(),
status: "ok".to_string(),
message: "Proxy settings are valid".to_string(),
created_at: now.clone(),
},
GatewayDiagnostic {
id: gen_id(),
category: "provider_latency".to_string(),
status: "warning".to_string(),
message: "No providers configured for latency testing".to_string(),
created_at: now,
},
])
}
@@ -0,0 +1,136 @@
use sea_orm::*;
use crate::crypto;
use crate::entity::gateway_keys;
use crate::error::{BoosApiError, Result};
use crate::types::{CreateGatewayKeyResult, GatewayKey};
use crate::utils::{gen_id, now_ts};
fn key_from_entity(m: gateway_keys::Model) -> GatewayKey {
GatewayKey {
id: m.id,
name: m.name,
key_hash: m.key_hash,
key_prefix: m.key_prefix,
enabled: m.enabled != 0,
created_at: m.created_at,
last_used_at: m.last_used_at,
has_encrypted_key: m.encrypted_key.is_some(),
}
}
pub async fn list_gateway_keys(db: &DatabaseConnection) -> Result<Vec<GatewayKey>> {
let rows = gateway_keys::Entity::find()
.order_by_desc(gateway_keys::Column::CreatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(key_from_entity).collect())
}
pub async fn create_gateway_key(
db: &DatabaseConnection,
name: &str,
master_key: Option<&[u8; 32]>,
) -> Result<CreateGatewayKeyResult> {
let id = gen_id();
let now = now_ts();
let plain_key = crypto::generate_gateway_key();
let key_hash = crypto::sha256_hash(&plain_key);
let key_prefix = crypto::key_prefix(&plain_key);
let encrypted_key = master_key
.map(|mk| crypto::encrypt_key(&plain_key, mk))
.transpose()?;
gateway_keys::ActiveModel {
id: Set(id.clone()),
name: Set(name.to_string()),
key_hash: Set(key_hash),
key_prefix: Set(key_prefix),
enabled: Set(1),
created_at: Set(now),
last_used_at: Set(None),
encrypted_key: Set(encrypted_key),
}
.insert(db)
.await?;
let row = gateway_keys::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("GatewayKey {}", id)))?;
Ok(CreateGatewayKeyResult {
gateway_key: key_from_entity(row),
plain_key,
})
}
pub async fn delete_gateway_key(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = gateway_keys::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("GatewayKey {}", id)));
}
Ok(())
}
pub async fn toggle_gateway_key(db: &DatabaseConnection, id: &str) -> Result<GatewayKey> {
let row = gateway_keys::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("GatewayKey {}", id)))?;
let new_enabled = if row.enabled != 0 { 0 } else { 1 };
let mut am: gateway_keys::ActiveModel = row.into();
am.enabled = Set(new_enabled);
am.update(db).await?;
let row = gateway_keys::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("GatewayKey {}", id)))?;
Ok(key_from_entity(row))
}
/// Verify an incoming API key against stored hashes. Returns the matching key if found.
pub async fn verify_key(db: &DatabaseConnection, plain_key: &str) -> Result<GatewayKey> {
let key_hash = crypto::sha256_hash(plain_key);
let row = gateway_keys::Entity::find()
.filter(gateway_keys::Column::KeyHash.eq(&key_hash))
.filter(gateway_keys::Column::Enabled.eq(1))
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound("Invalid or disabled gateway key".to_string()))?;
Ok(key_from_entity(row))
}
pub async fn update_last_used(db: &DatabaseConnection, id: &str) -> Result<()> {
if let Some(row) = gateway_keys::Entity::find_by_id(id).one(db).await? {
let mut am: gateway_keys::ActiveModel = row.into();
am.last_used_at = Set(Some(now_ts()));
am.update(db).await?;
}
Ok(())
}
/// Decrypt and return the plain key for a given key ID.
pub async fn get_plain_key(
db: &DatabaseConnection,
master_key: &[u8; 32],
key_id: &str,
) -> Result<String> {
let row = gateway_keys::Entity::find_by_id(key_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("GatewayKey {}", key_id)))?;
let encrypted = row.encrypted_key.ok_or_else(|| {
BoosApiError::Crypto("Key was created before encrypted storage was available".to_string())
})?;
crypto::decrypt_key(&encrypted, master_key)
}
@@ -0,0 +1,76 @@
use crate::entity::gateway_request_logs;
use crate::error::Result;
use crate::types::GatewayRequestLog;
use crate::utils::{gen_id, now_ts};
use sea_orm::*;
pub async fn record_request_log(
db: &DatabaseConnection,
key_id: &str,
key_name: &str,
method: &str,
path: &str,
model: Option<&str>,
provider_id: Option<&str>,
status_code: i32,
duration_ms: i32,
request_tokens: i32,
response_tokens: i32,
error_message: Option<&str>,
) -> Result<()> {
gateway_request_logs::ActiveModel {
id: Set(gen_id()),
key_id: Set(key_id.to_string()),
key_name: Set(key_name.to_string()),
method: Set(method.to_string()),
path: Set(path.to_string()),
model: Set(model.map(|s| s.to_string())),
provider_id: Set(provider_id.map(|s| s.to_string())),
status_code: Set(status_code),
duration_ms: Set(duration_ms),
request_tokens: Set(request_tokens),
response_tokens: Set(response_tokens),
error_message: Set(error_message.map(|s| s.to_string())),
created_at: Set(now_ts()),
}
.insert(db)
.await?;
Ok(())
}
pub async fn list_request_logs(
db: &DatabaseConnection,
limit: u64,
offset: u64,
) -> Result<Vec<GatewayRequestLog>> {
let rows = gateway_request_logs::Entity::find()
.order_by_desc(gateway_request_logs::Column::CreatedAt)
.limit(limit)
.offset(offset)
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| GatewayRequestLog {
id: r.id,
key_id: r.key_id,
key_name: r.key_name,
method: r.method,
path: r.path,
model: r.model,
provider_id: r.provider_id,
status_code: r.status_code,
duration_ms: r.duration_ms,
request_tokens: r.request_tokens,
response_tokens: r.response_tokens,
error_message: r.error_message,
created_at: r.created_at,
})
.collect())
}
pub async fn clear_request_logs(db: &DatabaseConnection) -> Result<u64> {
let result = gateway_request_logs::Entity::delete_many().exec(db).await?;
Ok(result.rows_affected)
}
@@ -0,0 +1,228 @@
use sea_orm::*;
use crate::entity::gateway_usage;
use crate::error::Result;
use crate::types::*;
use crate::utils::now_ts;
pub async fn record_usage(
db: &DatabaseConnection,
key_id: &str,
provider_id: &str,
model_id: Option<&str>,
request_tokens: u64,
response_tokens: u64,
) -> Result<()> {
gateway_usage::ActiveModel {
key_id: Set(key_id.to_string()),
provider_id: Set(provider_id.to_string()),
model_id: Set(model_id.map(|s| s.to_string())),
request_tokens: Set(request_tokens as i32),
response_tokens: Set(response_tokens as i32),
created_at: Set(now_ts()),
..Default::default()
}
.insert(db)
.await?;
Ok(())
}
pub async fn get_metrics(db: &DatabaseConnection) -> Result<GatewayMetrics> {
let today_start = chrono::Utc::now()
.date_naive()
.and_hms_opt(0, 0, 0)
.unwrap()
.and_utc()
.timestamp();
#[derive(Debug, FromQueryResult)]
struct MetricsRow {
total_requests: i64,
total_tokens: i64,
}
let all = MetricsRow::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT COUNT(*) as total_requests, \
COALESCE(SUM(request_tokens + response_tokens), 0) as total_tokens \
FROM gateway_usage",
))
.one(db)
.await?
.unwrap_or(MetricsRow { total_requests: 0, total_tokens: 0 });
let today = MetricsRow::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT COUNT(*) as total_requests, \
COALESCE(SUM(request_tokens + response_tokens), 0) as total_tokens \
FROM gateway_usage WHERE created_at >= ?",
[today_start.into()],
))
.one(db)
.await?
.unwrap_or(MetricsRow { total_requests: 0, total_tokens: 0 });
Ok(GatewayMetrics {
total_requests: all.total_requests as u64,
total_tokens: all.total_tokens as u64,
active_connections: 0, // runtime state, not tracked in DB
today_requests: today.total_requests as u64,
today_tokens: today.total_tokens as u64,
})
}
pub async fn get_usage_by_key(db: &DatabaseConnection) -> Result<Vec<UsageByKey>> {
#[derive(Debug, FromQueryResult)]
struct Row {
key_id: String,
key_name: String,
request_count: i64,
token_count: i64,
}
let rows = Row::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT gu.key_id, gk.name as key_name, \
COUNT(*) as request_count, \
COALESCE(SUM(gu.request_tokens + gu.response_tokens), 0) as token_count \
FROM gateway_usage gu \
JOIN gateway_keys gk ON gk.id = gu.key_id \
GROUP BY gu.key_id \
ORDER BY token_count DESC",
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByKey {
key_id: r.key_id,
key_name: r.key_name,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
})
.collect())
}
pub async fn get_usage_by_provider(db: &DatabaseConnection) -> Result<Vec<UsageByProvider>> {
#[derive(Debug, FromQueryResult)]
struct Row {
provider_id: String,
provider_name: String,
request_count: i64,
token_count: i64,
}
let rows = Row::find_by_statement(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT gu.provider_id, COALESCE(p.name, gu.provider_id) as provider_name, \
COUNT(*) as request_count, \
COALESCE(SUM(gu.request_tokens + gu.response_tokens), 0) as token_count \
FROM gateway_usage gu \
LEFT JOIN providers p ON p.id = gu.provider_id \
GROUP BY gu.provider_id \
ORDER BY token_count DESC",
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByProvider {
provider_id: r.provider_id,
provider_name: r.provider_name,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
})
.collect())
}
pub async fn get_usage_by_day(db: &DatabaseConnection, days: u32) -> Result<Vec<UsageByDay>> {
#[derive(Debug, FromQueryResult)]
struct Row {
date: String,
request_count: i64,
token_count: i64,
}
let since = now_ts() - (days as i64 * 86400);
let rows = Row::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT date(created_at, 'unixepoch') as date, \
COUNT(*) as request_count, \
COALESCE(SUM(request_tokens + response_tokens), 0) as token_count \
FROM gateway_usage \
WHERE created_at >= ? \
GROUP BY date \
ORDER BY date ASC",
[since.into()],
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| UsageByDay {
date: r.date,
request_count: r.request_count as u64,
token_count: r.token_count as u64,
})
.collect())
}
pub async fn get_connected_programs(db: &DatabaseConnection) -> Result<Vec<ConnectedProgram>> {
let today_start = chrono::Utc::now()
.date_naive()
.and_hms_opt(0, 0, 0)
.unwrap()
.and_utc()
.timestamp();
let active_threshold = now_ts() - 300; // active within last 5 minutes
#[derive(Debug, FromQueryResult)]
struct Row {
key_id: String,
key_name: String,
key_prefix: String,
today_requests: i64,
today_tokens: i64,
last_active_at: Option<i64>,
}
let rows = Row::find_by_statement(Statement::from_sql_and_values(
DatabaseBackend::Sqlite,
"SELECT gk.id as key_id, gk.name as key_name, gk.key_prefix, \
COALESCE(t.cnt, 0) as today_requests, \
COALESCE(t.tokens, 0) as today_tokens, \
gk.last_used_at as last_active_at \
FROM gateway_keys gk \
LEFT JOIN ( \
SELECT key_id, COUNT(*) as cnt, \
SUM(request_tokens + response_tokens) as tokens \
FROM gateway_usage WHERE created_at >= ? \
GROUP BY key_id \
) t ON t.key_id = gk.id \
WHERE gk.enabled = 1 \
ORDER BY gk.created_at DESC",
[today_start.into()],
))
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|r| ConnectedProgram {
key_id: r.key_id,
key_name: r.key_name,
key_prefix: r.key_prefix,
today_requests: r.today_requests as u64,
today_tokens: r.today_tokens as u64,
last_active_at: r.last_active_at,
is_active: r
.last_active_at
.map(|t| t >= active_threshold)
.unwrap_or(false),
})
.collect())
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,274 @@
use sea_orm::sea_query::Expr;
use sea_orm::*;
use crate::entity::{knowledge_bases, knowledge_documents};
use crate::error::{BoosApiError, Result};
use crate::types::{
CreateKnowledgeBaseInput, KnowledgeBase, KnowledgeDocument, UpdateKnowledgeBaseInput,
};
use crate::utils::gen_id;
fn model_to_kb(m: knowledge_bases::Model) -> KnowledgeBase {
KnowledgeBase {
id: m.id,
name: m.name,
description: m.description,
embedding_provider: m.embedding_provider,
enabled: m.enabled != 0,
icon_type: m.icon_type,
icon_value: m.icon_value,
sort_order: m.sort_order,
embedding_dimensions: m.embedding_dimensions,
retrieval_threshold: m.retrieval_threshold,
retrieval_top_k: m.retrieval_top_k,
rerank_provider: m.rerank_provider,
rerank_candidate_k: m.rerank_candidate_k,
chunk_size: m.chunk_size,
chunk_overlap: m.chunk_overlap,
separator: m.separator,
index_concurrency: m.index_concurrency,
index_interval_ms: m.index_interval_ms,
}
}
fn model_to_doc(m: knowledge_documents::Model) -> KnowledgeDocument {
KnowledgeDocument {
id: m.id,
knowledge_base_id: m.knowledge_base_id,
title: m.title,
source_path: m.source_path,
mime_type: m.mime_type,
size_bytes: m.size_bytes,
indexing_status: m.indexing_status,
doc_type: m.doc_type,
index_error: m.index_error,
}
}
pub async fn list_knowledge_bases(db: &DatabaseConnection) -> Result<Vec<KnowledgeBase>> {
let models = knowledge_bases::Entity::find()
.order_by_asc(knowledge_bases::Column::SortOrder)
.order_by_asc(knowledge_bases::Column::Name)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_kb).collect())
}
pub async fn get_knowledge_base(db: &DatabaseConnection, id: &str) -> Result<KnowledgeBase> {
let model = knowledge_bases::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("KnowledgeBase {}", id)))?;
Ok(model_to_kb(model))
}
pub async fn create_knowledge_base(
db: &DatabaseConnection,
input: CreateKnowledgeBaseInput,
) -> Result<KnowledgeBase> {
let id = gen_id();
let am = knowledge_bases::ActiveModel {
id: Set(id.clone()),
name: Set(input.name),
description: Set(input.description),
embedding_provider: Set(input.embedding_provider),
enabled: Set(if input.enabled.unwrap_or(true) { 1 } else { 0 }),
icon_type: Set(None),
icon_value: Set(None),
sort_order: Set(0),
embedding_dimensions: Set(None),
retrieval_threshold: Set(None),
retrieval_top_k: Set(None),
rerank_provider: Set(None),
rerank_candidate_k: Set(None),
chunk_size: Set(None),
chunk_overlap: Set(None),
separator: Set(None),
index_concurrency: Set(None),
index_interval_ms: Set(None),
};
am.insert(db).await?;
get_knowledge_base(db, &id).await
}
pub async fn update_knowledge_base(
db: &DatabaseConnection,
id: &str,
input: UpdateKnowledgeBaseInput,
) -> Result<KnowledgeBase> {
let model = knowledge_bases::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("KnowledgeBase {}", id)))?;
let existing = model_to_kb(model.clone());
let mut am: knowledge_bases::ActiveModel = model.into();
am.name = Set(input.name.unwrap_or(existing.name));
am.description = Set(input.description.or(existing.description));
am.embedding_provider = Set(input.embedding_provider.or(existing.embedding_provider));
am.enabled = Set(if input.enabled.unwrap_or(existing.enabled) {
1
} else {
0
});
if input.update_icon {
am.icon_type = Set(input.icon_type);
am.icon_value = Set(input.icon_value);
}
if input.update_embedding_dimensions {
am.embedding_dimensions = Set(input.embedding_dimensions);
}
if input.update_retrieval_threshold {
am.retrieval_threshold = Set(input.retrieval_threshold);
}
if input.update_retrieval_top_k {
am.retrieval_top_k = Set(input.retrieval_top_k);
}
if input.update_rerank_provider {
am.rerank_provider = Set(input.rerank_provider);
}
if input.update_rerank_candidate_k {
am.rerank_candidate_k = Set(input.rerank_candidate_k);
}
if input.update_chunk_size {
am.chunk_size = Set(input.chunk_size);
}
if input.update_chunk_overlap {
am.chunk_overlap = Set(input.chunk_overlap);
}
if input.update_separator {
am.separator = Set(input.separator);
}
if input.update_index_concurrency {
am.index_concurrency = Set(input.index_concurrency);
}
if input.update_index_interval_ms {
am.index_interval_ms = Set(input.index_interval_ms);
}
am.update(db).await?;
get_knowledge_base(db, id).await
}
pub async fn reorder_knowledge_bases(db: &DatabaseConnection, base_ids: &[String]) -> Result<()> {
for (i, id) in base_ids.iter().enumerate() {
knowledge_bases::Entity::update_many()
.col_expr(knowledge_bases::Column::SortOrder, Expr::value(i as i32))
.filter(knowledge_bases::Column::Id.eq(id))
.exec(db)
.await?;
}
Ok(())
}
pub async fn delete_knowledge_base(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = knowledge_bases::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("KnowledgeBase {}", id)));
}
Ok(())
}
pub async fn list_documents(
db: &DatabaseConnection,
base_id: &str,
) -> Result<Vec<KnowledgeDocument>> {
let models = knowledge_documents::Entity::find()
.filter(knowledge_documents::Column::KnowledgeBaseId.eq(base_id))
.order_by_asc(knowledge_documents::Column::Title)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_doc).collect())
}
pub async fn add_document(
db: &DatabaseConnection,
knowledge_base_id: &str,
title: &str,
source_path: &str,
mime_type: &str,
doc_type: Option<&str>,
) -> Result<KnowledgeDocument> {
let id = gen_id();
// Read actual file size from disk
let file_size = std::fs::metadata(source_path)
.map(|m| m.len() as i64)
.unwrap_or(0);
let am = knowledge_documents::ActiveModel {
id: Set(id.clone()),
knowledge_base_id: Set(knowledge_base_id.to_string()),
title: Set(title.to_string()),
source_path: Set(source_path.to_string()),
mime_type: Set(mime_type.to_string()),
size_bytes: Set(file_size),
doc_type: Set(doc_type.unwrap_or("file").to_string()),
..Default::default()
};
am.insert(db).await?;
let model = knowledge_documents::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("KnowledgeDocument {}", id)))?;
Ok(model_to_doc(model))
}
pub async fn update_document_status(db: &DatabaseConnection, id: &str, status: &str) -> Result<()> {
update_document_status_with_error(db, id, status, None).await
}
pub async fn update_document_status_with_error(
db: &DatabaseConnection,
id: &str,
status: &str,
error: Option<&str>,
) -> Result<()> {
let mut am: knowledge_documents::ActiveModel = knowledge_documents::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("KnowledgeDocument {}", id)))?
.into();
am.indexing_status = Set(status.to_string());
am.index_error = Set(error.map(|e| e.to_string()));
am.update(db).await?;
Ok(())
}
pub async fn delete_document(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = knowledge_documents::Entity::delete_by_id(id)
.exec(db)
.await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("KnowledgeDocument {}", id)));
}
Ok(())
}
/// Batch lookup document titles by IDs. Returns a map of document_id -> title.
pub async fn get_document_titles(
db: &DatabaseConnection,
doc_ids: &[String],
) -> Result<std::collections::HashMap<String, String>> {
if doc_ids.is_empty() {
return Ok(std::collections::HashMap::new());
}
let models = knowledge_documents::Entity::find()
.filter(knowledge_documents::Column::Id.is_in(doc_ids.iter().map(|s| s.as_str())))
.all(db)
.await?;
Ok(models.into_iter().map(|m| (m.id, m.title)).collect())
}
@@ -0,0 +1,515 @@
use sea_orm::*;
use crate::entity::{mcp_servers, tool_descriptors};
use crate::error::{BoosApiError, Result};
use crate::repo::settings;
use crate::types::{CreateMcpServerInput, McpServer, ToolDescriptor, UpdateMcpServerInput};
use crate::utils::gen_id;
// ── Builtin server definitions (not stored in DB) ───────────────────────
const BUILTIN_FETCH_ID: &str = "builtin-fetch";
const BUILTIN_SEARCH_FILE_ID: &str = "builtin-search-file";
struct BuiltinDef {
id: &'static str,
name: &'static str,
default_enabled: bool,
}
const BUILTIN_DEFS: &[BuiltinDef] = &[
BuiltinDef {
id: BUILTIN_FETCH_ID,
name: "@boosapi/fetch",
default_enabled: true,
},
BuiltinDef {
id: BUILTIN_SEARCH_FILE_ID,
name: "@boosapi/search-file",
default_enabled: false,
},
];
fn builtin_setting_key(name: &str) -> String {
format!("builtin_mcp:{name}:enabled")
}
fn make_builtin_server(def: &BuiltinDef, enabled: bool) -> McpServer {
McpServer {
id: def.id.to_string(),
name: def.name.to_string(),
transport: "builtin".to_string(),
command: None,
args_json: None,
endpoint: None,
env_json: None,
enabled,
permission_policy: "auto".to_string(),
source: "builtin".to_string(),
discover_timeout_secs: None,
execute_timeout_secs: None,
headers_json: None,
icon_type: None,
icon_value: None,
}
}
async fn get_builtin_enabled(db: &DatabaseConnection, name: &str, default: bool) -> bool {
match settings::get_setting(db, &builtin_setting_key(name)).await {
Ok(Some(v)) => v == "true",
_ => default,
}
}
/// Return all builtin servers with their persisted enabled state.
pub async fn list_builtin_servers(db: &DatabaseConnection) -> Vec<McpServer> {
let mut out = Vec::with_capacity(BUILTIN_DEFS.len());
for def in BUILTIN_DEFS {
let enabled = get_builtin_enabled(db, def.name, def.default_enabled).await;
out.push(make_builtin_server(def, enabled));
}
out
}
/// Check whether a server ID belongs to a builtin.
pub fn is_builtin_id(id: &str) -> bool {
BUILTIN_DEFS.iter().any(|d| d.id == id)
}
/// Toggle enabled state for a builtin server (persists to settings table).
pub async fn set_builtin_enabled(
db: &DatabaseConnection,
id: &str,
enabled: bool,
) -> Result<McpServer> {
let def = BUILTIN_DEFS
.iter()
.find(|d| d.id == id)
.ok_or_else(|| BoosApiError::NotFound(format!("Builtin server {id}")))?;
settings::set_setting(
db,
&builtin_setting_key(def.name),
if enabled { "true" } else { "false" },
)
.await?;
Ok(make_builtin_server(def, enabled))
}
/// Get a single builtin server by ID.
pub async fn get_builtin_server(db: &DatabaseConnection, id: &str) -> Result<McpServer> {
let def = BUILTIN_DEFS
.iter()
.find(|d| d.id == id)
.ok_or_else(|| BoosApiError::NotFound(format!("Builtin server {id}")))?;
let enabled = get_builtin_enabled(db, def.name, def.default_enabled).await;
Ok(make_builtin_server(def, enabled))
}
// ── DB-backed custom servers ────────────────────────────────────────────
fn model_to_mcp_server(m: mcp_servers::Model) -> McpServer {
McpServer {
id: m.id,
name: m.name,
transport: m.transport,
command: m.command,
args_json: m.args_json,
endpoint: m.endpoint,
env_json: m.env_json,
enabled: m.enabled != 0,
permission_policy: m.permission_policy,
source: m.source,
discover_timeout_secs: m.discover_timeout_secs,
execute_timeout_secs: m.execute_timeout_secs,
headers_json: m.headers_json,
icon_type: m.icon_type,
icon_value: m.icon_value,
}
}
pub async fn list_mcp_servers(db: &DatabaseConnection) -> Result<Vec<McpServer>> {
let mut servers = list_builtin_servers(db).await;
let custom_rows = mcp_servers::Entity::find()
.order_by_asc(mcp_servers::Column::Name)
.all(db)
.await?;
servers.extend(custom_rows.into_iter().map(model_to_mcp_server));
Ok(servers)
}
pub async fn get_mcp_server(db: &DatabaseConnection, id: &str) -> Result<McpServer> {
// Check builtins first
if is_builtin_id(id) {
return get_builtin_server(db, id).await;
}
let model = mcp_servers::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("McpServer {}", id)))?;
Ok(model_to_mcp_server(model))
}
pub async fn create_mcp_server(
db: &DatabaseConnection,
input: CreateMcpServerInput,
) -> Result<McpServer> {
let id = gen_id();
let args_json = input
.args
.as_ref()
.map(|a| serde_json::to_string(a).unwrap_or_default());
let env_json = input
.env
.as_ref()
.map(|e| serde_json::to_string(e).unwrap_or_default());
mcp_servers::ActiveModel {
id: Set(id.clone()),
name: Set(input.name),
transport: Set(input.transport),
command: Set(input.command),
args_json: Set(args_json),
endpoint: Set(input.endpoint),
env_json: Set(env_json),
enabled: Set(if input.enabled.unwrap_or(true) { 1 } else { 0 }),
permission_policy: Set(input.permission_policy.unwrap_or_else(|| "ask".to_string())),
source: Set(input.source.unwrap_or_else(|| "custom".to_string())),
discover_timeout_secs: Set(input.discover_timeout_secs),
execute_timeout_secs: Set(input.execute_timeout_secs),
headers_json: Set(input.headers_json),
icon_type: Set(input.icon_type),
icon_value: Set(input.icon_value),
}
.insert(db)
.await?;
get_mcp_server(db, &id).await
}
pub async fn update_mcp_server(
db: &DatabaseConnection,
id: &str,
input: UpdateMcpServerInput,
) -> Result<McpServer> {
// Builtin servers only support toggling enabled
if is_builtin_id(id) {
let enabled = input.enabled.unwrap_or(true);
return set_builtin_enabled(db, id, enabled).await;
}
let model = mcp_servers::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("McpServer {}", id)))?;
let mut am: mcp_servers::ActiveModel = model.into();
if let Some(name) = input.name {
if !name.is_empty() {
am.name = Set(name);
}
}
if let Some(transport) = input.transport {
if !transport.is_empty() {
am.transport = Set(transport);
}
}
if let Some(command) = input.command {
am.command = Set(command.filter(|v| !v.is_empty()));
}
if let Some(args) = input.args {
am.args_json = Set(args.map(|a| serde_json::to_string(&a).unwrap_or_default()));
}
if let Some(endpoint) = input.endpoint {
am.endpoint = Set(endpoint.filter(|v| !v.is_empty()));
}
if let Some(env) = input.env {
am.env_json = Set(env.map(|e| serde_json::to_string(&e).unwrap_or_default()));
}
if let Some(enabled) = input.enabled {
am.enabled = Set(if enabled { 1 } else { 0 });
}
if let Some(permission_policy) = input.permission_policy {
am.permission_policy = Set(permission_policy);
}
if let Some(discover_timeout_secs) = input.discover_timeout_secs {
am.discover_timeout_secs = Set(Some(discover_timeout_secs));
}
if let Some(execute_timeout_secs) = input.execute_timeout_secs {
am.execute_timeout_secs = Set(Some(execute_timeout_secs));
}
if let Some(headers_json) = input.headers_json {
am.headers_json = Set(headers_json.filter(|v| !v.is_empty()));
}
if let Some(icon_type) = input.icon_type {
am.icon_type = Set(icon_type.filter(|v| !v.is_empty()));
}
if let Some(icon_value) = input.icon_value {
am.icon_value = Set(icon_value.filter(|v| !v.is_empty()));
}
am.update(db).await?;
get_mcp_server(db, id).await
}
pub async fn delete_mcp_server(db: &DatabaseConnection, id: &str) -> Result<()> {
// Prevent deletion of built-in MCP servers
let server = get_mcp_server(db, id).await?;
if server.source == "builtin" {
return Err(BoosApiError::Gateway(
"Cannot delete built-in MCP server".to_string(),
));
}
let result = mcp_servers::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("McpServer {}", id)));
}
Ok(())
}
/// Return tool descriptors for a given MCP server.
pub async fn list_tools_for_server(
db: &DatabaseConnection,
server_id: &str,
) -> Result<Vec<ToolDescriptor>> {
// Builtins: resolve name from definition, no DB lookup needed
if let Some(def) = BUILTIN_DEFS.iter().find(|d| d.id == server_id) {
return Ok(builtin_tools(server_id, def.name));
}
// Custom servers: read from tool_descriptors table
let rows = tool_descriptors::Entity::find()
.filter(tool_descriptors::Column::ServerId.eq(server_id))
.order_by_asc(tool_descriptors::Column::Name)
.all(db)
.await?;
Ok(rows
.into_iter()
.map(|m| ToolDescriptor {
id: m.id,
server_id: m.server_id,
name: m.name,
description: m.description,
input_schema_json: m.input_schema_json,
})
.collect())
}
/// Save discovered tool descriptors for a server (replaces existing).
pub async fn save_tool_descriptors(
db: &DatabaseConnection,
server_id: &str,
tools: Vec<crate::mcp_client::DiscoveredTool>,
) -> Result<Vec<ToolDescriptor>> {
// Delete existing tools for this server
tool_descriptors::Entity::delete_many()
.filter(tool_descriptors::Column::ServerId.eq(server_id))
.exec(db)
.await?;
// Insert new tools
let mut result = Vec::with_capacity(tools.len());
for tool in tools {
let id = gen_id();
let input_schema_json = tool
.input_schema
.as_ref()
.map(|s| serde_json::to_string(s).unwrap_or_default());
tool_descriptors::ActiveModel {
id: Set(id.clone()),
server_id: Set(server_id.to_string()),
name: Set(tool.name.clone()),
description: Set(tool.description.clone()),
input_schema_json: Set(input_schema_json.clone()),
}
.insert(db)
.await?;
result.push(ToolDescriptor {
id,
server_id: server_id.to_string(),
name: tool.name,
description: tool.description,
input_schema_json,
});
}
Ok(result)
}
fn builtin_tools(server_id: &str, server_name: &str) -> Vec<ToolDescriptor> {
match server_name {
"@boosapi/fetch" => vec![
ToolDescriptor {
id: format!("{server_id}-fetch-url"),
server_id: server_id.to_string(),
name: "fetch_url".into(),
description: Some("Fetch a URL and return its content".into()),
input_schema_json: Some(r#"{"type":"object","properties":{"url":{"type":"string","description":"URL to fetch"},"max_length":{"type":"integer","description":"Maximum content length"}},"required":["url"]}"#.into()),
},
ToolDescriptor {
id: format!("{server_id}-fetch-markdown"),
server_id: server_id.to_string(),
name: "fetch_markdown".into(),
description: Some("Fetch a URL and convert the content to markdown".into()),
input_schema_json: Some(r#"{"type":"object","properties":{"url":{"type":"string","description":"URL to fetch"}},"required":["url"]}"#.into()),
},
],
"@boosapi/search-file" => vec![
ToolDescriptor {
id: format!("{server_id}-read-file"),
server_id: server_id.to_string(),
name: "read_file".into(),
description: Some("Read the contents of a file".into()),
input_schema_json: Some(r#"{"type":"object","properties":{"path":{"type":"string","description":"File path to read"}},"required":["path"]}"#.into()),
},
ToolDescriptor {
id: format!("{server_id}-list-directory"),
server_id: server_id.to_string(),
name: "list_directory".into(),
description: Some("List files and directories in a given path".into()),
input_schema_json: Some(r#"{"type":"object","properties":{"path":{"type":"string","description":"Directory path to list"}},"required":["path"]}"#.into()),
},
ToolDescriptor {
id: format!("{server_id}-search-files"),
server_id: server_id.to_string(),
name: "search_files".into(),
description: Some("Search for files matching a pattern".into()),
input_schema_json: Some(r#"{"type":"object","properties":{"path":{"type":"string","description":"Base directory"},"pattern":{"type":"string","description":"Search pattern"}},"required":["path","pattern"]}"#.into()),
},
],
_ => vec![],
}
}
/// Find which MCP server owns a given tool, searching across the provided server IDs.
pub async fn find_server_for_tool(
db: &DatabaseConnection,
tool_name: &str,
server_ids: &[String],
) -> Result<Option<(McpServer, ToolDescriptor)>> {
for server_id in server_ids {
if let Ok(tools) = list_tools_for_server(db, server_id).await {
if let Some(td) = tools.into_iter().find(|t| t.name == tool_name) {
if let Ok(server) = get_mcp_server(db, server_id).await {
return Ok(Some((server, td)));
}
}
}
}
Ok(None)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::create_test_pool;
use crate::types::UpdateMcpServerInput;
fn remote_server_input(headers_json: Option<String>) -> CreateMcpServerInput {
CreateMcpServerInput {
name: "Remote MCP".into(),
transport: "http".into(),
command: None,
args: None,
endpoint: Some("https://example.com/mcp".into()),
env: None,
enabled: Some(false),
permission_policy: Some("ask".into()),
source: Some("custom".into()),
discover_timeout_secs: Some(30),
execute_timeout_secs: Some(30),
headers_json,
icon_type: None,
icon_value: None,
}
}
#[tokio::test]
async fn update_mcp_server_preserves_headers_when_field_is_omitted() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let created = create_mcp_server(
db,
remote_server_input(Some(r#"{"Authorization":"Bearer old"}"#.into())),
)
.await
.unwrap();
let updated = update_mcp_server(
db,
&created.id,
UpdateMcpServerInput {
name: Some("Renamed MCP".into()),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(updated.name, "Renamed MCP");
assert_eq!(
updated.headers_json.as_deref(),
Some(r#"{"Authorization":"Bearer old"}"#)
);
}
#[tokio::test]
async fn update_mcp_server_replaces_headers_when_value_is_provided() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let created = create_mcp_server(
db,
remote_server_input(Some(r#"{"Authorization":"Bearer old"}"#.into())),
)
.await
.unwrap();
let updated = update_mcp_server(
db,
&created.id,
UpdateMcpServerInput {
headers_json: Some(Some(r#"{"Authorization":"Bearer new"}"#.into())),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(
updated.headers_json.as_deref(),
Some(r#"{"Authorization":"Bearer new"}"#)
);
}
#[tokio::test]
async fn update_mcp_server_clears_headers_when_null_is_provided() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let created = create_mcp_server(
db,
remote_server_input(Some(r#"{"Authorization":"Bearer old"}"#.into())),
)
.await
.unwrap();
update_mcp_server(
db,
&created.id,
UpdateMcpServerInput {
headers_json: Some(None),
..Default::default()
},
)
.await
.unwrap();
let fetched = get_mcp_server(db, &created.id).await.unwrap();
assert_eq!(fetched.headers_json, None);
}
}
@@ -0,0 +1,229 @@
use sea_orm::sea_query::Expr;
use sea_orm::*;
use crate::entity::{memory_items, memory_namespaces};
use crate::error::{BoosApiError, Result};
use crate::types::{
CreateMemoryItemInput, CreateMemoryNamespaceInput, MemoryItem, MemoryNamespace,
UpdateMemoryItemInput, UpdateMemoryNamespaceInput,
};
use crate::utils::gen_id;
fn model_to_namespace(m: memory_namespaces::Model) -> MemoryNamespace {
MemoryNamespace {
id: m.id,
name: m.name,
scope: m.scope,
embedding_provider: m.embedding_provider,
embedding_dimensions: m.embedding_dimensions,
retrieval_threshold: m.retrieval_threshold,
retrieval_top_k: m.retrieval_top_k,
icon_type: m.icon_type,
icon_value: m.icon_value,
sort_order: m.sort_order,
}
}
fn model_to_item(m: memory_items::Model) -> MemoryItem {
MemoryItem {
id: m.id,
namespace_id: m.namespace_id,
title: m.title,
content: m.content,
source: m.source,
index_status: m.index_status,
index_error: m.index_error,
updated_at: m.updated_at,
}
}
pub async fn list_namespaces(db: &DatabaseConnection) -> Result<Vec<MemoryNamespace>> {
let models = memory_namespaces::Entity::find()
.order_by_asc(memory_namespaces::Column::SortOrder)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_namespace).collect())
}
pub async fn get_namespace(db: &DatabaseConnection, id: &str) -> Result<MemoryNamespace> {
let model = memory_namespaces::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryNamespace {}", id)))?;
Ok(model_to_namespace(model))
}
pub async fn create_namespace(
db: &DatabaseConnection,
input: CreateMemoryNamespaceInput,
) -> Result<MemoryNamespace> {
let id = gen_id();
let am = memory_namespaces::ActiveModel {
id: Set(id.clone()),
name: Set(input.name),
scope: Set(input.scope),
embedding_provider: Set(input.embedding_provider),
embedding_dimensions: Set(input.embedding_dimensions),
retrieval_threshold: Set(input.retrieval_threshold),
retrieval_top_k: Set(input.retrieval_top_k),
icon_type: Set(input.icon_type),
icon_value: Set(input.icon_value),
sort_order: Set(0),
};
am.insert(db).await?;
get_namespace(db, &id).await
}
pub async fn delete_namespace(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = memory_namespaces::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("MemoryNamespace {}", id)));
}
Ok(())
}
pub async fn update_namespace(
db: &DatabaseConnection,
id: &str,
input: UpdateMemoryNamespaceInput,
) -> Result<MemoryNamespace> {
let model = memory_namespaces::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryNamespace {}", id)))?;
let mut am: memory_namespaces::ActiveModel = model.clone().into();
if let Some(name) = input.name {
am.name = Set(name);
}
if input.update_embedding_provider {
am.embedding_provider = Set(input.embedding_provider);
}
if input.update_embedding_dimensions {
am.embedding_dimensions = Set(input.embedding_dimensions);
}
if input.update_retrieval_threshold {
am.retrieval_threshold = Set(input.retrieval_threshold);
}
if input.update_retrieval_top_k {
am.retrieval_top_k = Set(input.retrieval_top_k);
}
if input.update_icon {
am.icon_type = Set(input.icon_type);
am.icon_value = Set(input.icon_value);
}
if let Some(sort_order) = input.sort_order {
am.sort_order = Set(sort_order);
}
am.update(db).await?;
get_namespace(db, id).await
}
pub async fn reorder_namespaces(db: &DatabaseConnection, namespace_ids: &[String]) -> Result<()> {
for (i, id) in namespace_ids.iter().enumerate() {
memory_namespaces::Entity::update_many()
.col_expr(memory_namespaces::Column::SortOrder, Expr::value(i as i32))
.filter(memory_namespaces::Column::Id.eq(id))
.exec(db)
.await?;
}
Ok(())
}
pub async fn list_items(db: &DatabaseConnection, namespace_id: &str) -> Result<Vec<MemoryItem>> {
let models = memory_items::Entity::find()
.filter(memory_items::Column::NamespaceId.eq(namespace_id))
.order_by_desc(memory_items::Column::UpdatedAt)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_item).collect())
}
pub async fn add_item(db: &DatabaseConnection, input: CreateMemoryItemInput) -> Result<MemoryItem> {
let id = gen_id();
let source = input.source.unwrap_or_else(|| "manual".to_string());
let am = memory_items::ActiveModel {
id: Set(id.clone()),
namespace_id: Set(input.namespace_id),
title: Set(input.title),
content: Set(input.content),
source: Set(source),
..Default::default()
};
am.insert(db).await?;
let model = memory_items::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryItem {}", id)))?;
Ok(model_to_item(model))
}
pub async fn delete_item(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = memory_items::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("MemoryItem {}", id)));
}
Ok(())
}
pub async fn update_item(
db: &DatabaseConnection,
id: &str,
input: UpdateMemoryItemInput,
) -> Result<MemoryItem> {
let model = memory_items::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryItem {}", id)))?;
let mut am: memory_items::ActiveModel = model.into();
if let Some(title) = input.title {
am.title = Set(title);
}
if let Some(content) = input.content {
am.content = Set(content);
// Content changed — reset index status to pending
am.index_status = Set("pending".to_string());
}
am.updated_at = Set(chrono::Utc::now().to_rfc3339());
am.update(db).await?;
let updated = memory_items::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryItem {}", id)))?;
Ok(model_to_item(updated))
}
pub async fn update_item_index_status(
db: &DatabaseConnection,
id: &str,
status: &str,
error: Option<&str>,
) -> Result<()> {
let model = memory_items::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("MemoryItem {}", id)))?;
let mut am: memory_items::ActiveModel = model.into();
am.index_status = Set(status.to_string());
am.index_error = Set(error.map(|e| e.to_string()));
am.update(db).await?;
Ok(())
}
@@ -0,0 +1,855 @@
use sea_orm::sea_query::Expr;
use sea_orm::*;
use std::collections::HashSet;
use crate::entity::messages;
use crate::error::{BoosApiError, Result};
use crate::types::{Attachment, ConversationStats, Message, MessagePage, MessageRole};
use crate::utils::{gen_id, now_ts};
fn parse_role(s: &str) -> MessageRole {
match s {
"system" => MessageRole::System,
"user" => MessageRole::User,
"tool" => MessageRole::Tool,
_ => MessageRole::Assistant,
}
}
fn role_str(role: &MessageRole) -> &'static str {
match role {
MessageRole::System => "system",
MessageRole::User => "user",
MessageRole::Assistant => "assistant",
MessageRole::Tool => "tool",
}
}
fn parse_attachment_list(raw: &str) -> Result<Vec<Attachment>> {
serde_json::from_str(raw)
.map_err(|e| BoosApiError::Validation(format!("Invalid message attachments JSON: {e}")))
}
fn stringify_attachment_list(attachments: &[Attachment]) -> Result<String> {
serde_json::to_string(attachments).map_err(|e| {
BoosApiError::Validation(format!("Failed to serialize message attachments: {e}"))
})
}
const STALE_PARTIAL_ASSISTANT_ERROR: &str = "BoosApi was closed while this response was running. This stale response has been marked as failed.";
fn message_from_entity(m: messages::Model) -> Result<Message> {
Ok(Message {
id: m.id,
conversation_id: m.conversation_id,
role: parse_role(&m.role),
content: m.content,
provider_id: m.provider_id,
model_id: m.model_id,
token_count: m.token_count.map(|v| v as u32),
prompt_tokens: m.prompt_tokens.map(|v| v as u32),
completion_tokens: m.completion_tokens.map(|v| v as u32),
attachments: parse_attachment_list(&m.attachments)?,
thinking: m.thinking,
created_at: m.created_at,
parent_message_id: m.parent_message_id,
version_index: m.version_index,
is_active: m.is_active == 1,
tool_calls_json: m.tool_calls_json,
tool_call_id: m.tool_call_id,
status: m.status,
tokens_per_second: m.tokens_per_second,
first_token_latency_ms: m.first_token_latency_ms,
})
}
pub async fn list_messages(db: &DatabaseConnection, conversation_id: &str) -> Result<Vec<Message>> {
let rows = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::IsActive.eq(1))
.order_by_asc(messages::Column::CreatedAt)
.all(db)
.await?;
rows.into_iter().map(message_from_entity).collect()
}
pub async fn list_messages_for_model_context(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<Message>> {
let rows = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(
Condition::any()
.add(messages::Column::IsActive.eq(1))
.add(
Condition::all()
.add(messages::Column::VersionIndex.eq(-1))
.add(messages::Column::Role.is_in(["assistant", "tool"])),
),
)
.order_by_asc(messages::Column::CreatedAt)
.order_by_asc(messages::Column::Id)
.all(db)
.await?;
rows.into_iter().map(message_from_entity).collect()
}
pub async fn mark_stale_partial_assistant_messages_failed(db: &DatabaseConnection) -> Result<u64> {
let rows = messages::Entity::find()
.filter(messages::Column::Role.eq("assistant"))
.filter(messages::Column::Status.eq("partial"))
.filter(messages::Column::IsActive.eq(1))
.all(db)
.await?;
let mut changed = 0;
for row in rows {
let mut am: messages::ActiveModel = row.clone().into();
am.status = Set("error".to_string());
am.content = Set(if row.content.contains(STALE_PARTIAL_ASSISTANT_ERROR) {
row.content
} else {
format!(
"{}\n\n<!-- boosapi-stream-error -->\n{}",
row.content, STALE_PARTIAL_ASSISTANT_ERROR
)
});
am.update(db).await?;
changed += 1;
}
Ok(changed)
}
pub async fn list_messages_page(
db: &DatabaseConnection,
conversation_id: &str,
limit: u64,
before_message_id: Option<&str>,
) -> Result<MessagePage> {
let total_active_count = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::IsActive.eq(1))
.count(db)
.await?;
let mut query = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::IsActive.eq(1));
if let Some(cursor_id) = before_message_id {
let cursor = messages::Entity::find_by_id(cursor_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", cursor_id)))?;
query = query.filter(
Condition::any()
.add(messages::Column::CreatedAt.lt(cursor.created_at))
.add(
Condition::all()
.add(messages::Column::CreatedAt.eq(cursor.created_at))
.add(messages::Column::Id.lt(cursor.id.clone())),
),
);
}
let mut rows = query
.order_by_desc(messages::Column::CreatedAt)
.order_by_desc(messages::Column::Id)
.limit(limit + 1)
.all(db)
.await?;
let has_older = rows.len() > limit as usize;
if has_older {
rows.truncate(limit as usize);
}
rows.reverse();
let messages = rows
.into_iter()
.map(message_from_entity)
.collect::<Result<Vec<_>>>()?;
let oldest_message_id = messages.first().map(|message| message.id.clone());
Ok(MessagePage {
messages,
has_older,
oldest_message_id,
total_active_count,
})
}
pub async fn create_message(
db: &DatabaseConnection,
conversation_id: &str,
role: MessageRole,
content: &str,
attachments: &[Attachment],
parent_message_id: Option<&str>,
version_index: i32,
) -> Result<Message> {
let id = gen_id();
let now = now_ts();
let role_s = role_str(&role);
let attachments_json = stringify_attachment_list(attachments)?;
messages::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(conversation_id.to_string()),
role: Set(role_s.to_string()),
content: Set(content.to_string()),
attachments: Set(attachments_json),
created_at: Set(now),
parent_message_id: Set(parent_message_id.map(|s| s.to_string())),
version_index: Set(version_index),
is_active: Set(1),
..Default::default()
}
.insert(db)
.await?;
let row = messages::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", id)))?;
message_from_entity(row)
}
pub async fn update_message_content(
db: &DatabaseConnection,
id: &str,
content: &str,
) -> Result<Message> {
let row = messages::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", id)))?;
let mut am: messages::ActiveModel = row.into();
am.content = Set(content.to_string());
am.update(db).await?;
let row = messages::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", id)))?;
message_from_entity(row)
}
/// Update token usage stats on an existing message.
pub async fn update_message_usage(
db: &DatabaseConnection,
id: &str,
prompt_tokens: Option<i64>,
completion_tokens: Option<i64>,
) -> Result<()> {
let row = messages::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", id)))?;
let mut am: messages::ActiveModel = row.into();
if let Some(pt) = prompt_tokens {
am.prompt_tokens = Set(Some(pt));
}
if let Some(ct) = completion_tokens {
am.completion_tokens = Set(Some(ct));
}
am.update(db).await?;
Ok(())
}
fn compare_version_priority(left: &messages::Model, right: &messages::Model) -> std::cmp::Ordering {
right
.version_index
.cmp(&left.version_index)
.then_with(|| right.created_at.cmp(&left.created_at))
.then_with(|| right.id.cmp(&left.id))
}
fn select_next_active_version(
deleted: &messages::Model,
versions: &[messages::Model],
) -> Option<messages::Model> {
let remaining: Vec<messages::Model> = versions
.iter()
.filter(|version| version.id != deleted.id)
.cloned()
.collect();
if remaining.is_empty() {
return None;
}
let same_model: Vec<messages::Model> = if let Some(model_id) = deleted.model_id.as_ref() {
remaining
.iter()
.filter(|version| version.model_id.as_ref() == Some(model_id))
.cloned()
.collect()
} else {
Vec::new()
};
let mut candidates = if same_model.is_empty() { remaining } else { same_model };
candidates.sort_by(compare_version_priority);
candidates.into_iter().next()
}
pub async fn delete_message(db: &DatabaseConnection, id: &str) -> Result<()> {
let target = messages::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", id)))?;
let txn = db.begin().await?;
if target.role == "assistant" && target.is_active == 1 {
if let Some(parent_message_id) = target.parent_message_id.as_ref() {
let sibling_versions = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(&target.conversation_id))
.filter(messages::Column::ParentMessageId.eq(parent_message_id))
.filter(messages::Column::Role.eq("assistant"))
.filter(messages::Column::VersionIndex.gte(0))
.all(&txn)
.await?;
if let Some(next_version) = select_next_active_version(&target, &sibling_versions) {
let mut next_active: messages::ActiveModel = next_version.into();
next_active.is_active = Set(1);
next_active.update(&txn).await?;
}
}
}
let result = messages::Entity::delete_by_id(id).exec(&txn).await?;
txn.commit().await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("Message {}", id)));
}
Ok(())
}
/// Delete all messages in a conversation.
pub async fn clear_conversation_messages(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<u64> {
let result = messages::Entity::delete_many()
.filter(messages::Column::ConversationId.eq(conversation_id))
.exec(db)
.await?;
Ok(result.rows_affected)
}
/// Delete all messages in a conversation created after the given timestamp (inclusive).
pub async fn delete_messages_after(
db: &DatabaseConnection,
conversation_id: &str,
created_at: i64,
) -> Result<u64> {
let result = messages::Entity::delete_many()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::CreatedAt.gte(created_at))
.exec(db)
.await?;
Ok(result.rows_affected)
}
pub async fn list_message_versions(
db: &DatabaseConnection,
conversation_id: &str,
parent_message_id: &str,
) -> Result<Vec<Message>> {
let rows = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::ParentMessageId.eq(parent_message_id))
.filter(messages::Column::Role.eq("assistant"))
.filter(messages::Column::VersionIndex.gte(0))
.order_by_asc(messages::Column::VersionIndex)
.all(db)
.await?;
if rows.is_empty() {
return Ok(Vec::new());
}
let candidate_ids: Vec<String> = rows.iter().map(|row| row.id.clone()).collect();
let tool_rows = messages::Entity::find()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::Role.eq("tool"))
.filter(messages::Column::ParentMessageId.is_in(candidate_ids.clone()))
.all(db)
.await?;
let tool_parent_ids: HashSet<String> = tool_rows
.into_iter()
.filter_map(|row| row.parent_message_id)
.collect();
rows.into_iter()
.filter(|row| !tool_parent_ids.contains(&row.id))
.map(message_from_entity)
.collect()
}
pub async fn set_active_version(
db: &DatabaseConnection,
conversation_id: &str,
parent_message_id: &str,
target_message_id: &str,
) -> Result<()> {
// Deactivate all versions for this parent
messages::Entity::update_many()
.filter(messages::Column::ConversationId.eq(conversation_id))
.filter(messages::Column::ParentMessageId.eq(parent_message_id))
.col_expr(messages::Column::IsActive, Expr::value(0))
.exec(db)
.await?;
// Activate target version
let row = messages::Entity::find_by_id(target_message_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("Message {}", target_message_id)))?;
let mut am: messages::ActiveModel = row.into();
am.is_active = Set(1);
am.update(db).await?;
Ok(())
}
pub async fn delete_message_group(db: &DatabaseConnection, user_message_id: &str) -> Result<u64> {
// Delete all assistant versions for this user message
let ai_result = messages::Entity::delete_many()
.filter(messages::Column::ParentMessageId.eq(user_message_id))
.exec(db)
.await?;
// Delete the user message itself
messages::Entity::delete_by_id(user_message_id)
.exec(db)
.await?;
Ok(ai_result.rows_affected + 1)
}
/// Create a message with role "tool" for storing tool execution results.
pub async fn create_tool_result_message(
db: &DatabaseConnection,
conversation_id: &str,
tool_call_id: &str,
content: &str,
parent_message_id: &str,
) -> Result<()> {
let id = crate::utils::gen_id();
crate::entity::messages::ActiveModel {
id: Set(id),
conversation_id: Set(conversation_id.to_string()),
role: Set("tool".to_string()),
content: Set(content.to_string()),
provider_id: Set(None),
model_id: Set(None),
token_count: Set(None),
prompt_tokens: Set(None),
completion_tokens: Set(None),
attachments: Set("[]".to_string()),
thinking: Set(None),
created_at: Set(crate::utils::now_ts()),
branch_id: Set(None),
parent_message_id: Set(Some(parent_message_id.to_string())),
// Tool-result scaffolding: excluded from history reload (paired with intermediate assistant).
version_index: Set(-1),
is_active: Set(0),
tool_calls_json: Set(None),
tool_call_id: Set(Some(tool_call_id.to_string())),
status: Set("complete".to_string()),
tokens_per_second: Set(None),
first_token_latency_ms: Set(None),
}
.insert(db)
.await?;
Ok(())
}
/// Create an assistant message that contains tool_calls (intermediate message in tool loop).
pub async fn create_assistant_tool_call_message(
db: &DatabaseConnection,
conversation_id: &str,
content: &str,
tool_calls_json: Option<&str>,
provider_id: &str,
model_id: &str,
parent_message_id: &str,
) -> Result<String> {
let id = crate::utils::gen_id();
crate::entity::messages::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(conversation_id.to_string()),
role: Set("assistant".to_string()),
content: Set(content.to_string()),
provider_id: Set(Some(provider_id.to_string())),
model_id: Set(Some(model_id.to_string())),
token_count: Set(None),
prompt_tokens: Set(None),
completion_tokens: Set(None),
attachments: Set("[]".to_string()),
thinking: Set(None),
created_at: Set(crate::utils::now_ts()),
branch_id: Set(None),
parent_message_id: Set(Some(parent_message_id.to_string())),
// Intermediate tool-call scaffolding message. Excluded from user-visible AI version pagination.
version_index: Set(-1),
is_active: Set(0),
tool_calls_json: Set(tool_calls_json.map(|s| s.to_string())),
tool_call_id: Set(None),
status: Set("complete".to_string()),
tokens_per_second: Set(None),
first_token_latency_ms: Set(None),
}
.insert(db)
.await?;
Ok(id)
}
pub async fn get_conversation_stats(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<ConversationStats> {
use sea_orm::{FromQueryResult, Statement};
#[derive(Debug, FromQueryResult)]
struct StatsRow {
total_messages: i64,
total_user_messages: i64,
total_assistant_messages: i64,
total_prompt_tokens: i64,
total_completion_tokens: i64,
avg_tokens_per_second: Option<f64>,
avg_first_token_latency_ms: Option<f64>,
avg_response_time_ms: Option<f64>,
}
let sql = r#"
SELECT
COUNT(*) AS total_messages,
SUM(CASE WHEN role = 'user' THEN 1 ELSE 0 END) AS total_user_messages,
SUM(CASE WHEN role = 'assistant' THEN 1 ELSE 0 END) AS total_assistant_messages,
COALESCE(SUM(prompt_tokens), 0) AS total_prompt_tokens,
COALESCE(SUM(completion_tokens), 0) AS total_completion_tokens,
AVG(CASE WHEN tokens_per_second IS NOT NULL AND tokens_per_second > 0 THEN tokens_per_second ELSE NULL END) AS avg_tokens_per_second,
AVG(CASE WHEN first_token_latency_ms IS NOT NULL THEN first_token_latency_ms ELSE NULL END) AS avg_first_token_latency_ms,
AVG(CASE
WHEN first_token_latency_ms IS NOT NULL AND tokens_per_second IS NOT NULL AND tokens_per_second > 0 AND completion_tokens IS NOT NULL AND completion_tokens > 0
THEN first_token_latency_ms + (completion_tokens * 1000.0 / tokens_per_second)
ELSE NULL
END) AS avg_response_time_ms
FROM messages
WHERE conversation_id = ? AND is_active = 1
"#;
let row = StatsRow::find_by_statement(Statement::from_sql_and_values(
db.get_database_backend(),
sql,
vec![conversation_id.into()],
))
.one(db)
.await?
.unwrap_or(StatsRow {
total_messages: 0,
total_user_messages: 0,
total_assistant_messages: 0,
total_prompt_tokens: 0,
total_completion_tokens: 0,
avg_tokens_per_second: None,
avg_first_token_latency_ms: None,
avg_response_time_ms: None,
});
let total_prompt = row.total_prompt_tokens as u64;
let total_completion = row.total_completion_tokens as u64;
Ok(ConversationStats {
total_messages: row.total_messages as u64,
total_user_messages: row.total_user_messages as u64,
total_assistant_messages: row.total_assistant_messages as u64,
total_prompt_tokens: total_prompt,
total_completion_tokens: total_completion,
total_tokens: total_prompt + total_completion,
avg_tokens_per_second: row.avg_tokens_per_second,
avg_first_token_latency_ms: row.avg_first_token_latency_ms,
avg_response_time_ms: row.avg_response_time_ms,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::create_test_pool;
use crate::repo::conversation;
#[tokio::test]
async fn create_message_round_trips_attachment_metadata() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let conv = conversation::create_conversation(db, "Attach Chat", "m1", "p1", None)
.await
.unwrap();
let msg = create_message(
db,
&conv.id,
MessageRole::User,
"See attached",
&[Attachment {
id: "att-1".into(),
file_name: "image.png".into(),
file_type: "image/png".into(),
file_path: "conv-1/image.png".into(),
file_size: 3,
data: None,
}],
None,
0,
)
.await
.unwrap();
assert_eq!(msg.attachments.len(), 1);
assert_eq!(msg.attachments[0].file_name, "image.png");
assert_eq!(msg.attachments[0].file_type, "image/png");
assert_eq!(msg.attachments[0].file_path, "conv-1/image.png");
assert_eq!(msg.attachments[0].file_size, 3);
}
#[tokio::test]
async fn list_messages_for_model_context_includes_inactive_tool_scaffolding_only() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let conv = conversation::create_conversation(db, "Tool Chat", "model-1", "prov-1", None)
.await
.unwrap();
let user_msg = create_message(db, &conv.id, MessageRole::User, "Run tool", &[], None, 0)
.await
.unwrap();
let visible_reply = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Done",
&[],
Some(&user_msg.id),
0,
)
.await
.unwrap();
let scaffold_id = create_assistant_tool_call_message(
db,
&conv.id,
"",
Some(r#"[{"id":"call-1","type":"function","function":{"name":"read_file","arguments":"{}"}}]"#),
"prov-1",
"model-1",
&user_msg.id,
)
.await
.unwrap();
create_tool_result_message(db, &conv.id, "call-1", "tool result", &scaffold_id)
.await
.unwrap();
let stale_version = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Old inactive visible reply",
&[],
Some(&user_msg.id),
1,
)
.await
.unwrap();
let stale_row = messages::Entity::find_by_id(&stale_version.id)
.one(db)
.await
.unwrap()
.unwrap();
let mut stale_am: messages::ActiveModel = stale_row.into();
stale_am.is_active = Set(0);
stale_am.update(db).await.unwrap();
let visible_messages = list_messages(db, &conv.id).await.unwrap();
assert_eq!(visible_messages.len(), 2);
assert!(visible_messages.iter().any(|message| message.id == visible_reply.id));
let context_messages = list_messages_for_model_context(db, &conv.id).await.unwrap();
let ids = context_messages
.iter()
.map(|message| message.id.as_str())
.collect::<Vec<_>>();
assert!(ids.contains(&user_msg.id.as_str()));
assert!(ids.contains(&visible_reply.id.as_str()));
assert!(ids.contains(&scaffold_id.as_str()));
assert!(!ids.contains(&stale_version.id.as_str()));
assert_eq!(
context_messages
.iter()
.filter(|message| message.role == MessageRole::Tool)
.count(),
1
);
}
#[tokio::test]
async fn delete_active_assistant_version_promotes_remaining_version() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let conv = conversation::create_conversation(db, "Version Chat", "model-1", "prov-1", None)
.await
.unwrap();
let user_msg = create_message(db, &conv.id, MessageRole::User, "Hello!", &[], None, 0)
.await
.unwrap();
let version_a = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Reply A",
&[],
Some(&user_msg.id),
0,
)
.await
.unwrap();
let version_b = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Reply B",
&[],
Some(&user_msg.id),
1,
)
.await
.unwrap();
set_active_version(db, &conv.id, &user_msg.id, &version_a.id)
.await
.unwrap();
delete_message(db, &version_a.id).await.unwrap();
let visible_messages = list_messages(db, &conv.id).await.unwrap();
assert_eq!(visible_messages.len(), 2);
let active_reply = visible_messages
.iter()
.find(|message| message.role == MessageRole::Assistant)
.expect("assistant reply should remain visible");
assert_eq!(active_reply.id, version_b.id);
assert!(active_reply.is_active);
let versions = list_message_versions(db, &conv.id, &user_msg.id)
.await
.unwrap();
assert_eq!(versions.len(), 1);
assert_eq!(versions[0].id, version_b.id);
assert!(versions[0].is_active);
}
#[tokio::test]
async fn stale_partial_assistant_messages_are_marked_error_on_startup() {
let h = create_test_pool().await.unwrap();
let db = &h.conn;
let conv = conversation::create_conversation(db, "Crash Chat", "model-1", "prov-1", None)
.await
.unwrap();
let user_msg = create_message(db, &conv.id, MessageRole::User, "Call MCP", &[], None, 0)
.await
.unwrap();
let active_partial = create_message(
db,
&conv.id,
MessageRole::Assistant,
":::mcp {\"name\":\"server\",\"tool\":\"fetch\"}\n",
&[],
Some(&user_msg.id),
0,
)
.await
.unwrap();
let complete = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Done",
&[],
Some(&user_msg.id),
1,
)
.await
.unwrap();
let inactive_partial = create_message(
db,
&conv.id,
MessageRole::Assistant,
"Old partial",
&[],
Some(&user_msg.id),
2,
)
.await
.unwrap();
for (id, status, is_active) in [
(&active_partial.id, "partial", 1),
(&complete.id, "complete", 1),
(&inactive_partial.id, "partial", 0),
] {
let row = messages::Entity::find_by_id(id).one(db).await.unwrap().unwrap();
let mut am: messages::ActiveModel = row.into();
am.status = Set(status.to_string());
am.is_active = Set(is_active);
am.update(db).await.unwrap();
}
let changed = mark_stale_partial_assistant_messages_failed(db).await.unwrap();
assert_eq!(changed, 1);
let active = messages::Entity::find_by_id(&active_partial.id)
.one(db)
.await
.unwrap()
.unwrap();
assert_eq!(active.status, "error");
assert!(active.content.contains("BoosApi was closed while this response was running"));
let still_complete = messages::Entity::find_by_id(&complete.id)
.one(db)
.await
.unwrap()
.unwrap();
assert_eq!(still_complete.status, "complete");
assert_eq!(still_complete.content, "Done");
let still_inactive_partial = messages::Entity::find_by_id(&inactive_partial.id)
.one(db)
.await
.unwrap()
.unwrap();
assert_eq!(still_inactive_partial.status, "partial");
assert_eq!(still_inactive_partial.content, "Old partial");
}
}
@@ -0,0 +1,82 @@
use sea_orm::*;
use sea_query::OnConflict;
use crate::entity::program_policies;
use crate::error::{BoosApiError, Result};
use crate::types::{ProgramPolicy, SaveProgramPolicyInput};
use crate::utils::gen_id;
fn model_to_program_policy(m: program_policies::Model) -> ProgramPolicy {
ProgramPolicy {
id: m.id,
program_name: m.program_name,
allowed_provider_ids_json: m.allowed_provider_ids_json,
allowed_model_ids_json: m.allowed_model_ids_json,
default_provider_id: m.default_provider_id,
default_model_id: m.default_model_id,
rate_limit_per_minute: m.rate_limit_per_minute.map(|v| v as i32),
}
}
pub async fn list_program_policies(db: &DatabaseConnection) -> Result<Vec<ProgramPolicy>> {
let rows = program_policies::Entity::find()
.order_by_asc(program_policies::Column::ProgramName)
.all(db)
.await?;
Ok(rows.into_iter().map(model_to_program_policy).collect())
}
pub async fn save_program_policy(
db: &DatabaseConnection,
input: &SaveProgramPolicyInput,
) -> Result<ProgramPolicy> {
let id = gen_id();
let allowed_provider_ids_json =
serde_json::to_string(&input.allowed_provider_ids).unwrap_or_else(|_| "[]".to_string());
let allowed_model_ids_json =
serde_json::to_string(&input.allowed_model_ids).unwrap_or_else(|_| "[]".to_string());
program_policies::Entity::insert(program_policies::ActiveModel {
id: Set(id.clone()),
program_name: Set(input.program_name.clone()),
allowed_provider_ids_json: Set(allowed_provider_ids_json.clone()),
allowed_model_ids_json: Set(allowed_model_ids_json.clone()),
default_provider_id: Set(input.default_provider_id.clone()),
default_model_id: Set(input.default_model_id.clone()),
rate_limit_per_minute: Set(input.rate_limit_per_minute.map(|v| v as i64)),
})
.on_conflict(
OnConflict::column(program_policies::Column::ProgramName)
.update_columns([
program_policies::Column::AllowedProviderIdsJson,
program_policies::Column::AllowedModelIdsJson,
program_policies::Column::DefaultProviderId,
program_policies::Column::DefaultModelId,
program_policies::Column::RateLimitPerMinute,
])
.to_owned(),
)
.exec(db)
.await?;
Ok(ProgramPolicy {
id,
program_name: input.program_name.clone(),
allowed_provider_ids_json,
allowed_model_ids_json,
default_provider_id: input.default_provider_id.clone(),
default_model_id: input.default_model_id.clone(),
rate_limit_per_minute: input.rate_limit_per_minute,
})
}
pub async fn delete_program_policy(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = program_policies::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("ProgramPolicy {}", id)));
}
Ok(())
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,133 @@
use sea_orm::*;
use crate::entity::provider_keys;
use crate::error::{BoosApiError, Result};
use crate::types::ProviderKey;
use crate::utils::{gen_id, now_ts};
fn key_from_entity(m: provider_keys::Model) -> ProviderKey {
ProviderKey {
id: m.id,
provider_id: m.provider_id,
key_encrypted: m.key_encrypted,
key_prefix: m.key_prefix,
enabled: m.enabled != 0,
last_validated_at: m.last_validated_at,
last_error: m.last_error,
rotation_index: m.rotation_index as u32,
created_at: m.created_at,
}
}
pub async fn list_keys_for_provider(db: &DatabaseConnection, provider_id: &str) -> Result<Vec<ProviderKey>> {
let rows = provider_keys::Entity::find()
.filter(provider_keys::Column::ProviderId.eq(provider_id))
.order_by_asc(provider_keys::Column::RotationIndex)
.all(db)
.await?;
Ok(rows.into_iter().map(key_from_entity).collect())
}
pub async fn add_provider_key(
db: &DatabaseConnection,
provider_id: &str,
key_encrypted: &str,
key_prefix: &str,
) -> Result<ProviderKey> {
let id = gen_id();
let now = now_ts();
let max_idx = provider_keys::Entity::find()
.filter(provider_keys::Column::ProviderId.eq(provider_id))
.select_only()
.column_as(provider_keys::Column::RotationIndex.max(), "m")
.into_tuple::<Option<i32>>()
.one(db)
.await?
.flatten();
let rotation_index = max_idx.unwrap_or(-1) + 1;
provider_keys::ActiveModel {
id: Set(id.clone()),
provider_id: Set(provider_id.to_string()),
key_encrypted: Set(key_encrypted.to_string()),
key_prefix: Set(key_prefix.to_string()),
enabled: Set(1),
last_validated_at: Set(None),
last_error: Set(None),
rotation_index: Set(rotation_index),
created_at: Set(now),
}
.insert(db)
.await?;
let row = provider_keys::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ProviderKey {}", id)))?;
Ok(key_from_entity(row))
}
pub async fn delete_provider_key(db: &DatabaseConnection, key_id: &str) -> Result<()> {
let result = provider_keys::Entity::delete_by_id(key_id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("ProviderKey {}", key_id)));
}
Ok(())
}
/// Toggles enabled state (NOT enabled). Different from provider.rs which takes an explicit bool.
pub async fn toggle_provider_key(db: &DatabaseConnection, key_id: &str) -> Result<ProviderKey> {
let row = provider_keys::Entity::find_by_id(key_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ProviderKey {}", key_id)))?;
let new_enabled = if row.enabled != 0 { 0 } else { 1 };
let mut am: provider_keys::ActiveModel = row.into();
am.enabled = Set(new_enabled);
am.update(db).await?;
let row = provider_keys::Entity::find_by_id(key_id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ProviderKey {}", key_id)))?;
Ok(key_from_entity(row))
}
/// Takes an optional error message. Different from provider.rs which takes a bool.
pub async fn update_key_validation(
db: &DatabaseConnection,
key_id: &str,
error: Option<&str>,
) -> Result<()> {
if let Some(row) = provider_keys::Entity::find_by_id(key_id).one(db).await? {
let mut am: provider_keys::ActiveModel = row.into();
am.last_validated_at = Set(Some(now_ts()));
am.last_error = Set(error.map(|s| s.to_string()));
am.update(db).await?;
}
Ok(())
}
pub async fn get_enabled_keys(db: &DatabaseConnection, provider_id: &str) -> Result<Vec<ProviderKey>> {
let rows = provider_keys::Entity::find()
.filter(provider_keys::Column::ProviderId.eq(provider_id))
.filter(provider_keys::Column::Enabled.eq(1))
.order_by_asc(provider_keys::Column::RotationIndex)
.all(db)
.await?;
Ok(rows.into_iter().map(key_from_entity).collect())
}
pub async fn update_rotation_index(db: &DatabaseConnection, key_id: &str, index: u32) -> Result<()> {
if let Some(row) = provider_keys::Entity::find_by_id(key_id).one(db).await? {
let mut am: provider_keys::ActiveModel = row.into();
am.rotation_index = Set(index as i32);
am.update(db).await?;
}
Ok(())
}
@@ -0,0 +1,119 @@
use sea_orm::*;
use crate::entity::search_providers;
use crate::error::{BoosApiError, Result};
use crate::types::{CreateSearchProviderInput, SearchProvider};
use crate::utils::gen_id;
fn model_to_search_provider(m: search_providers::Model) -> SearchProvider {
SearchProvider {
id: m.id,
name: m.name,
provider_type: m.provider_type,
endpoint: m.endpoint,
has_api_key: m.api_key_ref.as_ref().map_or(false, |k| !k.is_empty()),
enabled: m.enabled != 0,
region: m.region,
language: m.language,
safe_search: m.safe_search.map(|v| v != 0),
result_limit: m.result_limit,
timeout_ms: m.timeout_ms,
}
}
pub async fn list_search_providers(db: &DatabaseConnection) -> Result<Vec<SearchProvider>> {
let rows = search_providers::Entity::find()
.order_by_asc(search_providers::Column::Name)
.all(db)
.await?;
Ok(rows.into_iter().map(model_to_search_provider).collect())
}
pub async fn get_search_provider(db: &DatabaseConnection, id: &str) -> Result<SearchProvider> {
let model = search_providers::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("SearchProvider {}", id)))?;
Ok(model_to_search_provider(model))
}
pub async fn create_search_provider(
db: &DatabaseConnection,
input: CreateSearchProviderInput,
) -> Result<SearchProvider> {
let id = gen_id();
search_providers::ActiveModel {
id: Set(id.clone()),
name: Set(input.name),
provider_type: Set(input.provider_type),
endpoint: Set(input.endpoint),
api_key_ref: Set(input.api_key),
enabled: Set(if input.enabled.unwrap_or(true) { 1 } else { 0 }),
region: Set(input.region),
language: Set(input.language),
safe_search: Set(input.safe_search.map(|b| if b { 1i64 } else { 0 })),
result_limit: Set(input.result_limit.unwrap_or(10)),
timeout_ms: Set(input.timeout_ms.unwrap_or(5000)),
}
.insert(db)
.await?;
get_search_provider(db, &id).await
}
pub async fn update_search_provider(
db: &DatabaseConnection,
id: &str,
input: CreateSearchProviderInput,
) -> Result<SearchProvider> {
let model = search_providers::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("SearchProvider {}", id)))?;
let name = if input.name.is_empty() {
model.name.clone()
} else {
input.name
};
let provider_type = if input.provider_type.is_empty() {
model.provider_type.clone()
} else {
input.provider_type
};
let endpoint = input.endpoint.or(model.endpoint.clone());
let api_key_ref = input.api_key.or(model.api_key_ref.clone());
let enabled = input.enabled.unwrap_or(model.enabled != 0);
let region = input.region.or(model.region.clone());
let language = input.language.or(model.language.clone());
let safe_search = input.safe_search.or(model.safe_search.map(|v| v != 0));
let result_limit = input.result_limit.unwrap_or(model.result_limit);
let timeout_ms = input.timeout_ms.unwrap_or(model.timeout_ms);
let mut am: search_providers::ActiveModel = model.into();
am.name = Set(name);
am.provider_type = Set(provider_type);
am.endpoint = Set(endpoint);
am.api_key_ref = Set(api_key_ref);
am.enabled = Set(if enabled { 1 } else { 0 });
am.region = Set(region);
am.language = Set(language);
am.safe_search = Set(safe_search.map(|b| if b { 1i64 } else { 0 }));
am.result_limit = Set(result_limit);
am.timeout_ms = Set(timeout_ms);
am.update(db).await?;
get_search_provider(db, id).await
}
pub async fn delete_search_provider(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = search_providers::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("SearchProvider {}", id)));
}
Ok(())
}
@@ -0,0 +1,76 @@
use sea_orm::*;
use sea_query::OnConflict;
use crate::entity::settings;
use crate::error::{BoosApiError, Result};
use crate::types::AppSettings;
pub async fn get_settings(db: &DatabaseConnection) -> Result<AppSettings> {
let rows = settings::Entity::find().all(db).await?;
let mut map = serde_json::Map::new();
for row in &rows {
let val = serde_json::from_str::<serde_json::Value>(&row.value)
.unwrap_or_else(|_| serde_json::Value::String(row.value.clone()));
map.insert(row.key.clone(), val);
}
let settings: AppSettings =
serde_json::from_value(serde_json::Value::Object(map)).unwrap_or_default();
Ok(settings)
}
pub async fn save_settings(db: &DatabaseConnection, settings: &AppSettings) -> Result<()> {
let value = serde_json::to_value(settings).unwrap_or_default();
if let serde_json::Value::Object(map) = value {
db.transaction::<_, _, sea_orm::DbErr>(|txn| {
Box::pin(async move {
for (key, val) in map {
let val_str = match &val {
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
};
settings::Entity::insert(settings::ActiveModel {
key: Set(key),
value: Set(val_str),
})
.on_conflict(
OnConflict::column(settings::Column::Key)
.update_column(settings::Column::Value)
.to_owned(),
)
.exec(txn)
.await?;
}
Ok(())
})
})
.await
.map_err(|e| match e {
sea_orm::TransactionError::Connection(db_err) => BoosApiError::from(db_err),
sea_orm::TransactionError::Transaction(db_err) => BoosApiError::from(db_err),
})?;
}
Ok(())
}
pub async fn get_setting(db: &DatabaseConnection, key: &str) -> Result<Option<String>> {
let row = settings::Entity::find_by_id(key).one(db).await?;
Ok(row.map(|r| r.value))
}
pub async fn set_setting(db: &DatabaseConnection, key: &str, value: &str) -> Result<()> {
settings::Entity::insert(settings::ActiveModel {
key: Set(key.to_string()),
value: Set(value.to_string()),
})
.on_conflict(
OnConflict::column(settings::Column::Key)
.update_column(settings::Column::Value)
.to_owned(),
)
.exec(db)
.await?;
Ok(())
}
@@ -0,0 +1,46 @@
use sea_orm::*;
use sea_orm::prelude::Expr;
use std::collections::HashSet;
use crate::entity::skill_states;
use crate::error::Result;
use crate::utils::now_ts;
/// Get all disabled skill names from the database.
pub async fn get_disabled_skills(db: &DatabaseConnection) -> Result<HashSet<String>> {
let rows = skill_states::Entity::find()
.filter(skill_states::Column::Enabled.eq(0))
.all(db)
.await?;
Ok(rows.into_iter().map(|r| r.name).collect())
}
/// Set a skill's enabled state. Creates or updates the record.
pub async fn set_skill_enabled(
db: &DatabaseConnection,
name: &str,
enabled: bool,
) -> Result<()> {
let now = now_ts();
let existing = skill_states::Entity::find_by_id(name).one(db).await?;
if let Some(_) = existing {
skill_states::Entity::update_many()
.col_expr(
skill_states::Column::Enabled,
Expr::value(if enabled { 1 } else { 0 }),
)
.col_expr(skill_states::Column::UpdatedAt, Expr::value(now))
.filter(skill_states::Column::Name.eq(name))
.exec(db)
.await?;
} else {
let am = skill_states::ActiveModel {
name: Set(name.to_string()),
enabled: Set(if enabled { 1 } else { 0 }),
updated_at: Set(now),
};
am.insert(db).await?;
}
Ok(())
}
@@ -0,0 +1,127 @@
use sea_orm::*;
use serde::{Deserialize, Serialize};
use crate::entity::stored_files;
use crate::error::{BoosApiError, Result};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredFile {
pub id: String,
pub hash: String,
pub original_name: String,
pub mime_type: String,
pub size_bytes: i64,
pub storage_path: String,
pub conversation_id: Option<String>,
pub created_at: String,
}
fn model_to_stored_file(m: stored_files::Model) -> StoredFile {
StoredFile {
id: m.id,
hash: m.hash,
original_name: m.original_name,
mime_type: m.mime_type,
size_bytes: m.size_bytes,
storage_path: m.storage_path,
conversation_id: m.conversation_id,
created_at: m.created_at,
}
}
pub async fn create_stored_file(
db: &DatabaseConnection,
id: &str,
hash: &str,
original_name: &str,
mime_type: &str,
size_bytes: i64,
storage_path: &str,
conversation_id: Option<&str>,
) -> Result<StoredFile> {
let am = stored_files::ActiveModel {
id: Set(id.to_string()),
hash: Set(hash.to_string()),
original_name: Set(original_name.to_string()),
mime_type: Set(mime_type.to_string()),
size_bytes: Set(size_bytes),
storage_path: Set(storage_path.to_string()),
conversation_id: Set(conversation_id.map(|s| s.to_string())),
..Default::default()
};
am.insert(db).await?;
get_stored_file(db, id).await
}
pub async fn get_stored_file(db: &DatabaseConnection, id: &str) -> Result<StoredFile> {
let model = stored_files::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("StoredFile {}", id)))?;
Ok(model_to_stored_file(model))
}
pub async fn list_stored_files_by_conversation(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<StoredFile>> {
let models = stored_files::Entity::find()
.filter(stored_files::Column::ConversationId.eq(conversation_id))
.order_by_desc(stored_files::Column::CreatedAt)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_stored_file).collect())
}
pub async fn delete_stored_file(db: &DatabaseConnection, id: &str) -> Result<()> {
let result = stored_files::Entity::delete_by_id(id).exec(db).await?;
if result.rows_affected == 0 {
return Err(BoosApiError::NotFound(format!("StoredFile {}", id)));
}
Ok(())
}
pub async fn delete_stored_files_by_conversation(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<()> {
stored_files::Entity::delete_many()
.filter(stored_files::Column::ConversationId.eq(conversation_id))
.exec(db)
.await?;
Ok(())
}
pub async fn list_all_stored_files(db: &DatabaseConnection) -> Result<Vec<StoredFile>> {
let models = stored_files::Entity::find()
.order_by_desc(stored_files::Column::CreatedAt)
.all(db)
.await?;
Ok(models.into_iter().map(model_to_stored_file).collect())
}
pub async fn count_stored_files_with_storage_path(
db: &DatabaseConnection,
storage_path: &str,
) -> Result<u64> {
stored_files::Entity::find()
.filter(stored_files::Column::StoragePath.eq(storage_path))
.count(db)
.await
.map_err(Into::into)
}
pub async fn find_by_hash(db: &DatabaseConnection, hash: &str) -> Result<Option<StoredFile>> {
let model = stored_files::Entity::find()
.filter(stored_files::Column::Hash.eq(hash))
.one(db)
.await?;
Ok(model.map(model_to_stored_file))
}
@@ -0,0 +1,111 @@
use sea_orm::*;
use crate::entity::tool_executions;
use crate::error::{BoosApiError, Result};
use crate::types::ToolExecution;
use crate::utils::gen_id;
fn model_to_tool_execution(m: tool_executions::Model) -> ToolExecution {
ToolExecution {
id: m.id,
conversation_id: m.conversation_id,
message_id: m.message_id,
server_id: m.server_id,
tool_name: m.tool_name,
status: m.status,
input_preview: m.input_preview,
output_preview: m.output_preview,
error_message: m.error_message,
duration_ms: m.duration_ms,
created_at: m.created_at,
approval_status: m.approval_status,
}
}
pub async fn list_tool_executions(
db: &DatabaseConnection,
conversation_id: &str,
) -> Result<Vec<ToolExecution>> {
let rows = tool_executions::Entity::find()
.filter(tool_executions::Column::ConversationId.eq(conversation_id))
.order_by_desc(tool_executions::Column::CreatedAt)
.all(db)
.await?;
Ok(rows.into_iter().map(model_to_tool_execution).collect())
}
pub async fn create_tool_execution(
db: &DatabaseConnection,
conversation_id: &str,
message_id: Option<&str>,
server_id: &str,
tool_name: &str,
input_preview: Option<&str>,
approval_status: Option<&str>,
) -> Result<ToolExecution> {
let id = gen_id();
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
tool_executions::ActiveModel {
id: Set(id.clone()),
conversation_id: Set(conversation_id.to_string()),
message_id: Set(message_id.map(|s| s.to_string())),
server_id: Set(server_id.to_string()),
tool_name: Set(tool_name.to_string()),
status: Set("pending".to_string()),
input_preview: Set(input_preview.map(|s| s.to_string())),
output_preview: Set(None),
error_message: Set(None),
duration_ms: Set(None),
created_at: Set(now),
approval_status: Set(approval_status.map(|s| s.to_string())),
}
.insert(db)
.await?;
let model = tool_executions::Entity::find_by_id(&id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ToolExecution {}", id)))?;
Ok(model_to_tool_execution(model))
}
pub async fn update_tool_execution_status(
db: &DatabaseConnection,
id: &str,
status: &str,
output: Option<&str>,
error: Option<&str>,
) -> Result<()> {
let model = tool_executions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ToolExecution {}", id)))?;
let mut am: tool_executions::ActiveModel = model.into();
am.status = Set(status.to_string());
am.output_preview = Set(output.map(|s| s.to_string()));
am.error_message = Set(error.map(|s| s.to_string()));
am.update(db).await?;
Ok(())
}
pub async fn update_tool_execution_approval_status(
db: &DatabaseConnection,
id: &str,
approval_status: &str,
) -> Result<()> {
let model = tool_executions::Entity::find_by_id(id)
.one(db)
.await?
.ok_or_else(|| BoosApiError::NotFound(format!("ToolExecution {}", id)))?;
let mut am: tool_executions::ActiveModel = model.into();
am.approval_status = Set(Some(approval_status.to_string()));
am.update(db).await?;
Ok(())
}
@@ -0,0 +1,540 @@
use aws_config::BehaviorVersion;
use aws_credential_types::{provider::SharedCredentialsProvider, Credentials};
use aws_sdk_s3::config::retry::RetryConfig;
use aws_sdk_s3::config::timeout::TimeoutConfig;
use aws_sdk_s3::config::Region;
use aws_sdk_s3::error::{ProvideErrorMetadata, SdkError};
use aws_sdk_s3::primitives::{ByteStream, DateTime, DateTimeFormat};
use aws_sdk_s3::Client;
use serde::{Deserialize, Serialize};
use std::error::Error as StdError;
use std::fmt::Debug;
use std::path::Path;
use std::time::Duration;
use crate::error::{BoosApiError, Result};
use crate::webdav::parse_hostname_from_filename;
const S3_CONNECT_TIMEOUT_SECS: u64 = 10;
const S3_READ_TIMEOUT_SECS: u64 = 30;
const S3_OPERATION_ATTEMPT_TIMEOUT_SECS: u64 = 300;
const S3_OPERATION_TIMEOUT_SECS: u64 = 900;
const S3_MAX_ATTEMPTS: u32 = 5;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct S3Config {
pub bucket: String,
pub region: String,
pub prefix: String,
pub endpoint_url: Option<String>,
pub force_path_style: bool,
pub use_default_credentials: bool,
pub access_key_id: String,
pub secret_access_key: String,
pub session_token: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct S3FileInfo {
pub file_name: String,
pub size: i64,
pub last_modified: String,
pub hostname: String,
}
pub struct S3BackupClient {
client: Client,
config: S3Config,
}
impl S3BackupClient {
pub async fn new(config: S3Config) -> Result<Self> {
validate_config(&config)?;
let mut loader = aws_config::defaults(BehaviorVersion::latest())
.region(Region::new(config.region.trim().to_string()));
if !config.use_default_credentials {
let session_token = config
.session_token
.as_deref()
.map(str::trim)
.filter(|v| !v.is_empty())
.map(ToOwned::to_owned);
let credentials = Credentials::new(
config.access_key_id.trim(),
config.secret_access_key.trim(),
session_token,
None,
"boosapi-s3",
);
loader = loader.credentials_provider(SharedCredentialsProvider::new(credentials));
}
let shared_config = loader.load().await;
let mut s3_config = aws_sdk_s3::config::Builder::from(&shared_config)
.force_path_style(config.force_path_style)
.timeout_config(s3_timeout_config())
.retry_config(s3_retry_config());
if let Some(endpoint) = normalized_endpoint(&config) {
s3_config = s3_config.endpoint_url(endpoint);
}
Ok(Self {
client: Client::from_conf(s3_config.build()),
config,
})
}
pub async fn check_connection(&self) -> Result<bool> {
self.client
.head_bucket()
.bucket(bucket_name(&self.config))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format_s3_sdk_error("connection", &e)))?;
Ok(true)
}
pub async fn list_files(&self) -> Result<Vec<S3FileInfo>> {
let mut files = Vec::new();
let mut continuation_token = None;
let prefix = normalize_prefix(&self.config.prefix);
loop {
let mut request = self
.client
.list_objects_v2()
.bucket(bucket_name(&self.config))
.prefix(&prefix);
if let Some(token) = continuation_token {
request = request.continuation_token(token);
}
let response = request
.send()
.await
.map_err(|e| BoosApiError::Gateway(format_s3_sdk_error("list backups", &e)))?;
for object in response.contents() {
let Some(key) = object.key() else {
continue;
};
let Some(file_name) = backup_file_name_from_key(&prefix, key) else {
continue;
};
files.push(S3FileInfo {
file_name: file_name.clone(),
size: object.size().unwrap_or(0),
last_modified: object
.last_modified()
.map(format_last_modified)
.unwrap_or_default(),
hostname: parse_hostname_from_filename(&file_name),
});
}
if !response.is_truncated().unwrap_or(false) {
break;
}
continuation_token = response
.next_continuation_token()
.map(ToOwned::to_owned);
if continuation_token.is_none() {
break;
}
}
sort_backup_files(&mut files);
Ok(files)
}
pub async fn upload_file(&self, filename: &str, local_path: &Path) -> Result<()> {
let body = ByteStream::from_path(local_path)
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read file for S3 upload: {}", e)))?;
self.client
.put_object()
.bucket(bucket_name(&self.config))
.key(object_key(&self.config.prefix, filename))
.body(body)
.send()
.await
.map_err(|e| BoosApiError::Gateway(format_s3_sdk_error("upload", &e)))?;
Ok(())
}
pub async fn download_file(&self, filename: &str, local_path: &Path) -> Result<()> {
let response = self
.client
.get_object()
.bucket(bucket_name(&self.config))
.key(object_key(&self.config.prefix, filename))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format_s3_sdk_error("download", &e)))?;
let data = response
.body
.collect()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read S3 download: {}", e)))?
.into_bytes();
std::fs::write(local_path, &data)
.map_err(|e| BoosApiError::Gateway(format!("Failed to write download: {}", e)))?;
Ok(())
}
pub async fn delete_file(&self, filename: &str) -> Result<()> {
self.client
.delete_object()
.bucket(bucket_name(&self.config))
.key(object_key(&self.config.prefix, filename))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format_s3_sdk_error("delete", &e)))?;
Ok(())
}
}
fn s3_timeout_config() -> TimeoutConfig {
TimeoutConfig::builder()
.connect_timeout(Duration::from_secs(S3_CONNECT_TIMEOUT_SECS))
.read_timeout(Duration::from_secs(S3_READ_TIMEOUT_SECS))
.operation_attempt_timeout(Duration::from_secs(S3_OPERATION_ATTEMPT_TIMEOUT_SECS))
.operation_timeout(Duration::from_secs(S3_OPERATION_TIMEOUT_SECS))
.build()
}
fn s3_retry_config() -> RetryConfig {
RetryConfig::standard().with_max_attempts(S3_MAX_ATTEMPTS)
}
fn format_s3_sdk_error<E, R>(action: &str, err: &SdkError<E, R>) -> String
where
E: ProvideErrorMetadata + StdError + 'static,
R: Debug,
{
let prefix = format!("S3 {} failed", action);
match err {
SdkError::TimeoutError(_) => {
format!(
"{}: request timed out. Details: {}",
prefix,
sdk_error_chain(err)
)
}
SdkError::DispatchFailure(dispatch) => {
let kind = if dispatch.is_timeout() {
"timeout"
} else if dispatch.is_io() {
"I/O"
} else if dispatch.is_user() {
"request"
} else {
"transport"
};
let details = dispatch
.as_connector_error()
.map(|source| error_chain(source))
.unwrap_or_else(|| sdk_error_chain(err));
format!(
"{}: network request failed ({}). Check S3 endpoint, proxy, DNS, and network connectivity. Details: {}",
prefix, kind, details
)
}
SdkError::ServiceError(_) => {
let details = err
.as_service_error()
.and_then(service_error_details)
.unwrap_or_else(|| sdk_error_chain(err));
format!("{}: {}", prefix, details)
}
_ => format!("{}: {}", prefix, sdk_error_chain(err)),
}
}
fn service_error_details(error: &(impl ProvideErrorMetadata + StdError)) -> Option<String> {
match (error.code(), error.message()) {
(Some(code), Some(message)) if !message.trim().is_empty() => {
Some(format!("{}: {}", code, message))
}
(Some(code), _) => Some(code.to_string()),
(_, Some(message)) if !message.trim().is_empty() => Some(message.to_string()),
_ => {
let fallback = error.to_string();
(!fallback.trim().is_empty()).then_some(fallback)
}
}
}
fn sdk_error_chain<E, R>(err: &SdkError<E, R>) -> String
where
E: StdError + 'static,
R: Debug,
{
StdError::source(err)
.map(error_chain)
.unwrap_or_else(|| err.to_string())
}
fn error_chain(mut error: &(dyn StdError + 'static)) -> String {
let mut parts = vec![error.to_string()];
while let Some(source) = error.source() {
parts.push(source.to_string());
error = source;
}
parts.join(": ")
}
fn validate_config(config: &S3Config) -> Result<()> {
if config.bucket.trim().is_empty() {
return Err(BoosApiError::Gateway("S3 bucket is required".into()));
}
if config.region.trim().is_empty() {
return Err(BoosApiError::Gateway("S3 region is required".into()));
}
if !config.use_default_credentials {
if config.access_key_id.trim().is_empty() {
return Err(BoosApiError::Gateway("S3 access key ID is required".into()));
}
if config.secret_access_key.trim().is_empty() {
return Err(BoosApiError::Gateway(
"S3 secret access key is required".into(),
));
}
}
Ok(())
}
fn bucket_name(config: &S3Config) -> &str {
config.bucket.trim()
}
fn normalized_endpoint(config: &S3Config) -> Option<String> {
config
.endpoint_url
.as_deref()
.map(str::trim)
.filter(|v| !v.is_empty())
.map(ToOwned::to_owned)
}
pub(crate) fn normalize_prefix(prefix: &str) -> String {
let trimmed = prefix.trim().trim_matches('/');
if trimmed.is_empty() {
String::new()
} else {
format!("{}/", trimmed)
}
}
pub(crate) fn object_key(prefix: &str, filename: &str) -> String {
format!("{}{}", normalize_prefix(prefix), filename.trim_start_matches('/'))
}
pub(crate) fn backup_file_name_from_key(prefix: &str, key: &str) -> Option<String> {
let name = key.strip_prefix(prefix).unwrap_or(key);
if name.contains('/') || !name.starts_with("boosapi-backup-") || !name.ends_with(".zip") {
return None;
}
Some(name.to_string())
}
pub(crate) fn sort_backup_files(files: &mut [S3FileInfo]) {
files.sort_by(|a, b| b.file_name.cmp(&a.file_name));
}
fn format_last_modified(value: &DateTime) -> String {
value
.fmt(DateTimeFormat::DateTime)
.unwrap_or_else(|_| value.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn normalizes_s3_prefix_for_object_keys() {
assert_eq!(normalize_prefix(""), "");
assert_eq!(normalize_prefix("/boosapi/backups/"), "boosapi/backups/");
assert_eq!(object_key("/boosapi/backups/", "/backup.zip"), "boosapi/backups/backup.zip");
}
#[test]
fn filters_backup_files_from_prefixed_keys() {
let prefix = normalize_prefix("boosapi/backups");
assert_eq!(
backup_file_name_from_key(&prefix, "boosapi/backups/boosapi-backup-20260515_010203.host.zip"),
Some("boosapi-backup-20260515_010203.host.zip".to_string())
);
assert_eq!(
backup_file_name_from_key(&prefix, "boosapi/backups/nested/boosapi-backup-20260515_010203.host.zip"),
None
);
assert_eq!(
backup_file_name_from_key(&prefix, "boosapi/backups/not-a-backup.zip"),
None
);
assert_eq!(
backup_file_name_from_key(&prefix, "boosapi/backups/boosapi-backup-20260515_010203.host.db"),
None
);
}
#[test]
fn sorts_s3_backups_newest_first_by_filename() {
let mut files = vec![
S3FileInfo {
file_name: "boosapi-backup-20260514_010203.alpha.zip".to_string(),
size: 1,
last_modified: String::new(),
hostname: "alpha".to_string(),
},
S3FileInfo {
file_name: "boosapi-backup-20260515_010203.alpha.zip".to_string(),
size: 1,
last_modified: String::new(),
hostname: "alpha".to_string(),
},
];
sort_backup_files(&mut files);
assert_eq!(files[0].file_name, "boosapi-backup-20260515_010203.alpha.zip");
assert_eq!(files[1].file_name, "boosapi-backup-20260514_010203.alpha.zip");
}
#[test]
fn s3_client_uses_bounded_timeouts_and_extra_retries() {
let timeout = s3_timeout_config();
assert_eq!(
timeout.connect_timeout(),
Some(std::time::Duration::from_secs(S3_CONNECT_TIMEOUT_SECS))
);
assert_eq!(
timeout.read_timeout(),
Some(std::time::Duration::from_secs(S3_READ_TIMEOUT_SECS))
);
assert_eq!(
timeout.operation_attempt_timeout(),
Some(std::time::Duration::from_secs(
S3_OPERATION_ATTEMPT_TIMEOUT_SECS
))
);
assert_eq!(
timeout.operation_timeout(),
Some(std::time::Duration::from_secs(S3_OPERATION_TIMEOUT_SECS))
);
assert_eq!(s3_retry_config().max_attempts(), S3_MAX_ATTEMPTS);
}
#[test]
fn s3_dispatch_errors_are_actionable_for_users() {
let err = aws_sdk_s3::error::SdkError::<TestServiceError, ()>::dispatch_failure(
aws_sdk_s3::error::ConnectorError::io(Box::new(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"connection reset by peer",
))),
);
let message = format_s3_sdk_error("upload", &err);
assert!(message.contains("S3 upload failed: network request failed"));
assert!(message.contains("endpoint"));
assert!(message.contains("connection reset by peer"));
assert_ne!(message, "S3 upload failed: dispatch failure");
}
#[test]
fn s3_timeout_errors_are_actionable_for_users() {
let err = aws_sdk_s3::error::SdkError::<TestServiceError, ()>::timeout_error(
std::io::Error::new(std::io::ErrorKind::TimedOut, "deadline elapsed"),
);
let message = format_s3_sdk_error("upload", &err);
assert!(message.contains("S3 upload failed: request timed out"));
assert!(message.contains("deadline elapsed"));
}
#[test]
fn s3_service_errors_keep_service_details() {
let err = aws_sdk_s3::error::SdkError::service_error(
TestServiceError::new("AccessDenied", "access denied by policy"),
(),
);
let message = format_s3_sdk_error("upload", &err);
assert!(message.contains("S3 upload failed: AccessDenied"));
assert!(message.contains("access denied by policy"));
}
#[tokio::test]
async fn explicit_credentials_require_access_key_and_secret() {
let config = S3Config {
bucket: "bucket".to_string(),
region: "us-east-1".to_string(),
use_default_credentials: false,
..Default::default()
};
let err = match S3BackupClient::new(config).await {
Ok(_) => panic!("config without explicit credentials should fail"),
Err(err) => err.to_string(),
};
assert!(err.contains("S3 access key ID is required"));
}
#[tokio::test]
async fn default_credentials_still_require_bucket_and_region() {
let config = S3Config {
use_default_credentials: true,
..Default::default()
};
let err = match S3BackupClient::new(config).await {
Ok(_) => panic!("config without bucket should fail"),
Err(err) => err.to_string(),
};
assert!(err.contains("S3 bucket is required"));
}
#[derive(Debug)]
struct TestServiceError {
meta: aws_sdk_s3::error::ErrorMetadata,
}
impl TestServiceError {
fn new(code: &str, message: &str) -> Self {
Self {
meta: aws_sdk_s3::error::ErrorMetadata::builder()
.code(code)
.message(message)
.build(),
}
}
}
impl std::fmt::Display for TestServiceError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "test service error")
}
}
impl std::error::Error for TestServiceError {}
impl aws_sdk_s3::error::ProvideErrorMetadata for TestServiceError {
fn meta(&self) -> &aws_sdk_s3::error::ErrorMetadata {
&self.meta
}
}
}
+384
View File
@@ -0,0 +1,384 @@
use std::time::Instant;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use crate::error::{BoosApiError, Result};
// ── Response types ────────────────────────────────────────
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchResult {
pub title: String,
pub content: String,
pub url: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchResponse {
pub ok: bool,
pub query: String,
pub results: Vec<SearchResult>,
pub latency_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TestResult {
pub ok: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub latency_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub result_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
// ── Default endpoints ─────────────────────────────────────
pub fn default_endpoint(provider_type: &str) -> &'static str {
match provider_type {
"tavily" => "https://api.tavily.com/search",
"zhipu" => "https://open.bigmodel.cn/api/paas/v4/web_search",
"bocha" => "https://api.bochaai.com/v1/web-search",
_ => "",
}
}
// ── Main search dispatch ──────────────────────────────────
pub async fn execute_search(
provider_type: &str,
endpoint: Option<&str>,
api_key: &str,
query: &str,
max_results: i32,
timeout_ms: i32,
) -> Result<SearchResponse> {
let start = Instant::now();
let result = match provider_type {
"tavily" => search_tavily(endpoint, api_key, query, max_results, timeout_ms).await,
"zhipu" => search_zhipu(endpoint, api_key, query, max_results, timeout_ms).await,
"bocha" => search_bocha(endpoint, api_key, query, max_results, timeout_ms).await,
_ => {
return Err(BoosApiError::Validation(format!(
"Unsupported provider type: {}",
provider_type
)));
}
};
let latency = start.elapsed().as_millis() as u64;
match result {
Ok(results) => Ok(SearchResponse {
ok: true,
query: query.to_string(),
results,
latency_ms: latency,
error: None,
}),
Err(e) => Ok(SearchResponse {
ok: false,
query: query.to_string(),
results: vec![],
latency_ms: latency,
error: Some(e.to_string()),
}),
}
}
pub async fn test_provider(
provider_type: &str,
endpoint: Option<&str>,
api_key: &str,
timeout_ms: i32,
) -> TestResult {
let resp = execute_search(provider_type, endpoint, api_key, "test", 3, timeout_ms).await;
match resp {
Ok(r) if r.ok => TestResult {
ok: true,
latency_ms: Some(r.latency_ms),
result_count: Some(r.results.len()),
error: None,
},
Ok(r) => TestResult {
ok: false,
latency_ms: Some(r.latency_ms),
result_count: None,
error: r.error,
},
Err(e) => TestResult {
ok: false,
latency_ms: None,
result_count: None,
error: Some(e.to_string()),
},
}
}
// ── Tavily ────────────────────────────────────────────────
// POST {endpoint}
// Body: { api_key, query, max_results }
// Response: { results: [{ title, content, url }] }
#[derive(Serialize)]
struct TavilyRequest<'a> {
api_key: &'a str,
query: &'a str,
max_results: i32,
}
#[derive(Deserialize)]
struct TavilyResponse {
results: Vec<TavilyResult>,
}
#[derive(Deserialize)]
struct TavilyResult {
title: Option<String>,
content: Option<String>,
url: Option<String>,
}
async fn search_tavily(
endpoint: Option<&str>,
api_key: &str,
query: &str,
max_results: i32,
timeout_ms: i32,
) -> Result<Vec<SearchResult>> {
let url = endpoint.unwrap_or("https://api.tavily.com/search");
let client = Client::builder()
.timeout(std::time::Duration::from_millis(timeout_ms as u64))
.build()
.map_err(|e| BoosApiError::Provider(format!("HTTP client error: {e}")))?;
let body = TavilyRequest {
api_key,
query,
max_results: max_results.max(1),
};
let resp = client
.post(url)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| BoosApiError::Provider(format!("Tavily request failed: {e}")))?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(BoosApiError::Provider(format!(
"Tavily API error {status}: {text}"
)));
}
let data: TavilyResponse = resp
.json()
.await
.map_err(|e| BoosApiError::Provider(format!("Tavily response parse error: {e}")))?;
Ok(data
.results
.into_iter()
.take(max_results as usize)
.map(|r| SearchResult {
title: r.title.unwrap_or_else(|| "No title".to_string()),
content: r.content.unwrap_or_default(),
url: r.url.unwrap_or_default(),
})
.collect())
}
// ── Zhipu (智谱) ─────────────────────────────────────────
// POST {endpoint}
// Header: Authorization: Bearer {apiKey}
// Body: { search_query, search_engine: "search_std" }
// Response: { search_result: [{ title, content, link }] }
#[derive(Serialize)]
struct ZhipuRequest<'a> {
search_query: &'a str,
search_engine: &'a str,
}
#[derive(Deserialize)]
struct ZhipuResponse {
search_result: Option<Vec<ZhipuResult>>,
}
#[derive(Deserialize)]
struct ZhipuResult {
title: Option<String>,
content: Option<String>,
link: Option<String>,
}
async fn search_zhipu(
endpoint: Option<&str>,
api_key: &str,
query: &str,
max_results: i32,
timeout_ms: i32,
) -> Result<Vec<SearchResult>> {
let url = endpoint.unwrap_or("https://open.bigmodel.cn/api/paas/v4/web_search");
let client = Client::builder()
.timeout(std::time::Duration::from_millis(timeout_ms as u64))
.build()
.map_err(|e| BoosApiError::Provider(format!("HTTP client error: {e}")))?;
let body = ZhipuRequest {
search_query: query,
search_engine: "search_std",
};
let resp = client
.post(url)
.header("Authorization", format!("Bearer {api_key}"))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| BoosApiError::Provider(format!("Zhipu request failed: {e}")))?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(BoosApiError::Provider(format!(
"Zhipu API error {status}: {text}"
)));
}
let data: ZhipuResponse = resp
.json()
.await
.map_err(|e| BoosApiError::Provider(format!("Zhipu response parse error: {e}")))?;
let results = data.search_result.unwrap_or_default();
Ok(results
.into_iter()
.take(max_results as usize)
.map(|r| SearchResult {
title: r.title.unwrap_or_else(|| "No title".to_string()),
content: r.content.unwrap_or_default(),
url: r.link.unwrap_or_default(),
})
.collect())
}
// ── Bocha (博查) ──────────────────────────────────────────
// POST {endpoint}
// Header: Authorization: Bearer {apiKey}
// Body: { query, count, summary: true, page: 1 }
// Response: { code, data: { webPages: { value: [{ name, url, snippet, summary }] } } }
#[derive(Serialize)]
struct BochaRequest<'a> {
query: &'a str,
count: i32,
summary: bool,
page: i32,
}
#[derive(Deserialize)]
struct BochaResponse {
code: Option<i32>,
msg: Option<String>,
data: Option<BochaData>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct BochaData {
web_pages: Option<BochaWebPages>,
}
#[derive(Deserialize)]
struct BochaWebPages {
value: Option<Vec<BochaWebResult>>,
}
#[derive(Deserialize)]
struct BochaWebResult {
name: Option<String>,
url: Option<String>,
snippet: Option<String>,
summary: Option<String>,
}
async fn search_bocha(
endpoint: Option<&str>,
api_key: &str,
query: &str,
max_results: i32,
timeout_ms: i32,
) -> Result<Vec<SearchResult>> {
let url = endpoint.unwrap_or("https://api.bochaai.com/v1/web-search");
let client = Client::builder()
.timeout(std::time::Duration::from_millis(timeout_ms as u64))
.build()
.map_err(|e| BoosApiError::Provider(format!("HTTP client error: {e}")))?;
let body = BochaRequest {
query,
count: max_results.max(1),
summary: true,
page: 1,
};
let resp = client
.post(url)
.header("Authorization", format!("Bearer {api_key}"))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| BoosApiError::Provider(format!("Bocha request failed: {e}")))?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(BoosApiError::Provider(format!(
"Bocha API error {status}: {text}"
)));
}
let data: BochaResponse = resp
.json()
.await
.map_err(|e| BoosApiError::Provider(format!("Bocha response parse error: {e}")))?;
if data.code.unwrap_or(0) != 200 {
return Err(BoosApiError::Provider(format!(
"Bocha search failed: {}",
data.msg.unwrap_or_else(|| "Unknown error".to_string())
)));
}
let results = data
.data
.and_then(|d| d.web_pages)
.and_then(|wp| wp.value)
.unwrap_or_default();
Ok(results
.into_iter()
.take(max_results as usize)
.map(|r| SearchResult {
title: r.name.unwrap_or_else(|| "No title".to_string()),
content: r.summary.or(r.snippet).unwrap_or_default(),
url: r.url.unwrap_or_default(),
})
.collect())
}
@@ -0,0 +1,115 @@
use crate::storage_paths;
use serde::Serialize;
use std::path::Path;
#[derive(Debug, Serialize, Clone)]
pub struct BucketStats {
pub bucket: String,
pub file_count: u64,
pub total_bytes: u64,
}
#[derive(Debug, Serialize, Clone)]
pub struct StorageInventory {
pub buckets: Vec<BucketStats>,
pub documents_root: String,
}
pub fn scan_storage() -> StorageInventory {
let root = storage_paths::documents_root();
let buckets = ["images", "files", "backups"]
.iter()
.map(|name| {
let dir = root.join(name);
let (count, bytes) = if dir.exists() {
count_dir_contents(&dir)
} else {
(0, 0)
};
BucketStats {
bucket: name.to_string(),
file_count: count,
total_bytes: bytes,
}
})
.collect();
StorageInventory {
buckets,
documents_root: root.to_string_lossy().to_string(),
}
}
fn count_dir_contents(dir: &Path) -> (u64, u64) {
let mut count = 0u64;
let mut bytes = 0u64;
if let Ok(entries) = std::fs::read_dir(dir) {
for entry in entries.flatten() {
if let Ok(meta) = entry.metadata() {
if meta.is_file() {
count += 1;
bytes += meta.len();
}
}
}
}
(count, bytes)
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
#[test]
fn test_scan_empty_storage() {
let tmp = tempfile::tempdir().unwrap();
let images = tmp.path().join("images");
let files = tmp.path().join("files");
let backups = tmp.path().join("backups");
fs::create_dir_all(&images).unwrap();
fs::create_dir_all(&files).unwrap();
fs::create_dir_all(&backups).unwrap();
for name in &["images", "files", "backups"] {
let (count, bytes) = count_dir_contents(&tmp.path().join(name));
assert_eq!(count, 0, "{name} should have 0 files");
assert_eq!(bytes, 0, "{name} should have 0 bytes");
}
}
#[test]
fn test_scan_with_files() {
let tmp = tempfile::tempdir().unwrap();
let images = tmp.path().join("images");
fs::create_dir_all(&images).unwrap();
// Write two files of known sizes
fs::write(images.join("a.png"), vec![0u8; 1024]).unwrap();
fs::write(images.join("b.jpg"), vec![0u8; 2048]).unwrap();
let (count, bytes) = count_dir_contents(&images);
assert_eq!(count, 2);
assert_eq!(bytes, 3072);
}
#[test]
fn test_count_dir_nonexistent() {
let tmp = tempfile::tempdir().unwrap();
let missing = tmp.path().join("nonexistent");
let (count, bytes) = count_dir_contents(&missing);
assert_eq!(count, 0);
assert_eq!(bytes, 0);
}
#[test]
fn test_count_dir_skips_subdirectories() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("mixed");
fs::create_dir_all(dir.join("subdir")).unwrap();
fs::write(dir.join("file.txt"), vec![0u8; 512]).unwrap();
let (count, bytes) = count_dir_contents(&dir);
assert_eq!(count, 1, "should only count files, not subdirs");
assert_eq!(bytes, 512);
}
}
@@ -0,0 +1,502 @@
//! One-time migration: move files from legacy ~/.boosapi/files/ to ~/Documents/BoosAPI Desktop/
use std::collections::HashMap;
use std::path::Path;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
use crate::entity::{messages, stored_files};
use crate::storage_paths::{build_relative_path, documents_root};
/// Summary of what the migration did.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct MigrationReport {
pub files_moved: usize,
pub files_skipped: usize,
pub files_missing: usize,
pub db_records_updated: usize,
pub messages_updated: usize,
}
/// Migrates files from legacy ~/.boosapi/files/ to ~/Documents/BoosAPI Desktop/{images,files}/
/// Updates stored_files.storage_path and messages.attachments JSON in the database.
/// Idempotent: safe to run multiple times.
pub async fn migrate_to_documents_root(
db: &DatabaseConnection,
legacy_files_dir: &Path,
) -> Result<MigrationReport, String> {
let target = documents_root();
run_migration(db, legacy_files_dir, &target).await
}
/// Internal entry point with an explicit target root, for testability.
async fn run_migration(
db: &DatabaseConnection,
legacy_dir: &Path,
target_root: &Path,
) -> Result<MigrationReport, String> {
// 1. Ensure target directories
for sub in &["images", "files"] {
std::fs::create_dir_all(target_root.join(sub))
.map_err(|e| format!("create target dir: {e}"))?;
}
let mut report = MigrationReport::default();
let mut path_map: HashMap<String, String> = HashMap::new();
// 2. Query all stored_files
let rows = stored_files::Entity::find()
.all(db)
.await
.map_err(|e| format!("query stored_files: {e}"))?;
// 3. Process each stored file
for row in rows {
// a. Already in new format — skip
if row.storage_path.starts_with("images/") || row.storage_path.starts_with("files/") {
report.files_skipped += 1;
continue;
}
// b. Old absolute path
let old_abs = legacy_dir.join(&row.storage_path);
// c. New relative path
let new_rel = build_relative_path(&row.original_name, &row.mime_type, &row.hash);
// d. New absolute path
let new_abs = target_root.join(&new_rel);
// Remember mapping for message-attachment update
path_map.insert(row.storage_path.clone(), new_rel.clone());
// e. Copy if source exists, f. warn if missing
if old_abs.exists() {
if let Some(p) = new_abs.parent() {
std::fs::create_dir_all(p).map_err(|e| format!("mkdir: {e}"))?;
}
std::fs::copy(&old_abs, &new_abs)
.map_err(|e| format!("copy {}: {e}", old_abs.display()))?;
report.files_moved += 1;
} else {
tracing::warn!(path = %old_abs.display(), "source file missing during migration");
report.files_missing += 1;
}
// Update DB record (even when source is missing, for consistency)
let mut am: stored_files::ActiveModel = row.into();
am.storage_path = Set(new_rel);
am.update(db)
.await
.map_err(|e| format!("update stored_files: {e}"))?;
report.db_records_updated += 1;
}
// 4. Update message attachment paths
let msgs = messages::Entity::find()
.filter(messages::Column::Attachments.ne("[]"))
.filter(messages::Column::Attachments.ne(""))
.all(db)
.await
.map_err(|e| format!("query messages: {e}"))?;
for msg in msgs {
let Ok(mut atts) = serde_json::from_str::<Vec<serde_json::Value>>(&msg.attachments) else {
continue;
};
let mut changed = false;
for att in &mut atts {
let Some(fp) = att
.get("file_path")
.and_then(|v| v.as_str())
.map(String::from)
else {
continue;
};
if fp.starts_with("images/") || fp.starts_with("files/") {
continue;
}
let new_fp = if let Some(p) = path_map.get(&fp) {
p.clone()
} else {
// Fallback: compute from attachment metadata
let name = att
.get("file_name")
.and_then(|v| v.as_str())
.unwrap_or("file");
let mime = att
.get("file_type")
.and_then(|v| v.as_str())
.unwrap_or("application/octet-stream");
let hash = std::path::Path::new(&fp)
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("unknown");
build_relative_path(name, mime, hash)
};
att["file_path"] = serde_json::Value::String(new_fp);
changed = true;
}
if changed {
let json =
serde_json::to_string(&atts).map_err(|e| format!("serialize attachments: {e}"))?;
let mut am: messages::ActiveModel = msg.into();
am.attachments = Set(json);
am.update(db)
.await
.map_err(|e| format!("update message: {e}"))?;
report.messages_updated += 1;
}
}
Ok(report)
}
#[cfg(test)]
mod tests {
use super::*;
use sea_orm::{ConnectionTrait, Database, DbBackend, Statement};
use std::fs;
const CREATE_STORED_FILES: &str = "
CREATE TABLE stored_files (
id TEXT PRIMARY KEY,
hash TEXT NOT NULL,
original_name TEXT NOT NULL,
mime_type TEXT NOT NULL,
size_bytes INTEGER NOT NULL,
storage_path TEXT NOT NULL,
conversation_id TEXT,
created_at TEXT NOT NULL
)";
const CREATE_MESSAGES: &str = "
CREATE TABLE messages (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
provider_id TEXT,
model_id TEXT,
token_count INTEGER,
attachments TEXT NOT NULL DEFAULT '[]',
thinking TEXT,
created_at INTEGER NOT NULL DEFAULT 0,
branch_id TEXT,
parent_message_id TEXT,
version_index INTEGER NOT NULL DEFAULT 0,
is_active INTEGER NOT NULL DEFAULT 1,
tool_calls_json TEXT,
tool_call_id TEXT
)";
async fn test_db() -> DatabaseConnection {
let db = Database::connect("sqlite::memory:").await.unwrap();
for ddl in [CREATE_STORED_FILES, CREATE_MESSAGES] {
db.execute(Statement::from_string(DbBackend::Sqlite, ddl))
.await
.unwrap();
}
db
}
struct TestDirs {
_root: tempfile::TempDir,
legacy: std::path::PathBuf,
target: std::path::PathBuf,
}
fn test_dirs() -> TestDirs {
let root = tempfile::tempdir().unwrap();
let legacy = root.path().join("legacy");
let target = root.path().join("target");
fs::create_dir_all(&legacy).unwrap();
fs::create_dir_all(&target).unwrap();
TestDirs {
_root: root,
legacy,
target,
}
}
async fn insert_stored_file(
db: &DatabaseConnection,
id: &str,
hash: &str,
name: &str,
mime: &str,
path: &str,
conv: Option<&str>,
) {
let am = stored_files::ActiveModel {
id: Set(id.into()),
hash: Set(hash.into()),
original_name: Set(name.into()),
mime_type: Set(mime.into()),
size_bytes: Set(100),
storage_path: Set(path.into()),
conversation_id: Set(conv.map(String::from)),
created_at: Set("2024-01-01".into()),
};
am.insert(db).await.unwrap();
}
async fn insert_message(db: &DatabaseConnection, id: &str, conv: &str, attachments: &str) {
let am = messages::ActiveModel {
id: Set(id.into()),
conversation_id: Set(conv.into()),
role: Set("user".into()),
content: Set("hi".into()),
provider_id: Set(None),
model_id: Set(None),
token_count: Set(None),
prompt_tokens: Set(None),
completion_tokens: Set(None),
attachments: Set(attachments.into()),
thinking: Set(None),
created_at: Set(0),
branch_id: Set(None),
parent_message_id: Set(None),
version_index: Set(0),
is_active: Set(1),
tool_calls_json: Set(None),
tool_call_id: Set(None),
status: Set("complete".into()),
tokens_per_second: Set(None),
first_token_latency_ms: Set(None),
};
am.insert(db).await.unwrap();
}
#[tokio::test]
async fn empty_db_is_noop() {
let db = test_db().await;
let dirs = test_dirs();
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r, MigrationReport::default());
}
#[tokio::test]
async fn already_migrated_paths_are_skipped() {
let db = test_db().await;
let dirs = test_dirs();
insert_stored_file(
&db,
"f1",
"aaa",
"pic.png",
"image/png",
"images/aaa_pic.png",
None,
)
.await;
insert_stored_file(
&db,
"f2",
"bbb",
"doc.pdf",
"application/pdf",
"files/bbb_doc.pdf",
None,
)
.await;
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r.files_skipped, 2);
assert_eq!(r.files_moved, 0);
assert_eq!(r.db_records_updated, 0);
}
#[tokio::test]
async fn migrates_image_file() {
let db = test_db().await;
let dirs = test_dirs();
let old_dir = dirs.legacy.join("conv1");
fs::create_dir_all(&old_dir).unwrap();
fs::write(old_dir.join("abcdef123456789.png"), b"PNG_DATA").unwrap();
insert_stored_file(
&db,
"f1",
"abcdef123456789",
"photo.png",
"image/png",
"conv1/abcdef123456789.png",
Some("conv1"),
)
.await;
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r.files_moved, 1);
assert_eq!(r.db_records_updated, 1);
// File copied to new location
let expected = dirs.target.join("images/abcdef123456_photo.png");
assert!(expected.exists(), "target file should exist");
assert_eq!(fs::read(&expected).unwrap(), b"PNG_DATA");
// Legacy file still exists (non-destructive)
assert!(old_dir.join("abcdef123456789.png").exists());
// DB updated
let f = stored_files::Entity::find_by_id("f1")
.one(&db)
.await
.unwrap()
.unwrap();
assert_eq!(f.storage_path, "images/abcdef123456_photo.png");
}
#[tokio::test]
async fn migrates_non_image_file() {
let db = test_db().await;
let dirs = test_dirs();
let old_dir = dirs.legacy.join("conv1");
fs::create_dir_all(&old_dir).unwrap();
fs::write(old_dir.join("fedcba654321abc.pdf"), b"PDF").unwrap();
insert_stored_file(
&db,
"f1",
"fedcba654321abc",
"report.pdf",
"application/pdf",
"conv1/fedcba654321abc.pdf",
Some("conv1"),
)
.await;
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r.files_moved, 1);
let expected = dirs.target.join("files/fedcba654321_report.pdf");
assert!(expected.exists());
let f = stored_files::Entity::find_by_id("f1")
.one(&db)
.await
.unwrap()
.unwrap();
assert_eq!(f.storage_path, "files/fedcba654321_report.pdf");
}
#[tokio::test]
async fn missing_source_still_updates_db() {
let db = test_db().await;
let dirs = test_dirs();
// No file on disk — only a DB record
insert_stored_file(
&db,
"f1",
"abcdef123456789",
"photo.png",
"image/png",
"conv1/abcdef123456789.png",
Some("conv1"),
)
.await;
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r.files_missing, 1);
assert_eq!(r.files_moved, 0);
assert_eq!(r.db_records_updated, 1);
let f = stored_files::Entity::find_by_id("f1")
.one(&db)
.await
.unwrap()
.unwrap();
assert_eq!(f.storage_path, "images/abcdef123456_photo.png");
}
#[tokio::test]
async fn idempotent_second_run_skips() {
let db = test_db().await;
let dirs = test_dirs();
let old_dir = dirs.legacy.join("conv1");
fs::create_dir_all(&old_dir).unwrap();
fs::write(old_dir.join("abcdef123456789.png"), b"IMG").unwrap();
insert_stored_file(
&db,
"f1",
"abcdef123456789",
"photo.png",
"image/png",
"conv1/abcdef123456789.png",
Some("conv1"),
)
.await;
let r1 = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r1.files_moved, 1);
assert_eq!(r1.db_records_updated, 1);
let r2 = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r2.files_skipped, 1);
assert_eq!(r2.files_moved, 0);
assert_eq!(r2.db_records_updated, 0);
}
#[tokio::test]
async fn updates_message_attachments() {
let db = test_db().await;
let dirs = test_dirs();
let old_dir = dirs.legacy.join("conv1");
fs::create_dir_all(&old_dir).unwrap();
fs::write(old_dir.join("abcdef123456789.png"), b"IMG").unwrap();
insert_stored_file(
&db,
"f1",
"abcdef123456789",
"photo.png",
"image/png",
"conv1/abcdef123456789.png",
Some("conv1"),
)
.await;
let att_json = r#"[{"id":"a1","file_type":"image/png","file_name":"photo.png","file_path":"conv1/abcdef123456789.png","file_size":100}]"#;
insert_message(&db, "m1", "conv1", att_json).await;
let r = run_migration(&db, &dirs.legacy, &dirs.target)
.await
.unwrap();
assert_eq!(r.messages_updated, 1);
let m = messages::Entity::find_by_id("m1")
.one(&db)
.await
.unwrap()
.unwrap();
let atts: Vec<serde_json::Value> = serde_json::from_str(&m.attachments).unwrap();
assert_eq!(
atts[0]["file_path"].as_str().unwrap(),
"images/abcdef123456_photo.png"
);
}
}
@@ -0,0 +1,326 @@
use std::path::{Path, PathBuf};
use std::sync::RwLock;
static DOCUMENTS_ROOT_OVERRIDE: RwLock<Option<PathBuf>> = RwLock::new(None);
/// Initialise the custom documents root from a stored setting.
/// Call once during app startup; ignored if `custom` is `None`.
pub fn init_documents_root(custom: Option<PathBuf>) {
if let Some(path) = custom {
if let Ok(mut guard) = DOCUMENTS_ROOT_OVERRIDE.write() {
*guard = Some(path);
}
}
}
/// Replace the documents root at runtime (e.g. after the user picks a new
/// directory). Takes effect immediately for all subsequent `documents_root()`
/// calls.
pub fn set_documents_root(path: PathBuf) {
if let Ok(mut guard) = DOCUMENTS_ROOT_OVERRIDE.write() {
*guard = Some(path);
}
}
/// Clear any custom override so `documents_root()` falls back to the default.
pub fn clear_documents_root_override() {
if let Ok(mut guard) = DOCUMENTS_ROOT_OVERRIDE.write() {
*guard = None;
}
}
/// Returns the active documents root — custom override if set, otherwise the
/// platform default (`~/Documents/BoosAPI Desktop/`).
pub fn documents_root() -> PathBuf {
if let Ok(guard) = DOCUMENTS_ROOT_OVERRIDE.read() {
if let Some(ref custom) = *guard {
return custom.clone();
}
}
default_documents_root()
}
/// The platform default documents root: `~/Documents/BoosAPI Desktop/`.
pub fn default_documents_root() -> PathBuf {
dirs::document_dir()
.expect("Could not determine Documents directory")
.join("BoosAPI Desktop")
}
/// Returns the typed subdirectory for a given MIME type.
/// - "image/*" → "images"
/// - everything else → "files"
/// - "backup" sentinel → "backups"
pub fn file_type_bucket(mime_type: &str) -> &'static str {
if mime_type == "backup" {
"backups"
} else if mime_type.starts_with("image/") {
"images"
} else {
"files"
}
}
/// Resolves a relative storage path to absolute under documents root.
pub fn resolve_documents_path(relative_path: &str) -> PathBuf {
documents_root().join(relative_path)
}
/// Generates a storage-ready relative path for a new file.
/// Format: "{bucket}/{hash_prefix}_{sanitized_name}"
pub fn build_relative_path(original_name: &str, mime_type: &str, hash: &str) -> String {
let bucket = file_type_bucket(mime_type);
let hash_prefix = &hash[..hash.len().min(12)];
let sanitized = sanitize_filename(original_name);
format!("{}/{}_{}", bucket, hash_prefix, sanitized)
}
/// Validates a relative path (no traversal, no absolute, lowercase dir).
pub fn validate_relative_path(path: &str) -> Result<(), String> {
if path.is_empty() {
return Err("path must not be empty".to_string());
}
if path.starts_with('/') || path.starts_with('\\') {
return Err("path must not be absolute".to_string());
}
if path.contains("..") {
return Err("path must not contain '..' traversal".to_string());
}
Ok(())
}
/// Ensures the documents root and subdirectories exist.
pub fn ensure_documents_dirs() -> std::io::Result<()> {
let root = documents_root();
for sub in &["images", "files", "backups"] {
std::fs::create_dir_all(root.join(sub))?;
}
Ok(())
}
/// Sanitize a filename: replace spaces with `_`, remove special chars, keep extension.
fn sanitize_filename(name: &str) -> String {
let path = Path::new(name);
let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("file");
let ext = path.extension().and_then(|e| e.to_str());
let sanitized_stem: String = stem
.chars()
.filter_map(|c| {
if c.is_alphanumeric() || c == '_' || c == '-' {
Some(c)
} else if c == ' ' {
Some('_')
} else {
None
}
})
.collect();
let sanitized_stem = if sanitized_stem.is_empty() {
"file".to_string()
} else {
sanitized_stem
};
match ext {
Some(e) => format!("{}.{}", sanitized_stem, e),
None => sanitized_stem,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn documents_root_ends_with_boosapi() {
let root = default_documents_root();
assert!(
root.ends_with("BoosAPI Desktop"),
"Expected path ending with 'BoosApi', got {:?}",
root
);
// Should be under the platform Documents directory
let parent = root.parent().unwrap();
let parent_name = parent.file_name().unwrap().to_str().unwrap();
assert_eq!(
parent_name, "Documents",
"Expected parent 'Documents', got {}",
parent_name
);
}
#[test]
fn documents_root_is_absolute() {
let root = default_documents_root();
assert!(root.is_absolute(), "Expected absolute path, got {:?}", root);
}
// -- file_type_bucket tests --
#[test]
fn bucket_image_png() {
assert_eq!(file_type_bucket("image/png"), "images");
}
#[test]
fn bucket_image_jpeg() {
assert_eq!(file_type_bucket("image/jpeg"), "images");
}
#[test]
fn bucket_image_gif() {
assert_eq!(file_type_bucket("image/gif"), "images");
}
#[test]
fn bucket_image_webp() {
assert_eq!(file_type_bucket("image/webp"), "images");
}
#[test]
fn bucket_application_pdf() {
assert_eq!(file_type_bucket("application/pdf"), "files");
}
#[test]
fn bucket_text_plain() {
assert_eq!(file_type_bucket("text/plain"), "files");
}
#[test]
fn bucket_application_zip() {
assert_eq!(file_type_bucket("application/zip"), "files");
}
#[test]
fn bucket_backup_sentinel() {
assert_eq!(file_type_bucket("backup"), "backups");
}
// -- resolve_documents_path tests --
#[test]
fn resolve_joins_correctly() {
let resolved = resolve_documents_path("images/abc123.jpg");
let root = documents_root();
assert_eq!(resolved, root.join("images").join("abc123.jpg"));
}
#[test]
fn resolve_single_file() {
let resolved = resolve_documents_path("readme.txt");
let root = documents_root();
assert_eq!(resolved, root.join("readme.txt"));
}
// -- build_relative_path tests --
#[test]
fn build_path_image() {
let result = build_relative_path("photo.jpg", "image/png", "abcdef123456789xyz");
assert_eq!(result, "images/abcdef123456_photo.jpg");
}
#[test]
fn build_path_pdf() {
let result = build_relative_path("report.pdf", "application/pdf", "fedcba987654321abc");
assert_eq!(result, "files/fedcba987654_report.pdf");
}
#[test]
fn build_path_sanitizes_spaces() {
let result = build_relative_path("my photo file.png", "image/png", "aabbccddee11");
assert_eq!(result, "images/aabbccddee11_my_photo_file.png");
}
#[test]
fn build_path_sanitizes_special_chars() {
let result = build_relative_path("file@#$%name!.txt", "text/plain", "112233445566");
assert_eq!(result, "files/112233445566_filename.txt");
}
#[test]
fn build_path_truncates_hash_to_12() {
let result = build_relative_path("a.txt", "text/plain", "123456789012345");
assert!(result.starts_with("files/123456789012_"));
}
#[test]
fn build_path_short_hash() {
let result = build_relative_path("a.txt", "text/plain", "abc");
assert_eq!(result, "files/abc_a.txt");
}
// -- validate_relative_path tests --
#[test]
fn validate_accepts_valid_path() {
assert!(validate_relative_path("images/abc123.jpg").is_ok());
}
#[test]
fn validate_accepts_files_path() {
assert!(validate_relative_path("files/report.pdf").is_ok());
}
#[test]
fn validate_rejects_traversal() {
let result = validate_relative_path("images/../etc/passwd");
assert!(result.is_err());
assert!(result.unwrap_err().contains(".."));
}
#[test]
fn validate_rejects_absolute_unix() {
let result = validate_relative_path("/etc/passwd");
assert!(result.is_err());
}
#[test]
fn validate_rejects_absolute_windows() {
let result = validate_relative_path("\\Windows\\system32");
assert!(result.is_err());
}
#[test]
fn validate_rejects_empty() {
let result = validate_relative_path("");
assert!(result.is_err());
}
#[test]
fn validate_rejects_dotdot_only() {
let result = validate_relative_path("..");
assert!(result.is_err());
}
// -- ensure_documents_dirs tests --
#[test]
fn ensure_dirs_creates_structure() {
// Use a temp dir to avoid side effects
let tmp = std::env::temp_dir().join("boosapi_test_ensure_dirs");
let _ = std::fs::remove_dir_all(&tmp);
// We test the logic by verifying the dirs exist after calling ensure_documents_dirs.
// Since ensure_documents_dirs uses documents_root(), we test it indirectly:
// just verify it doesn't error and the root exists afterward.
let result = ensure_documents_dirs();
assert!(
result.is_ok(),
"ensure_documents_dirs failed: {:?}",
result.err()
);
let root = documents_root();
assert!(
root.exists(),
"documents root should exist after ensure_documents_dirs"
);
assert!(root.join("images").exists(), "images/ should exist");
assert!(root.join("files").exists(), "files/ should exist");
assert!(root.join("backups").exists(), "backups/ should exist");
}
}
@@ -0,0 +1,183 @@
/// A chunk of text with its position index.
#[derive(Debug, Clone)]
pub struct TextChunk {
pub index: i32,
pub content: String,
}
/// Default chunk size in characters (~500 tokens).
pub const DEFAULT_CHUNK_SIZE: usize = 2000;
/// Default overlap in characters (~50 tokens).
pub const DEFAULT_OVERLAP: usize = 200;
/// Find the nearest char boundary at or before the given byte position.
fn floor_char_boundary(s: &str, index: usize) -> usize {
if index >= s.len() {
s.len()
} else {
let mut i = index;
while i > 0 && !s.is_char_boundary(i) {
i -= 1;
}
i
}
}
/// Find the nearest char boundary at or after the given byte position.
fn ceil_char_boundary(s: &str, index: usize) -> usize {
if index >= s.len() {
s.len()
} else {
let mut i = index;
while i < s.len() && !s.is_char_boundary(i) {
i += 1;
}
i
}
}
/// Split text into overlapping chunks, breaking at paragraph/sentence boundaries.
pub fn chunk_text(text: &str, chunk_size: usize, overlap: usize) -> Vec<TextChunk> {
let text = text.trim();
if text.is_empty() {
return vec![];
}
if text.len() <= chunk_size {
return vec![TextChunk {
index: 0,
content: text.to_string(),
}];
}
let mut chunks = Vec::new();
let mut start = 0;
while start < text.len() {
// Snap to char boundary to avoid slicing inside multi-byte chars (e.g. CJK)
let end = floor_char_boundary(text, (start + chunk_size).min(text.len()));
// Find a good break point near `end`
let actual_end = if end >= text.len() {
text.len()
} else {
find_break_point(text, start, end)
};
let chunk_content = text[start..actual_end].trim();
if !chunk_content.is_empty() {
chunks.push(TextChunk {
index: chunks.len() as i32,
content: chunk_content.to_string(),
});
}
// Move start forward by (chunk_size - overlap), but at least 1 char
let advance = if actual_end - start > overlap {
actual_end - start - overlap
} else {
actual_end - start
};
// Snap new start to a char boundary
start = ceil_char_boundary(text, start + advance.max(1));
// If remaining text is tiny, it's already covered by the last chunk's overlap
if start >= text.len() || text.len() - start < overlap {
break;
}
}
chunks
}
/// Find a good break point near `target` position, searching backwards from target.
/// Prefers: paragraph break (\n\n) > line break (\n) > sentence end (. ! ?) > space
fn find_break_point(text: &str, start: usize, target: usize) -> usize {
let search_range = &text[start..target];
let min_chunk = (target - start) / 2; // Don't break before half the chunk
// Try paragraph break
if let Some(pos) = search_range.rfind("\n\n") {
if pos >= min_chunk {
return start + pos + 2; // After the double newline
}
}
// Try line break
if let Some(pos) = search_range.rfind('\n') {
if pos >= min_chunk {
return start + pos + 1;
}
}
// Try sentence end
let bytes = search_range.as_bytes();
for i in (min_chunk..bytes.len()).rev() {
if matches!(bytes[i], b'.' | b'!' | b'?') {
// Check it's followed by a space or end
if i + 1 >= bytes.len() || bytes[i + 1] == b' ' || bytes[i + 1] == b'\n' {
return start + i + 1;
}
}
}
// Try word break (space)
if let Some(pos) = search_range.rfind(' ') {
if pos >= min_chunk {
return start + pos + 1;
}
}
// No good break found, just cut at target
target
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_empty_text() {
assert!(chunk_text("", 100, 20).is_empty());
}
#[test]
fn test_short_text() {
let chunks = chunk_text("Hello world", 100, 20);
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].content, "Hello world");
}
#[test]
fn test_chunking_preserves_content() {
let text = "A".repeat(500);
let chunks = chunk_text(&text, 200, 50);
assert!(chunks.len() > 1);
// First chunk should be roughly 200 chars
assert!(chunks[0].content.len() <= 200);
}
#[test]
fn test_chunking_cjk_no_panic() {
// CJK characters are 3 bytes each in UTF-8.
// A chunk_size of 100 bytes lands inside a multi-byte char → must not panic.
let text = "中".repeat(200); // 600 bytes
let chunks = chunk_text(&text, 100, 20);
assert!(chunks.len() > 1);
for chunk in &chunks {
// Every chunk must be valid UTF-8 (no partial chars)
assert!(chunk.content.is_char_boundary(0));
assert!(chunk.content.is_char_boundary(chunk.content.len()));
}
}
#[test]
fn test_chunking_mixed_ascii_cjk() {
let text = "Hello世界!这是一段混合中英文的测试文本。".repeat(50);
let chunks = chunk_text(&text, 150, 30);
assert!(!chunks.is_empty());
for chunk in &chunks {
assert!(!chunk.content.is_empty());
}
}
}
@@ -0,0 +1,91 @@
//! Lightweight token estimation for context window management.
//!
//! Uses character-based heuristics rather than a full tokenizer to keep
//! dependencies minimal. Accuracy is sufficient for threshold detection
//! (within ~15% of actual token counts for mixed-language text).
/// Estimate the number of tokens in a text string.
///
/// Heuristic:
/// - ASCII / Latin characters: ~4 characters per token
/// - CJK / fullwidth characters: ~1.5 characters per token
/// - Each message carries ~4 tokens of overhead (role, delimiters)
pub fn estimate_tokens(text: &str) -> usize {
if text.is_empty() {
return 0;
}
let mut ascii_chars: usize = 0;
let mut cjk_chars: usize = 0;
for ch in text.chars() {
if is_cjk(ch) {
cjk_chars += 1;
} else {
ascii_chars += 1;
}
}
// English-ish: ~4 chars/token; CJK: ~1.5 chars/token
let ascii_tokens = (ascii_chars + 3) / 4; // ceil division
let cjk_tokens = (cjk_chars * 2 + 2) / 3; // ceil(cjk * 2/3)
ascii_tokens + cjk_tokens
}
/// Estimate tokens for an entire chat message (content + role overhead).
pub fn estimate_message_tokens(role: &str, content: &str) -> usize {
const PER_MESSAGE_OVERHEAD: usize = 4; // role, delimiters, etc.
estimate_tokens(role) + estimate_tokens(content) + PER_MESSAGE_OVERHEAD
}
/// Check if a character is in a CJK Unicode block.
fn is_cjk(ch: char) -> bool {
matches!(ch,
'\u{4E00}'..='\u{9FFF}' // CJK Unified Ideographs
| '\u{3400}'..='\u{4DBF}' // CJK Extension A
| '\u{F900}'..='\u{FAFF}' // CJK Compat Ideographs
| '\u{3000}'..='\u{303F}' // CJK Symbols and Punctuation
| '\u{FF00}'..='\u{FFEF}' // Halfwidth and Fullwidth Forms
| '\u{AC00}'..='\u{D7AF}' // Hangul Syllables
| '\u{3040}'..='\u{309F}' // Hiragana
| '\u{30A0}'..='\u{30FF}' // Katakana
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_empty_string() {
assert_eq!(estimate_tokens(""), 0);
}
#[test]
fn test_english_text() {
// "Hello, world!" = 13 chars → ~3 tokens (actual: 4 with GPT-4)
let tokens = estimate_tokens("Hello, world!");
assert!(tokens >= 2 && tokens <= 6);
}
#[test]
fn test_chinese_text() {
// "你好世界" = 4 CJK chars → ~3 tokens (actual: 2-4 with GPT-4)
let tokens = estimate_tokens("你好世界");
assert!(tokens >= 2 && tokens <= 5);
}
#[test]
fn test_mixed_text() {
let tokens = estimate_tokens("Hello 你好 world 世界");
assert!(tokens >= 4 && tokens <= 10);
}
#[test]
fn test_message_overhead() {
let tokens = estimate_message_tokens("user", "Hi");
// "user" + "Hi" + 4 overhead ≈ 6-7
assert!(tokens >= 5);
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,7 @@
pub fn gen_id() -> String {
uuid::Uuid::new_v4().to_string()
}
pub fn now_ts() -> i64 {
chrono::Utc::now().timestamp()
}
@@ -0,0 +1,855 @@
use sea_orm::{ConnectionTrait, DatabaseConnection, DbBackend, DbErr, Statement};
use crate::error::{BoosApiError, Result};
/// Register the sqlite-vec extension globally.
///
/// Must be called **once** before any SQLite connection is opened.
pub fn register_sqlite_vec_extension() {
unsafe {
libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute(
sqlite_vec::sqlite3_vec_init as *const (),
)));
}
}
/// A single embedding record for storage in the vector database.
#[derive(Debug, Clone)]
pub struct EmbeddingRecord {
/// Unique chunk identifier
pub id: String,
/// Parent document identifier
pub document_id: String,
/// Position of this chunk within the document
pub chunk_index: i32,
/// Text content of the chunk
pub content: String,
/// Embedding vector
pub embedding: Vec<f32>,
}
/// A result returned from vector similarity search.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct VectorSearchResult {
pub id: String,
pub document_id: String,
pub chunk_index: i32,
pub content: String,
/// Distance score (lower is more similar for L2 distance)
pub score: f32,
#[serde(
default,
rename = "rerankScore",
skip_serializing_if = "Option::is_none"
)]
pub rerank_score: Option<f32>,
/// Whether this chunk has an embedding in vec0
pub has_embedding: bool,
}
/// sqlite-vec–backed vector store for knowledge base embeddings.
///
/// Each knowledge base gets two tables in the shared SQLite database:
/// - `vec_{id}_meta` — chunk metadata (id, document_id, content, …)
/// - `vec_{id}` — vec0 virtual table holding the embedding vectors
pub struct VectorStore {
db: DatabaseConnection,
}
impl VectorStore {
/// Create a VectorStore that uses an existing sea-orm connection.
pub fn new(db: DatabaseConnection) -> Self {
Self { db }
}
/// Sanitised table-name prefix for a collection.
fn collection_name(collection_id: &str) -> String {
format!("vec_{}", collection_id.replace('-', "_"))
}
/// Ensure both the metadata and vec0 tables exist for a collection.
pub async fn ensure_collection(&self, collection_id: &str, dimensions: usize) -> Result<()> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
self.exec(&format!(
"CREATE TABLE IF NOT EXISTS {meta_table} (
rowid INTEGER PRIMARY KEY AUTOINCREMENT,
id TEXT NOT NULL UNIQUE,
document_id TEXT NOT NULL,
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL
)"
))
.await?;
self.exec(&format!(
"CREATE INDEX IF NOT EXISTS idx_{name}_doc ON {meta_table}(document_id)"
))
.await?;
if self.table_exists(&name).await?
&& self.collection_dimensions(&name).await? != Some(dimensions)
{
let _ = self.exec(&format!("DROP TABLE IF EXISTS {name}")).await;
}
self.exec(&format!(
"CREATE VIRTUAL TABLE IF NOT EXISTS {name} USING vec0(embedding float[{dimensions}])"
))
.await?;
Ok(())
}
/// Upsert embedding records for a single document.
///
/// All existing embeddings for the document (identified by `document_id` of
/// the first record) are deleted before the new records are inserted.
pub async fn upsert_embeddings(
&self,
collection_id: &str,
records: Vec<EmbeddingRecord>,
) -> Result<()> {
if records.is_empty() {
return Ok(());
}
let dimensions = records[0].embedding.len();
for (i, record) in records.iter().enumerate() {
if record.embedding.len() != dimensions {
return Err(BoosApiError::Provider(format!(
"Embedding dimension mismatch at record {}: got {} but expected {}",
i,
record.embedding.len(),
dimensions
)));
}
}
self.ensure_collection(collection_id, dimensions).await?;
let name = Self::collection_name(collection_id);
let doc_id = &records[0].document_id;
// Delete previous embeddings for this document.
self.delete_rows_by_document(&name, doc_id).await?;
// Determine the next safe rowid to avoid UNIQUE conflicts.
// We must check both _meta AND vec0 tables because orphan rows
// can exist in vec0 after a previous crash/panic mid-insert.
let meta_max = self
.db
.query_one(Statement::from_string(
DbBackend::Sqlite,
format!("SELECT COALESCE(MAX(rowid), 0) AS max_rid FROM {name}_meta"),
))
.await
.map_err(Self::wrap)?
.and_then(|r| r.try_get::<i64>("", "max_rid").ok())
.unwrap_or(0);
let vec_max = self
.db
.query_one(Statement::from_string(
DbBackend::Sqlite,
format!("SELECT COALESCE(MAX(rowid), 0) AS max_rid FROM {name}"),
))
.await
.ok()
.flatten()
.and_then(|r| r.try_get::<i64>("", "max_rid").ok())
.unwrap_or(0);
let mut next_rowid: i64 = meta_max.max(vec_max) + 1;
// Insert new records with explicit correlated rowids.
for record in &records {
let vec_json = Self::embedding_to_json(&record.embedding);
let rid = next_rowid;
next_rowid += 1;
// Insert embedding into vec0 with explicit rowid
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("INSERT INTO {name} (rowid, embedding) VALUES ($1, $2)"),
vec![rid.into(), vec_json.into()],
))
.await
.map_err(Self::wrap)?;
// Insert meta with the same rowid
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!(
"INSERT INTO {name}_meta (rowid, id, document_id, chunk_index, content) \
VALUES ($1, $2, $3, $4, $5)"
),
vec![
rid.into(),
record.id.clone().into(),
record.document_id.clone().into(),
record.chunk_index.into(),
record.content.clone().into(),
],
))
.await
.map_err(Self::wrap)?;
}
Ok(())
}
/// Add a single chunk to an existing collection.
/// Returns the generated chunk ID.
pub async fn add_single_chunk(
&self,
collection_id: &str,
document_id: &str,
content: &str,
embedding: &[f32],
) -> Result<String> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Err(BoosApiError::NotFound("Collection not found".into()));
}
self.ensure_collection(collection_id, embedding.len())
.await?;
// Determine next chunk_index for this document
let max_index = self
.db
.query_one(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("SELECT COALESCE(MAX(chunk_index), -1) AS max_idx FROM {meta_table} WHERE document_id = $1"),
vec![document_id.to_string().into()],
))
.await
.map_err(Self::wrap)?
.and_then(|r| r.try_get::<i32>("", "max_idx").ok())
.unwrap_or(-1);
let chunk_index = max_index + 1;
let chunk_id = format!("{}_{}", document_id, chunk_index);
// Determine next safe rowid
let meta_max = self
.db
.query_one(Statement::from_string(
DbBackend::Sqlite,
format!("SELECT COALESCE(MAX(rowid), 0) AS max_rid FROM {meta_table}"),
))
.await
.map_err(Self::wrap)?
.and_then(|r| r.try_get::<i64>("", "max_rid").ok())
.unwrap_or(0);
let vec_max = self
.db
.query_one(Statement::from_string(
DbBackend::Sqlite,
format!("SELECT COALESCE(MAX(rowid), 0) AS max_rid FROM {name}"),
))
.await
.ok()
.flatten()
.and_then(|r| r.try_get::<i64>("", "max_rid").ok())
.unwrap_or(0);
let rid: i64 = meta_max.max(vec_max) + 1;
let vec_json = Self::embedding_to_json(embedding);
// Insert embedding into vec0
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("INSERT INTO {name} (rowid, embedding) VALUES ($1, $2)"),
vec![rid.into(), vec_json.into()],
))
.await
.map_err(Self::wrap)?;
// Insert meta
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!(
"INSERT INTO {meta_table} (rowid, id, document_id, chunk_index, content) \
VALUES ($1, $2, $3, $4, $5)"
),
vec![
rid.into(),
chunk_id.clone().into(),
document_id.to_string().into(),
chunk_index.into(),
content.to_string().into(),
],
))
.await
.map_err(Self::wrap)?;
Ok(chunk_id)
}
/// Search for the most similar vectors in a knowledge base.
///
/// Returns up to `top_k` results ordered by ascending distance.
/// If the collection does not exist yet, an empty vec is returned.
pub async fn search(
&self,
knowledge_base_id: &str,
query_embedding: Vec<f32>,
top_k: usize,
) -> Result<Vec<VectorSearchResult>> {
let name = Self::collection_name(knowledge_base_id);
if !self.table_exists(&format!("{name}_meta")).await? {
tracing::warn!("Vector store: table {name}_meta does not exist, returning empty");
return Ok(vec![]);
}
let vec_json = Self::embedding_to_json(&query_embedding);
let sql = format!(
"SELECT m.id, m.document_id, m.chunk_index, m.content, v.distance \
FROM {name} v \
JOIN {name}_meta m ON m.rowid = v.rowid \
WHERE v.embedding MATCH $1 AND k = $2 \
ORDER BY v.distance"
);
let rows = match self
.db
.query_all(Statement::from_sql_and_values(
DbBackend::Sqlite,
&sql,
vec![vec_json.into(), (top_k as i64).into()],
))
.await
{
Ok(r) => r,
Err(e) => {
tracing::error!("Vector store: search query failed for {name}: {e}");
return Ok(vec![]);
}
};
let mut results = Vec::with_capacity(rows.len());
for row in &rows {
results.push(VectorSearchResult {
id: row.try_get("", "id").map_err(Self::wrap)?,
document_id: row.try_get("", "document_id").map_err(Self::wrap)?,
chunk_index: row.try_get("", "chunk_index").map_err(Self::wrap)?,
content: row.try_get("", "content").map_err(Self::wrap)?,
score: row
.try_get::<f64>("", "distance")
.map(|v| v as f32)
.map_err(Self::wrap)?,
rerank_score: None,
has_embedding: true,
});
}
Ok(results)
}
/// Delete all embeddings belonging to a specific document.
pub async fn delete_document_embeddings(
&self,
knowledge_base_id: &str,
document_id: &str,
) -> Result<()> {
let name = Self::collection_name(knowledge_base_id);
if !self.table_exists(&format!("{name}_meta")).await? {
return Ok(());
}
self.delete_rows_by_document(&name, document_id).await
}
/// Drop both tables for a knowledge base.
///
/// Silently succeeds if the tables do not exist.
pub async fn delete_collection(&self, knowledge_base_id: &str) -> Result<()> {
let name = Self::collection_name(knowledge_base_id);
let _ = self.exec(&format!("DROP TABLE IF EXISTS {name}")).await;
let _ = self
.exec(&format!("DROP TABLE IF EXISTS {name}_meta"))
.await;
Ok(())
}
/// Clear only the embedding vectors (vec0), keeping chunk metadata (_meta) intact.
/// This allows re-embedding without losing user edits or manually added chunks.
pub async fn clear_embeddings(&self, collection_id: &str) -> Result<()> {
let name = Self::collection_name(collection_id);
let _ = self.exec(&format!("DROP TABLE IF EXISTS {name}")).await;
Ok(())
}
/// List all chunk metadata with rowids for re-embedding.
/// Returns (rowid, chunk_id, content) tuples.
pub async fn list_all_chunks(&self, collection_id: &str) -> Result<Vec<(i64, String, String)>> {
self.list_chunks_raw(collection_id, None).await
}
/// List chunks (rowid, id, content) for a specific document.
pub async fn list_document_chunks_raw(
&self,
collection_id: &str,
document_id: &str,
) -> Result<Vec<(i64, String, String)>> {
self.list_chunks_raw(collection_id, Some(document_id)).await
}
/// Internal helper: list chunks with optional document_id filter.
async fn list_chunks_raw(
&self,
collection_id: &str,
document_id: Option<&str>,
) -> Result<Vec<(i64, String, String)>> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Ok(vec![]);
}
let rows = if let Some(doc_id) = document_id {
self.db
.query_all(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("SELECT rowid, id, content FROM \"{meta_table}\" WHERE document_id = $1 ORDER BY rowid"),
vec![doc_id.to_string().into()],
))
.await
.map_err(Self::wrap)?
} else {
self.db
.query_all(Statement::from_string(
DbBackend::Sqlite,
format!("SELECT rowid, id, content FROM {meta_table} ORDER BY rowid"),
))
.await
.map_err(Self::wrap)?
};
let mut result = Vec::new();
for row in &rows {
let rid: i64 = row.try_get("", "rowid").map_err(Self::wrap)?;
let id: String = row.try_get("", "id").map_err(Self::wrap)?;
let content: String = row.try_get("", "content").map_err(Self::wrap)?;
result.push((rid, id, content));
}
Ok(result)
}
/// Re-insert embeddings for existing chunks (used after clear_embeddings).
/// The vec0 table must already exist (or be recreated with correct dimensions).
pub async fn reinsert_embeddings(
&self,
collection_id: &str,
entries: Vec<(i64, Vec<f32>)>, // (rowid, embedding)
) -> Result<()> {
self.upsert_document_embeddings(collection_id, entries)
.await
}
/// Insert or replace embeddings for specific rowids.
/// Creates vec0 if needed, deletes existing rows, then inserts new embeddings.
pub async fn upsert_document_embeddings(
&self,
collection_id: &str,
entries: Vec<(i64, Vec<f32>)>,
) -> Result<()> {
if entries.is_empty() {
return Ok(());
}
let dimensions = entries[0].1.len();
let name = Self::collection_name(collection_id);
self.ensure_collection(collection_id, dimensions).await?;
for (rid, embedding) in &entries {
// Delete existing row if present (ignore errors — may not exist)
let _ = self
.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {name} WHERE rowid = $1"),
vec![(*rid).into()],
))
.await;
let vec_json = Self::embedding_to_json(embedding);
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("INSERT INTO {name} (rowid, embedding) VALUES ($1, $2)"),
vec![(*rid).into(), vec_json.into()],
))
.await
.map_err(Self::wrap)?;
}
Ok(())
}
/// Delete a single chunk by its id from both vec0 and metadata tables.
pub async fn delete_chunk(&self, collection_id: &str, chunk_id: &str) -> Result<()> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Ok(());
}
// Get the rowid from _meta
let row = self
.db
.query_one(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("SELECT rowid FROM {meta_table} WHERE id = $1"),
vec![chunk_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
if let Some(row) = row {
let rid: i64 = row.try_get("", "rowid").map_err(Self::wrap)?;
// Delete from vec0
let _ = self
.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {name} WHERE rowid = $1"),
vec![rid.into()],
))
.await;
// Delete from _meta
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {meta_table} WHERE id = $1"),
vec![chunk_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
}
Ok(())
}
/// Update the text content of a single chunk in the metadata table.
pub async fn update_chunk_content(
&self,
collection_id: &str,
chunk_id: &str,
new_content: &str,
) -> Result<()> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Err(BoosApiError::NotFound("Collection not found".into()));
}
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("UPDATE {meta_table} SET content = $1 WHERE id = $2"),
vec![new_content.to_string().into(), chunk_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
Ok(())
}
/// Update the embedding vector for a single chunk identified by its chunk id.
pub async fn update_chunk_embedding(
&self,
collection_id: &str,
chunk_id: &str,
embedding: &[f32],
) -> Result<()> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Err(BoosApiError::NotFound("Collection not found".into()));
}
// Get the rowid from _meta
let row = self
.db
.query_one(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("SELECT rowid FROM {meta_table} WHERE id = $1"),
vec![chunk_id.to_string().into()],
))
.await
.map_err(Self::wrap)?
.ok_or_else(|| BoosApiError::NotFound(format!("Chunk {} not found", chunk_id)))?;
let rid: i64 = row.try_get("", "rowid").map_err(Self::wrap)?;
let vec_json = Self::embedding_to_json(embedding);
self.ensure_collection(collection_id, embedding.len())
.await?;
let _ = self
.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {name} WHERE rowid = $1"),
vec![rid.into()],
))
.await;
// Insert embedding into vec0. DELETE+INSERT also restores rows after
// clearing a collection while preserving chunk metadata.
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("INSERT INTO {name} (rowid, embedding) VALUES ($1, $2)"),
vec![rid.into(), vec_json.into()],
))
.await
.map_err(Self::wrap)?;
Ok(())
}
// ── private helpers ─────────────────────────────────────────────────
/// Delete rows from both vec0 and metadata tables by `document_id`.
async fn delete_rows_by_document(&self, table_name: &str, document_id: &str) -> Result<()> {
let rows = self
.db
.query_all(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("SELECT rowid FROM {table_name}_meta WHERE document_id = $1"),
vec![document_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
for row in &rows {
let rid: i64 = row.try_get("", "rowid").map_err(Self::wrap)?;
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {table_name} WHERE rowid = $1"),
vec![rid.into()],
))
.await
.map_err(Self::wrap)?;
}
self.db
.execute(Statement::from_sql_and_values(
DbBackend::Sqlite,
&format!("DELETE FROM {table_name}_meta WHERE document_id = $1"),
vec![document_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
Ok(())
}
/// Convert an embedding vector to a JSON array string for sqlite-vec.
fn embedding_to_json(embedding: &[f32]) -> String {
format!(
"[{}]",
embedding
.iter()
.map(|v| v.to_string())
.collect::<Vec<_>>()
.join(",")
)
}
/// Check whether a regular table exists in the database.
async fn table_exists(&self, table_name: &str) -> Result<bool> {
let row = self
.db
.query_one(Statement::from_sql_and_values(
DbBackend::Sqlite,
"SELECT name FROM sqlite_master WHERE type='table' AND name=$1",
vec![table_name.to_string().into()],
))
.await
.map_err(Self::wrap)?;
Ok(row.is_some())
}
async fn collection_dimensions(&self, table_name: &str) -> Result<Option<usize>> {
let row = self
.db
.query_one(Statement::from_sql_and_values(
DbBackend::Sqlite,
"SELECT sql FROM sqlite_master WHERE type='table' AND name=$1",
vec![table_name.to_string().into()],
))
.await
.map_err(Self::wrap)?;
let Some(row) = row else {
return Ok(None);
};
let sql: String = row.try_get("", "sql").map_err(Self::wrap)?;
Ok(parse_vec_dimension(&sql))
}
/// List all chunks stored for a specific document within a collection.
pub async fn list_document_chunks(
&self,
collection_id: &str,
document_id: &str,
) -> Result<Vec<VectorSearchResult>> {
let name = Self::collection_name(collection_id);
let meta_table = format!("{name}_meta");
if !self.table_exists(&meta_table).await? {
return Ok(vec![]);
}
let vec_exists = self.table_exists(&name).await?;
let sql = if vec_exists {
format!(
"SELECT m.id, m.document_id, m.chunk_index, m.content, \
CASE WHEN v.rowid IS NOT NULL THEN 1 ELSE 0 END AS has_embedding \
FROM \"{meta_table}\" m \
LEFT JOIN \"{name}\" v ON m.rowid = v.rowid \
WHERE m.document_id = $1 ORDER BY m.chunk_index"
)
} else {
format!(
"SELECT id, document_id, chunk_index, content, 0 AS has_embedding \
FROM \"{meta_table}\" WHERE document_id = $1 ORDER BY chunk_index"
)
};
let rows = self
.db
.query_all(Statement::from_sql_and_values(
DbBackend::Sqlite,
&sql,
vec![document_id.to_string().into()],
))
.await
.map_err(Self::wrap)?;
let mut results = Vec::with_capacity(rows.len());
for row in &rows {
let has_emb: i32 = row.try_get("", "has_embedding").unwrap_or(0);
results.push(VectorSearchResult {
id: row.try_get("", "id").map_err(Self::wrap)?,
document_id: row.try_get("", "document_id").map_err(Self::wrap)?,
chunk_index: row.try_get("", "chunk_index").map_err(Self::wrap)?,
content: row.try_get("", "content").map_err(Self::wrap)?,
score: 0.0,
rerank_score: None,
has_embedding: has_emb != 0,
});
}
Ok(results)
}
/// Shorthand for executing a statement with no parameters.
async fn exec(&self, sql: &str) -> Result<()> {
self.db
.execute(Statement::from_string(DbBackend::Sqlite, sql))
.await
.map_err(Self::wrap)?;
Ok(())
}
fn wrap(e: DbErr) -> BoosApiError {
BoosApiError::Provider(format!("Vector store error: {e}"))
}
}
fn parse_vec_dimension(sql: &str) -> Option<usize> {
let marker = "float[";
let start = sql.find(marker)? + marker.len();
let end = sql[start..].find(']')? + start;
sql[start..end].parse().ok()
}
#[cfg(test)]
mod tests {
use super::*;
use sea_orm::{ConnectOptions, Database};
async fn test_store() -> VectorStore {
register_sqlite_vec_extension();
let mut opts = ConnectOptions::new("sqlite::memory:");
opts.sqlx_logging(false);
let db = Database::connect(opts).await.unwrap();
VectorStore::new(db)
}
fn record(id: &str, document_id: &str, dimensions: usize) -> EmbeddingRecord {
EmbeddingRecord {
id: id.to_string(),
document_id: document_id.to_string(),
chunk_index: 0,
content: "content".to_string(),
embedding: vec![0.0; dimensions],
}
}
#[tokio::test]
async fn clear_embeddings_allows_reinsert_with_new_dimensions() {
let store = test_store().await;
let collection_id = "kb_dimension_change";
store
.upsert_embeddings(collection_id, vec![record("chunk-1", "doc-1", 3072)])
.await
.unwrap();
store.clear_embeddings(collection_id).await.unwrap();
store
.upsert_document_embeddings(collection_id, vec![(1, vec![0.0; 4096])])
.await
.unwrap();
}
#[tokio::test]
async fn existing_collection_allows_reinsert_with_new_dimensions() {
let store = test_store().await;
let collection_id = "kb_existing_dimension_change";
store
.upsert_embeddings(collection_id, vec![record("chunk-1", "doc-1", 3072)])
.await
.unwrap();
store
.upsert_document_embeddings(collection_id, vec![(1, vec![0.0; 4096])])
.await
.unwrap();
}
#[test]
fn parse_vec_dimension_reads_float_size() {
assert_eq!(
parse_vec_dimension("CREATE VIRTUAL TABLE vec_kb USING vec0(embedding float[4096])"),
Some(4096)
);
assert_eq!(parse_vec_dimension("CREATE TABLE vec_kb(id TEXT)"), None);
}
}
+728
View File
@@ -0,0 +1,728 @@
use reqwest::{Client, Method, StatusCode};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::future::Future;
use std::io::{Read, Write};
use std::path::Path;
use zip::write::SimpleFileOptions;
use crate::error::{BoosApiError, Result};
// === Types ===
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WebDavConfig {
pub host: String,
pub username: String,
pub password: String,
pub path: String,
pub accept_invalid_certs: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WebDavFileInfo {
pub file_name: String,
pub size: i64,
pub last_modified: String,
pub hostname: String,
}
pub struct BackupZipContents {
pub db_path: std::path::PathBuf,
pub metadata: serde_json::Value,
pub has_documents: bool,
pub has_workspace: bool,
pub master_key_path: Option<std::path::PathBuf>,
}
// === WebDAV Client ===
pub struct WebDavClient {
client: Client,
config: WebDavConfig,
}
impl WebDavClient {
pub fn new(config: WebDavConfig) -> Result<Self> {
let client = Client::builder()
.danger_accept_invalid_certs(config.accept_invalid_certs)
.timeout(std::time::Duration::from_secs(300))
.build()
.map_err(|e| BoosApiError::Gateway(format!("Failed to create HTTP client: {}", e)))?;
Ok(Self { client, config })
}
fn base_url(&self) -> String {
let host = self.config.host.trim_end_matches('/');
let path = self.config.path.trim_matches('/');
if path.is_empty() {
format!("{}/", host)
} else {
format!("{}/{}/", host, path)
}
}
fn file_url(&self, filename: &str) -> String {
format!("{}{}", self.base_url(), filename)
}
/// Check connection and auto-create remote directory if missing.
pub async fn check_connection(&self) -> Result<bool> {
let url = self.base_url();
let method = Method::from_bytes(b"PROPFIND")
.map_err(|e| BoosApiError::Gateway(format!("Invalid method: {}", e)))?;
let response = self
.client
.request(method, &url)
.basic_auth(&self.config.username, Some(&self.config.password))
.header("Depth", "0")
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV connection failed: {}", e)))?;
match response.status() {
StatusCode::MULTI_STATUS | StatusCode::OK => Ok(true),
StatusCode::NOT_FOUND => {
self.mkdir().await?;
Ok(true)
}
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => Err(BoosApiError::Gateway(
"WebDAV authentication failed".to_string(),
)),
status => Err(BoosApiError::Gateway(format!(
"WebDAV error: HTTP {}",
status
))),
}
}
/// Create the remote directory tree.
async fn mkdir(&self) -> Result<()> {
let host = self.config.host.trim_end_matches('/');
let path = self.config.path.trim_matches('/');
if path.is_empty() {
return Ok(());
}
let parts: Vec<&str> = path.split('/').filter(|p| !p.is_empty()).collect();
let mut current = String::new();
for part in parts {
current = if current.is_empty() {
part.to_string()
} else {
format!("{}/{}", current, part)
};
let url = format!("{}/{}/", host, current);
let method = Method::from_bytes(b"MKCOL")
.map_err(|e| BoosApiError::Gateway(format!("Invalid method: {}", e)))?;
let response = self
.client
.request(method, &url)
.basic_auth(&self.config.username, Some(&self.config.password))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV MKCOL failed: {}", e)))?;
// CREATED=success, METHOD_NOT_ALLOWED=already exists
match response.status() {
StatusCode::CREATED | StatusCode::OK | StatusCode::METHOD_NOT_ALLOWED => {}
status => {
return Err(BoosApiError::Gateway(format!(
"WebDAV mkdir failed for '{}': HTTP {}",
current, status
)));
}
}
}
Ok(())
}
/// List boosapi backup .zip files in the remote directory.
pub async fn list_files(&self) -> Result<Vec<WebDavFileInfo>> {
run_after_directory_ready(
|| self.check_connection(),
|| async {
let url = self.base_url();
let method = Method::from_bytes(b"PROPFIND")
.map_err(|e| BoosApiError::Gateway(format!("Invalid method: {}", e)))?;
let body = r#"<?xml version="1.0" encoding="utf-8"?>
<D:propfind xmlns:D="DAV:">
<D:prop>
<D:getcontentlength/>
<D:getlastmodified/>
<D:resourcetype/>
</D:prop>
</D:propfind>"#;
let response = self
.client
.request(method, &url)
.basic_auth(&self.config.username, Some(&self.config.password))
.header("Depth", "1")
.header("Content-Type", "application/xml; charset=utf-8")
.body(body)
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV PROPFIND failed: {}", e)))?;
if response.status() != StatusCode::MULTI_STATUS && !response.status().is_success()
{
return Err(BoosApiError::Gateway(format!(
"WebDAV list failed: HTTP {}",
response.status()
)));
}
let text = response
.text()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read response: {}", e)))?;
parse_propfind_response(&text)
},
)
.await
}
/// Upload a local file to the remote directory.
pub async fn upload_file(&self, filename: &str, local_path: &Path) -> Result<()> {
run_after_directory_ready(
|| self.check_connection(),
|| async {
let data = std::fs::read(local_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read file: {}", e)))?;
let url = self.file_url(filename);
let response = self
.client
.put(&url)
.basic_auth(&self.config.username, Some(&self.config.password))
.header("Content-Type", "application/octet-stream")
.body(data)
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV upload failed: {}", e)))?;
match response.status() {
StatusCode::CREATED | StatusCode::OK | StatusCode::NO_CONTENT => Ok(()),
status => Err(BoosApiError::Gateway(format!(
"WebDAV upload failed: HTTP {}",
status
))),
}
},
)
.await
}
/// Download a file from the remote directory to a local path.
pub async fn download_file(&self, filename: &str, local_path: &Path) -> Result<()> {
let url = self.file_url(filename);
let response = self
.client
.get(&url)
.basic_auth(&self.config.username, Some(&self.config.password))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV download failed: {}", e)))?;
if !response.status().is_success() {
return Err(BoosApiError::Gateway(format!(
"WebDAV download failed: HTTP {}",
response.status()
)));
}
let data = response
.bytes()
.await
.map_err(|e| BoosApiError::Gateway(format!("Failed to read download: {}", e)))?;
std::fs::write(local_path, &data)
.map_err(|e| BoosApiError::Gateway(format!("Failed to write file: {}", e)))?;
Ok(())
}
/// Delete a file from the remote directory.
pub async fn delete_file(&self, filename: &str) -> Result<()> {
let url = self.file_url(filename);
let response = self
.client
.delete(&url)
.basic_auth(&self.config.username, Some(&self.config.password))
.send()
.await
.map_err(|e| BoosApiError::Gateway(format!("WebDAV delete failed: {}", e)))?;
match response.status() {
StatusCode::OK | StatusCode::NO_CONTENT | StatusCode::NOT_FOUND => Ok(()),
status => Err(BoosApiError::Gateway(format!(
"WebDAV delete failed: HTTP {}",
status
))),
}
}
}
// === ZIP Backup Utilities ===
/// Create a backup ZIP containing the database snapshot and metadata.
pub fn create_backup_zip(
db_path: &Path,
documents_dir: Option<&Path>,
workspace_dir: Option<&Path>,
master_key_path: Option<&Path>,
dest_zip: &Path,
app_version: &str,
object_counts_json: &str,
) -> Result<()> {
let file = std::fs::File::create(dest_zip)
.map_err(|e| BoosApiError::Gateway(format!("Failed to create ZIP file: {}", e)))?;
let mut zip = zip::ZipWriter::new(file);
let options = SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated);
// boosapi.db
let db_data = std::fs::read(db_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read database: {}", e)))?;
let db_checksum = format!("{:x}", Sha256::digest(&db_data));
zip.start_file("boosapi.db", options)
.map_err(|e| BoosApiError::Gateway(format!("ZIP error: {}", e)))?;
zip.write_all(&db_data)
.map_err(|e| BoosApiError::Gateway(format!("ZIP write error: {}", e)))?;
// metadata.json
let metadata = serde_json::json!({
"version": 1,
"app_version": app_version,
"created_at": chrono::Utc::now().to_rfc3339(),
"platform": std::env::consts::OS,
"arch": std::env::consts::ARCH,
"hostname": get_hostname(),
"db_checksum": db_checksum,
"include_documents": documents_dir.is_some(),
"include_workspace": workspace_dir.is_some(),
"object_counts": object_counts_json,
});
let metadata_json = serde_json::to_string_pretty(&metadata)
.map_err(|e| BoosApiError::Gateway(format!("JSON error: {}", e)))?;
zip.start_file("metadata.json", SimpleFileOptions::default())
.map_err(|e| BoosApiError::Gateway(format!("ZIP error: {}", e)))?;
zip.write_all(metadata_json.as_bytes())
.map_err(|e| BoosApiError::Gateway(format!("ZIP write error: {}", e)))?;
// Optional: master.key for cross-device restore
if let Some(key_path) = master_key_path {
if key_path.exists() {
let key_data = std::fs::read(key_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read master.key: {}", e)))?;
zip.start_file("master.key", options)
.map_err(|e| BoosApiError::Gateway(format!("ZIP error: {}", e)))?;
zip.write_all(&key_data)
.map_err(|e| BoosApiError::Gateway(format!("ZIP write error: {}", e)))?;
}
}
// Optional: documents/ directory
if let Some(docs_dir) = documents_dir {
if docs_dir.exists() {
add_directory_to_zip(&mut zip, docs_dir, "documents", options)?;
}
}
// Optional: workspace/ directory
if let Some(ws_dir) = workspace_dir {
if ws_dir.exists() {
add_directory_to_zip(&mut zip, ws_dir, "workspace", options)?;
}
}
zip.finish()
.map_err(|e| BoosApiError::Gateway(format!("ZIP finalize error: {}", e)))?;
Ok(())
}
/// Extract a backup ZIP and return its contents.
pub fn extract_backup_zip(zip_path: &Path, dest_dir: &Path) -> Result<BackupZipContents> {
let file = std::fs::File::open(zip_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to open ZIP: {}", e)))?;
let mut archive = zip::ZipArchive::new(file)
.map_err(|e| BoosApiError::Gateway(format!("Invalid ZIP file: {}", e)))?;
std::fs::create_dir_all(dest_dir)
.map_err(|e| BoosApiError::Gateway(format!("Failed to create temp dir: {}", e)))?;
let mut db_path = None;
let mut metadata = None;
let mut has_documents = false;
let mut has_workspace = false;
let mut master_key_path = None;
for i in 0..archive.len() {
let mut entry = archive
.by_index(i)
.map_err(|e| BoosApiError::Gateway(format!("ZIP read error: {}", e)))?;
let name = entry.name().to_string();
if name.contains("..") {
continue; // path traversal protection
}
if name == "boosapi.db" {
let path = dest_dir.join("boosapi.db");
let mut outfile = std::fs::File::create(&path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract db: {}", e)))?;
std::io::copy(&mut entry, &mut outfile)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract db: {}", e)))?;
db_path = Some(path);
} else if name == "metadata.json" {
let mut contents = String::new();
entry
.read_to_string(&mut contents)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read metadata: {}", e)))?;
metadata = Some(
serde_json::from_str::<serde_json::Value>(&contents)
.map_err(|e| BoosApiError::Gateway(format!("Invalid metadata JSON: {}", e)))?,
);
} else if name == "master.key" {
let path = dest_dir.join("master.key");
let mut outfile = std::fs::File::create(&path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract master.key: {}", e)))?;
std::io::copy(&mut entry, &mut outfile)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract master.key: {}", e)))?;
master_key_path = Some(path);
} else if name.starts_with("documents/") && !entry.is_dir() {
has_documents = true;
let path = dest_dir.join(&name);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).ok();
}
let mut outfile = std::fs::File::create(&path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract file: {}", e)))?;
std::io::copy(&mut entry, &mut outfile)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract file: {}", e)))?;
} else if name.starts_with("workspace/") && !entry.is_dir() {
has_workspace = true;
let path = dest_dir.join(&name);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).ok();
}
let mut outfile = std::fs::File::create(&path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract file: {}", e)))?;
std::io::copy(&mut entry, &mut outfile)
.map_err(|e| BoosApiError::Gateway(format!("Failed to extract file: {}", e)))?;
}
}
Ok(BackupZipContents {
db_path: db_path.ok_or_else(|| BoosApiError::Gateway("No boosapi.db in backup ZIP".into()))?,
metadata: metadata
.ok_or_else(|| BoosApiError::Gateway("No metadata.json in backup ZIP".into()))?,
has_documents,
has_workspace,
master_key_path,
})
}
/// Verify the checksum of an extracted database against metadata.
pub fn verify_db_checksum(db_path: &Path, expected_checksum: &str) -> Result<bool> {
let data = std::fs::read(db_path)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read db for checksum: {}", e)))?;
let actual = format!("{:x}", Sha256::digest(&data));
Ok(actual == expected_checksum)
}
/// Generate a backup filename with timestamp and hostname.
pub fn generate_backup_filename() -> String {
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S");
let hostname = get_hostname();
format!("boosapi-backup-{}.{}.zip", timestamp, hostname)
}
/// Parse hostname from a backup filename.
pub fn parse_hostname_from_filename(filename: &str) -> String {
// Format: boosapi-backup-YYYYMMDD_HHMMSS.hostname.zip
let name = filename.trim_end_matches(".zip");
if let Some(rest) = name.strip_prefix("boosapi-backup-") {
// rest = "YYYYMMDD_HHMMSS.hostname"
if let Some(dot_pos) = rest.find('.') {
return rest[dot_pos + 1..].to_string();
}
}
"unknown".to_string()
}
pub fn documents_sync_root() -> std::path::PathBuf {
crate::storage_paths::documents_root()
}
pub fn sync_status_timestamp() -> String {
chrono::Utc::now().to_rfc3339()
}
async fn run_after_directory_ready<T, Check, CheckFut, Action, ActionFut>(
check: Check,
action: Action,
) -> Result<T>
where
Check: FnOnce() -> CheckFut,
CheckFut: Future<Output = Result<bool>>,
Action: FnOnce() -> ActionFut,
ActionFut: Future<Output = Result<T>>,
{
check().await?;
action().await
}
// === Internal Helpers ===
fn get_hostname() -> String {
std::process::Command::new("hostname")
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "unknown".to_string())
}
fn add_directory_to_zip<W: Write + std::io::Seek>(
zip: &mut zip::ZipWriter<W>,
dir: &Path,
prefix: &str,
options: SimpleFileOptions,
) -> Result<()> {
let mut files = Vec::new();
collect_files(dir, &mut files)?;
for file_path in files {
let rel = file_path
.strip_prefix(dir)
.map_err(|e| BoosApiError::Gateway(format!("Path error: {}", e)))?;
let zip_path = format!("{}/{}", prefix, rel.to_string_lossy());
zip.start_file(&zip_path, options)
.map_err(|e| BoosApiError::Gateway(format!("ZIP error: {}", e)))?;
let data = std::fs::read(&file_path)
.map_err(|e| BoosApiError::Gateway(format!("Read error: {}", e)))?;
zip.write_all(&data)
.map_err(|e| BoosApiError::Gateway(format!("ZIP write error: {}", e)))?;
}
Ok(())
}
fn collect_files(dir: &Path, files: &mut Vec<std::path::PathBuf>) -> Result<()> {
if !dir.is_dir() {
return Ok(());
}
for entry in std::fs::read_dir(dir)
.map_err(|e| BoosApiError::Gateway(format!("Failed to read directory: {}", e)))?
{
let entry = entry.map_err(|e| BoosApiError::Gateway(format!("Dir entry error: {}", e)))?;
let path = entry.path();
if path.is_dir() {
collect_files(&path, files)?;
} else {
files.push(path);
}
}
Ok(())
}
/// Parse WebDAV PROPFIND XML response to extract file information.
fn parse_propfind_response(xml: &str) -> Result<Vec<WebDavFileInfo>> {
let mut files = Vec::new();
let response_blocks = split_xml_responses(xml);
for block in response_blocks {
let lower_block = block.to_lowercase();
// Skip collections (directories)
if lower_block.contains("<d:collection") || lower_block.contains("<collection") {
continue;
}
let href = extract_xml_value(&block, "href").unwrap_or_default();
if href.is_empty() || href.ends_with('/') {
continue;
}
let file_name = url_decode(href.split('/').last().unwrap_or(""));
if file_name.is_empty() || !file_name.ends_with(".zip") {
continue;
}
if !file_name.starts_with("boosapi-backup-") {
continue;
}
let size: i64 = extract_xml_value(&block, "getcontentlength")
.and_then(|s| s.parse().ok())
.unwrap_or(0);
let last_modified = extract_xml_value(&block, "getlastmodified").unwrap_or_default();
let hostname = parse_hostname_from_filename(&file_name);
files.push(WebDavFileInfo {
file_name,
size,
last_modified,
hostname,
});
}
// Newest first (filenames contain timestamps)
files.sort_by(|a, b| b.file_name.cmp(&a.file_name));
Ok(files)
}
fn split_xml_responses(xml: &str) -> Vec<String> {
let mut blocks = Vec::new();
let lower = xml.to_lowercase();
let tag_patterns = ["d:response", "response"];
for tag in &tag_patterns {
let open1 = format!("<{}>", tag);
let open2 = format!("<{} ", tag);
let close = format!("</{}>", tag);
let mut pos = 0;
while pos < lower.len() {
let start = lower[pos..]
.find(&open1)
.or_else(|| lower[pos..].find(&open2));
if let Some(s) = start {
let abs_start = pos + s;
if let Some(end) = lower[abs_start..].find(&close) {
let abs_end = abs_start + end + close.len();
blocks.push(xml[abs_start..abs_end].to_string());
pos = abs_end;
} else {
break;
}
} else {
break;
}
}
if !blocks.is_empty() {
break;
}
}
blocks
}
fn extract_xml_value(xml: &str, tag_local_name: &str) -> Option<String> {
let lower = xml.to_lowercase();
let tag = tag_local_name.to_lowercase();
let patterns = [
(format!("<d:{}>", tag), format!("</d:{}>", tag)),
(format!("<{}>", tag), format!("</{}>", tag)),
];
for (open, close) in &patterns {
if let Some(start) = lower.find(open) {
let content_start = start + open.len();
if let Some(end) = lower[content_start..].find(close) {
return Some(xml[content_start..content_start + end].trim().to_string());
}
}
}
None
}
fn url_decode(s: &str) -> String {
let mut result = String::with_capacity(s.len());
let bytes = s.as_bytes();
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%' && i + 2 < bytes.len() {
let h1 = (bytes[i + 1] as char).to_digit(16);
let h2 = (bytes[i + 2] as char).to_digit(16);
if let (Some(h1), Some(h2)) = (h1, h2) {
result.push((h1 * 16 + h2) as u8 as char);
i += 3;
continue;
}
}
result.push(bytes[i] as char);
i += 1;
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
#[tokio::test]
async fn run_after_directory_ready_checks_before_action() {
let events = Arc::new(Mutex::new(Vec::new()));
let check_events = events.clone();
let action_events = events.clone();
let result = run_after_directory_ready(
move || async move {
check_events.lock().unwrap().push("check");
Ok(true)
},
move || async move {
action_events.lock().unwrap().push("action");
Ok::<_, BoosApiError>("done")
},
)
.await;
assert!(matches!(result, Ok("done")));
assert_eq!(*events.lock().unwrap(), vec!["check", "action"]);
}
#[tokio::test]
async fn run_after_directory_ready_skips_action_when_check_fails() {
let events = Arc::new(Mutex::new(Vec::new()));
let check_events = events.clone();
let action_events = events.clone();
let result: Result<&'static str> = run_after_directory_ready(
move || async move {
check_events.lock().unwrap().push("check");
Err(BoosApiError::Gateway("probe failed".into()))
},
move || async move {
action_events.lock().unwrap().push("action");
Ok("done")
},
)
.await;
assert!(result.is_err(), "check failures must stop the action");
assert_eq!(*events.lock().unwrap(), vec!["check"]);
}
#[test]
fn documents_sync_root_matches_documents_root() {
assert_eq!(
documents_sync_root(),
crate::storage_paths::documents_root()
);
}
#[test]
fn sync_status_timestamp_is_rfc3339() {
let timestamp = sync_status_timestamp();
assert!(
chrono::DateTime::parse_from_rfc3339(&timestamp).is_ok(),
"sync status timestamps should be RFC3339 so the frontend can render them directly, got: {timestamp}"
);
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,298 @@
use std::path::{Path, PathBuf};
use boosapi_core::repo::codex_session_visibility::repair_codex_session_visibility_at;
use sea_orm::{ConnectionTrait, Database, DbBackend, Statement};
use serde_json::{json, Value};
fn codex_home_root() -> PathBuf {
std::env::temp_dir().join(format!(
"boosapi-codex-session-visibility-{}-{}",
std::process::id(),
uuid::Uuid::new_v4()
))
}
fn write_text(path: &Path, content: &str) {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).expect("create parent dir");
}
std::fs::write(path, content).expect("write text fixture");
}
fn write_json(path: &Path, value: &Value) {
let content = serde_json::to_string_pretty(value).expect("serialize json fixture");
write_text(path, &content);
}
fn rollout_path(codex_home: &Path, dir: &str, day: &str, thread: &str) -> PathBuf {
codex_home
.join(dir)
.join("2026")
.join("06")
.join(day)
.join(format!("rollout-{thread}.jsonl"))
}
fn write_rollout(path: &Path, thread_id: &str, provider: &str, cwd: &str, include_user: bool) {
let timestamp = "2026-06-09T00:00:00.000Z";
let mut lines = vec![json!({
"timestamp": timestamp, "type": "session_meta",
"payload": { "id": thread_id, "timestamp": timestamp, "cwd": cwd,
"source": "cli", "cli_version": "0.121.0", "model_provider": provider }
})
.to_string()];
if include_user {
lines.push(
json!({
"timestamp": timestamp, "type": "event_msg",
"payload": { "type": "user_message", "message": "hello" }
})
.to_string(),
);
}
lines.push(
json!({
"timestamp": timestamp, "type": "event_msg",
"payload": { "type": "assistant_message", "encrypted_content": "ciphertext" }
})
.to_string(),
);
write_text(path, &format!("{}\n", lines.join("\n")));
}
fn first_line_json(path: &Path) -> Value {
let content = std::fs::read_to_string(path).expect("read rollout");
let first_line = content.lines().next().expect("rollout first line");
serde_json::from_str(first_line).expect("parse rollout first line")
}
async fn create_state_db(codex_home: &Path) {
let db_path = codex_home.join("state_5.sqlite");
let db = Database::connect(format!("sqlite://{}?mode=rwc", db_path.display()))
.await
.expect("connect state db");
db.execute(Statement::from_string(
DbBackend::Sqlite,
r#"CREATE TABLE threads (
id TEXT PRIMARY KEY,
model_provider TEXT,
cwd TEXT NOT NULL DEFAULT '',
archived INTEGER NOT NULL DEFAULT 0,
has_user_event INTEGER NOT NULL DEFAULT 0,
first_user_message TEXT NOT NULL DEFAULT '',
updated_at_ms INTEGER NOT NULL DEFAULT 0
)"#
.to_string(),
))
.await
.expect("create threads");
db.execute(Statement::from_string(
DbBackend::Sqlite,
r#"INSERT INTO threads
(id, model_provider, cwd, archived, has_user_event, first_user_message, updated_at_ms)
VALUES
('thread-a', 'openai', '', 0, 0, 'hello', 123),
('thread-b', 'old-provider', '', 1, 0, 'archived', 456)"#
.to_string(),
))
.await
.expect("insert threads");
}
async fn read_thread_rows(codex_home: &Path) -> Vec<(String, String, String, i64, i64)> {
let db_path = codex_home.join("state_5.sqlite");
let db = Database::connect(format!("sqlite://{}?mode=ro", db_path.display()))
.await
.expect("connect state db read");
let rows = db
.query_all(Statement::from_string(
DbBackend::Sqlite,
"SELECT id, model_provider, cwd, has_user_event, updated_at_ms FROM threads ORDER BY id"
.to_string(),
))
.await
.expect("query threads");
rows.into_iter()
.map(|row| {
(
row.try_get::<String>("", "id").expect("id"),
row.try_get::<String>("", "model_provider")
.expect("model_provider"),
row.try_get::<String>("", "cwd").expect("cwd"),
row.try_get::<i64>("", "has_user_event")
.expect("has_user_event"),
row.try_get::<i64>("", "updated_at_ms")
.expect("updated_at_ms"),
)
})
.collect()
}
#[tokio::test]
async fn codex_session_visibility_syncs_rollouts_sqlite_and_global_state() {
let codex_home = codex_home_root();
let active_rollout = rollout_path(&codex_home, "sessions", "09", "thread-a");
let archived_rollout = rollout_path(&codex_home, "archived_sessions", "08", "thread-b");
write_text(
&codex_home.join("config.toml"),
r#"model_provider = "any"
[model_providers.any]
base_url = "http://localhost:1234/v1"
"#,
);
write_json(
&codex_home.join("auth.json"),
&json!({ "OPENAI_API_KEY": "keep-me" }),
);
write_rollout(
&active_rollout,
"thread-a",
"openai",
"/tmp/project-alpha",
true,
);
write_rollout(
&archived_rollout,
"thread-b",
"old-provider",
"/tmp/project-alpha",
false,
);
write_json(
&codex_home.join(".codex-global-state.json"),
&json!({
"electron-saved-workspace-roots": ["/tmp/project-alpha/"],
"project-order": [],
"active-workspace-roots": "/tmp/project-alpha/"
}),
);
create_state_db(&codex_home).await;
let result = repair_codex_session_visibility_at(&codex_home)
.await
.expect("repair succeeds");
assert_eq!(result.target_provider, "any");
assert_eq!(result.changed_session_files, 2);
assert_eq!(result.skipped_locked_session_files, 0);
assert_eq!(result.sqlite_provider_rows_updated, 2);
assert_eq!(result.sqlite_user_event_rows_updated, 1);
assert_eq!(result.sqlite_cwd_rows_updated, 2);
assert_eq!(result.updated_workspace_roots, 1);
assert!(result
.backup_dir
.as_deref()
.expect("backup dir")
.contains("boosapi-session-visibility"));
assert!(result
.encrypted_content_warning
.as_deref()
.expect("encrypted content warning")
.contains("encrypted_content"));
assert_eq!(
first_line_json(&active_rollout)["payload"]["model_provider"],
"any"
);
assert_eq!(
first_line_json(&archived_rollout)["payload"]["model_provider"],
"any"
);
let active_content = std::fs::read_to_string(&active_rollout).expect("read active rollout");
assert!(active_content.contains("encrypted_content"));
let rows = read_thread_rows(&codex_home).await;
assert_eq!(
rows,
vec![
(
"thread-a".into(),
"any".into(),
"/tmp/project-alpha".into(),
1,
123
),
(
"thread-b".into(),
"any".into(),
"/tmp/project-alpha".into(),
0,
456
),
]
);
assert_eq!(
std::fs::read_to_string(codex_home.join("auth.json")).expect("read auth"),
"{\n \"OPENAI_API_KEY\": \"keep-me\"\n}"
);
let global_state = serde_json::from_str::<Value>(
&std::fs::read_to_string(codex_home.join(".codex-global-state.json"))
.expect("read global state"),
)
.expect("parse global state");
assert_eq!(
global_state["electron-saved-workspace-roots"],
json!(["/tmp/project-alpha"])
);
assert!(codex_home.join(".codex-global-state.json.bak").exists());
let _ = std::fs::remove_dir_all(codex_home);
}
#[tokio::test]
async fn codex_session_visibility_succeeds_when_sqlite_is_missing() {
let codex_home = codex_home_root();
let rollout = rollout_path(&codex_home, "sessions", "09", "thread-a");
write_text(
&codex_home.join("config.toml"),
"sandbox_mode = \"workspace-write\"\n",
);
write_rollout(
&rollout,
"thread-a",
"old-provider",
"/tmp/project-beta",
true,
);
let result = repair_codex_session_visibility_at(&codex_home)
.await
.expect("repair succeeds without sqlite");
assert_eq!(result.target_provider, "openai");
assert_eq!(result.changed_session_files, 1);
assert!(!result.sqlite_present);
assert_eq!(
first_line_json(&rollout)["payload"]["model_provider"],
"openai"
);
let _ = std::fs::remove_dir_all(codex_home);
}
#[tokio::test]
async fn codex_session_visibility_does_not_rewrite_rollout_when_sqlite_is_malformed() {
let codex_home = codex_home_root();
let rollout = rollout_path(&codex_home, "sessions", "09", "thread-a");
write_text(
&codex_home.join("config.toml"),
"model_provider = \"any\"\n",
);
write_rollout(
&rollout,
"thread-a",
"old-provider",
"/tmp/project-gamma",
true,
);
write_text(&codex_home.join("state_5.sqlite"), "not a sqlite database");
let error = repair_codex_session_visibility_at(&codex_home)
.await
.expect_err("malformed sqlite stops repair");
assert!(error.to_string().contains("state_5.sqlite"));
assert_eq!(
first_line_json(&rollout)["payload"]["model_provider"],
"old-provider"
);
let _ = std::fs::remove_dir_all(codex_home);
}

Some files were not shown because too many files have changed in this diff Show More