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:
@@ -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: ®ex::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: ®ex::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: ®ex::Captures| {
|
||||
let src = caps.get(1).map_or("", |m| m.as_str());
|
||||
format!("", src)
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Bold
|
||||
let re_b = Regex::new(r"(?is)<(b|strong)\s*[^>]*>(.*?)</(b|strong)>").unwrap();
|
||||
text = re_b
|
||||
.replace_all(&text, |caps: ®ex::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: ®ex::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: ®ex::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 = [
|
||||
("&", "&"),
|
||||
("<", "<"),
|
||||
(">", ">"),
|
||||
(""", "\""),
|
||||
("'", "'"),
|
||||
("'", "'"),
|
||||
(" ", " "),
|
||||
("'", "'"),
|
||||
("/", "/"),
|
||||
("—", "—"),
|
||||
("–", "–"),
|
||||
("«", "«"),
|
||||
("»", "»"),
|
||||
("•", "•"),
|
||||
("…", "…"),
|
||||
("©", "©"),
|
||||
("®", "®"),
|
||||
("™", "™"),
|
||||
];
|
||||
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: ®ex::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: ®ex::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))
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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, ¤t_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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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(×tamp).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
Reference in New Issue
Block a user