forked from boos/BoosAPI
feat(desktop): embed BoosAPI Desktop client as a subproject
Vendor the rebranded BoosAPI Desktop client (Tauri + Vue) into metapi as a monorepo subdirectory under desktop/. Source only — node_modules, dist, Cargo target and the upstream .git history are excluded. desktop/ keeps its own independent Cargo workspace (desktop/src-tauri with crates/*); it is not part of the backend boosapi crate, so a `cargo check` at the metapi root still compiles only the backend. Build the desktop app from desktop/src-tauri per desktop/README.md. Note: app icons are still placeholder assets pending the final BoosAPI logo. The updater signing keys live outside the repo and are not included.
This commit is contained in:
boos
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "boosapi-agent"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
open-agent-sdk = { path = "../open-agent-sdk" }
|
||||
boosapi-core = { path = "../core" }
|
||||
boosapi-providers = { path = "../providers" }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
async-trait = "0.1"
|
||||
tracing = "0.1"
|
||||
tauri = "2"
|
||||
futures = "0.3"
|
||||
@@ -0,0 +1,609 @@
|
||||
//! BoosApiProviderBridge: adapts BoosApi's ProviderAdapter to the SDK's LLMProvider trait.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use open_agent_sdk::api::provider::ProviderRequest;
|
||||
use open_agent_sdk::api::ApiError;
|
||||
use open_agent_sdk::types::{ImageContentSource, ToolResultContentBlock};
|
||||
use open_agent_sdk::{
|
||||
ApiType, ContentBlock, LLMProvider, Message, MessageRole, ProviderResponse, SDKMessage, Usage,
|
||||
};
|
||||
|
||||
use boosapi_core::types::{
|
||||
ChatContent, ChatMessage, ChatRequest, ChatTool, ChatToolFunction, ContentPart, ImageUrl,
|
||||
ModelParamOverrides, TokenUsage, ToolCall, ToolCallFunction,
|
||||
};
|
||||
use boosapi_providers::{ProviderAdapter, ProviderRequestContext};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Bridge between BoosApi providers and the open-agent-sdk LLMProvider interface.
|
||||
pub struct BoosApiProviderBridge {
|
||||
adapter: Arc<dyn ProviderAdapter>,
|
||||
ctx: ProviderRequestContext,
|
||||
api_type: ApiType,
|
||||
model_param_overrides: Option<ModelParamOverrides>,
|
||||
app: Option<tauri::AppHandle>,
|
||||
conversation_id: Option<String>,
|
||||
}
|
||||
|
||||
impl BoosApiProviderBridge {
|
||||
pub fn new(
|
||||
adapter: Arc<dyn ProviderAdapter>,
|
||||
ctx: ProviderRequestContext,
|
||||
provider_type: &str,
|
||||
) -> Result<Self, String> {
|
||||
let api_type = match provider_type {
|
||||
"anthropic" => ApiType::AnthropicMessages,
|
||||
"openai" => ApiType::OpenAICompletions,
|
||||
"gemini" => ApiType::OpenAICompletions,
|
||||
"custom" => ApiType::OpenAICompletions,
|
||||
"openai_responses" => ApiType::OpenAICompletions,
|
||||
"deepseek" => ApiType::OpenAICompletions,
|
||||
"xai" => ApiType::OpenAICompletions,
|
||||
"glm" => ApiType::OpenAICompletions,
|
||||
"siliconflow" => ApiType::OpenAICompletions,
|
||||
other => {
|
||||
tracing::warn!(
|
||||
"Unknown provider type '{}', defaulting to OpenAI compat",
|
||||
other
|
||||
);
|
||||
ApiType::OpenAICompletions
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
adapter,
|
||||
ctx,
|
||||
api_type,
|
||||
model_param_overrides: None,
|
||||
app: None,
|
||||
conversation_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_model_param_overrides(mut self, overrides: Option<ModelParamOverrides>) -> Self {
|
||||
self.model_param_overrides = overrides;
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach a Tauri AppHandle for streaming text chunks to the frontend.
|
||||
pub fn with_app(mut self, app: tauri::AppHandle, conversation_id: String) -> Self {
|
||||
self.app = Some(app);
|
||||
self.conversation_id = Some(conversation_id);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LLMProvider for BoosApiProviderBridge {
|
||||
fn api_type(&self) -> ApiType {
|
||||
self.api_type.clone()
|
||||
}
|
||||
|
||||
async fn create_message(
|
||||
&self,
|
||||
request: ProviderRequest<'_>,
|
||||
stream_tx: Option<tokio::sync::mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, ApiError> {
|
||||
let chat_request = convert_request(request, self.model_param_overrides.as_ref());
|
||||
|
||||
let mut stream = self.adapter.chat_stream(&self.ctx, chat_request);
|
||||
let mut accumulated_text = String::new();
|
||||
let mut accumulated_thinking = String::new();
|
||||
let mut final_tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut final_usage: Option<TokenUsage> = None;
|
||||
|
||||
while let Some(result) = stream.next().await {
|
||||
match result {
|
||||
Ok(chunk) => {
|
||||
if let Some(ref text) = chunk.content {
|
||||
if !text.is_empty() {
|
||||
accumulated_text.push_str(text);
|
||||
|
||||
// Emit streaming text delta via SDK channel
|
||||
// (agent.rs will forward to frontend as agent-stream-text)
|
||||
if let Some(ref tx) = stream_tx {
|
||||
let _ = tx.try_send(SDKMessage::TextDelta { text: text.clone() });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref thinking) = chunk.thinking {
|
||||
if !thinking.is_empty() {
|
||||
accumulated_thinking.push_str(thinking);
|
||||
|
||||
// Emit streaming thinking delta via SDK channel
|
||||
// (agent.rs will forward to frontend as agent-stream-thinking)
|
||||
if let Some(ref tx) = stream_tx {
|
||||
let _ = tx.try_send(SDKMessage::ThinkingDelta {
|
||||
thinking: thinking.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if chunk.tool_calls.is_some() {
|
||||
final_tool_calls.clone_from(&chunk.tool_calls);
|
||||
}
|
||||
|
||||
if chunk.usage.is_some() {
|
||||
final_usage.clone_from(&chunk.usage);
|
||||
}
|
||||
|
||||
if chunk.done {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(classify_provider_error(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let usage = final_usage.unwrap_or(TokenUsage {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
});
|
||||
|
||||
let response = boosapi_core::types::ChatResponse {
|
||||
id: String::new(),
|
||||
model: String::new(),
|
||||
content: accumulated_text,
|
||||
thinking: if accumulated_thinking.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_thinking)
|
||||
},
|
||||
usage,
|
||||
tool_calls: final_tool_calls,
|
||||
};
|
||||
|
||||
Ok(convert_response(response))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SDK ProviderRequest → BoosApi ChatRequest
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn convert_request(
|
||||
request: ProviderRequest<'_>,
|
||||
model_param_overrides: Option<&ModelParamOverrides>,
|
||||
) -> ChatRequest {
|
||||
let messages: Vec<ChatMessage> = request
|
||||
.messages
|
||||
.iter()
|
||||
.flat_map(convert_sdk_message_to_chat_messages)
|
||||
.collect();
|
||||
|
||||
let tools: Option<Vec<ChatTool>> = request.tools.as_ref().map(|tools| {
|
||||
tools
|
||||
.iter()
|
||||
.map(|t| ChatTool {
|
||||
r#type: "function".to_string(),
|
||||
function: ChatToolFunction {
|
||||
name: t.name.clone(),
|
||||
description: Some(t.description.clone()),
|
||||
parameters: Some(t.input_schema.clone()),
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
let system_text: Option<String> = request.system.as_ref().map(|blocks| {
|
||||
blocks
|
||||
.iter()
|
||||
.map(|b| b.text.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
});
|
||||
|
||||
let mut final_messages = Vec::new();
|
||||
if let Some(sys) = system_text {
|
||||
final_messages.push(ChatMessage {
|
||||
role: "system".to_string(),
|
||||
content: ChatContent::Text(sys),
|
||||
reasoning_content: None,
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
final_messages.extend(messages);
|
||||
|
||||
let force_model_max_tokens = model_param_overrides
|
||||
.and_then(|overrides| overrides.force_max_tokens)
|
||||
== Some(true);
|
||||
let max_tokens = if request.max_tokens > 0 {
|
||||
Some(request.max_tokens as u32)
|
||||
} else if force_model_max_tokens {
|
||||
model_param_overrides
|
||||
.and_then(|overrides| overrides.max_tokens)
|
||||
.or(Some(4096))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
ChatRequest {
|
||||
model: request.model.to_string(),
|
||||
messages: final_messages,
|
||||
stream: true,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
max_tokens,
|
||||
tools,
|
||||
thinking_budget: request
|
||||
.thinking
|
||||
.as_ref()
|
||||
.and_then(|t| t.budget_tokens.map(|b| b as u32)),
|
||||
thinking_level: None,
|
||||
reasoning_profile: model_param_overrides
|
||||
.and_then(|overrides| overrides.reasoning_profile.clone()),
|
||||
use_max_completion_tokens: model_param_overrides
|
||||
.and_then(|overrides| overrides.use_max_completion_tokens),
|
||||
thinking_param_style: model_param_overrides
|
||||
.and_then(|overrides| overrides.thinking_param_style.clone()),
|
||||
extra_body: model_param_overrides.and_then(|overrides| overrides.extra_body.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single SDK Message into one or more BoosApi ChatMessages.
|
||||
/// ToolResult content blocks become separate tool-role messages.
|
||||
fn convert_sdk_message_to_chat_messages(msg: &Message) -> Vec<ChatMessage> {
|
||||
let role = match msg.role {
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
};
|
||||
|
||||
let mut text_parts: Vec<String> = Vec::new();
|
||||
let mut thinking_parts: Vec<String> = Vec::new();
|
||||
let mut tool_calls: Vec<ToolCall> = Vec::new();
|
||||
let mut tool_results: Vec<(String, String)> = Vec::new();
|
||||
let mut image_parts: Vec<ImageContentSource> = Vec::new();
|
||||
|
||||
for block in &msg.content {
|
||||
match block {
|
||||
ContentBlock::Text { text } => {
|
||||
text_parts.push(text.clone());
|
||||
}
|
||||
ContentBlock::ToolUse { id, name, input } => {
|
||||
tool_calls.push(ToolCall {
|
||||
id: id.clone(),
|
||||
call_type: "function".to_string(),
|
||||
function: ToolCallFunction {
|
||||
name: name.clone(),
|
||||
arguments: serde_json::to_string(input).unwrap_or_default(),
|
||||
},
|
||||
});
|
||||
}
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error: _,
|
||||
} => {
|
||||
let text = content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ToolResultContentBlock::Text { text } => Some(text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
tool_results.push((tool_use_id.clone(), text));
|
||||
}
|
||||
ContentBlock::Thinking { thinking, .. } => {
|
||||
if msg.role == MessageRole::Assistant && !thinking.is_empty() {
|
||||
thinking_parts.push(thinking.clone());
|
||||
}
|
||||
}
|
||||
ContentBlock::Image { source } => {
|
||||
image_parts.push(source.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut result = Vec::new();
|
||||
|
||||
if role == "assistant" {
|
||||
let content = if text_parts.is_empty() {
|
||||
ChatContent::Text(String::new())
|
||||
} else {
|
||||
ChatContent::Text(text_parts.join(""))
|
||||
};
|
||||
|
||||
result.push(ChatMessage {
|
||||
role: "assistant".to_string(),
|
||||
content,
|
||||
reasoning_content: if thinking_parts.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(thinking_parts.join(""))
|
||||
},
|
||||
tool_calls: if tool_calls.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_calls)
|
||||
},
|
||||
tool_call_id: None,
|
||||
});
|
||||
} else if !tool_results.is_empty() {
|
||||
for (tool_use_id, text) in tool_results {
|
||||
result.push(ChatMessage {
|
||||
role: "tool".to_string(),
|
||||
content: ChatContent::Text(text),
|
||||
reasoning_content: None,
|
||||
tool_calls: None,
|
||||
tool_call_id: Some(tool_use_id),
|
||||
});
|
||||
}
|
||||
} else {
|
||||
let content = if !image_parts.is_empty() {
|
||||
let mut parts: Vec<ContentPart> = text_parts
|
||||
.iter()
|
||||
.map(|t| ContentPart {
|
||||
r#type: "text".to_string(),
|
||||
text: Some(t.clone()),
|
||||
image_url: None,
|
||||
})
|
||||
.collect();
|
||||
for img in &image_parts {
|
||||
parts.push(ContentPart {
|
||||
r#type: "image_url".to_string(),
|
||||
text: None,
|
||||
image_url: Some(ImageUrl {
|
||||
url: format!("data:{};base64,{}", img.media_type, img.data),
|
||||
}),
|
||||
});
|
||||
}
|
||||
ChatContent::Multipart(parts)
|
||||
} else {
|
||||
ChatContent::Text(text_parts.join(""))
|
||||
};
|
||||
|
||||
result.push(ChatMessage {
|
||||
role: role.to_string(),
|
||||
content,
|
||||
reasoning_content: None,
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// BoosApi ChatResponse → SDK ProviderResponse
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn convert_response(response: boosapi_core::types::ChatResponse) -> ProviderResponse {
|
||||
let mut content_blocks: Vec<ContentBlock> = Vec::new();
|
||||
|
||||
if let Some(thinking) = &response.thinking {
|
||||
if !thinking.is_empty() {
|
||||
content_blocks.push(ContentBlock::Thinking {
|
||||
thinking: thinking.clone(),
|
||||
signature: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if !response.content.is_empty() {
|
||||
content_blocks.push(ContentBlock::Text {
|
||||
text: response.content.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(tool_calls) = &response.tool_calls {
|
||||
for tc in tool_calls {
|
||||
let input: Value = serde_json::from_str(&tc.function.arguments).unwrap_or(Value::Null);
|
||||
content_blocks.push(ContentBlock::ToolUse {
|
||||
id: tc.id.clone(),
|
||||
name: tc.function.name.clone(),
|
||||
input,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let stop_reason = if response
|
||||
.tool_calls
|
||||
.as_ref()
|
||||
.map_or(false, |tc| !tc.is_empty())
|
||||
{
|
||||
Some("tool_use".to_string())
|
||||
} else {
|
||||
Some("end_turn".to_string())
|
||||
};
|
||||
|
||||
ProviderResponse {
|
||||
message: Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: content_blocks,
|
||||
},
|
||||
usage: Usage {
|
||||
input_tokens: response.usage.prompt_tokens as u64,
|
||||
output_tokens: response.usage.completion_tokens as u64,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
},
|
||||
stop_reason,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Error classification
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Parse an HTTP status code from provider error messages like "OpenAI API error 400 Bad Request: ..."
|
||||
fn parse_http_status(err: &str) -> Option<u16> {
|
||||
for pattern in &["API error ", "error "] {
|
||||
if let Some(pos) = err.find(pattern) {
|
||||
let after = &err[pos + pattern.len()..];
|
||||
if let Some(end) = after.find(|c: char| !c.is_ascii_digit()) {
|
||||
if end > 0 {
|
||||
if let Ok(status) = after[..end].parse::<u16>() {
|
||||
if (100..600).contains(&status) {
|
||||
return Some(status);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Classify a provider error into the appropriate SDK ApiError variant,
|
||||
/// so the retry logic only retries truly transient errors.
|
||||
fn classify_provider_error(e: boosapi_core::error::BoosApiError) -> ApiError {
|
||||
let err_str = e.to_string();
|
||||
if let Some(status) = parse_http_status(&err_str) {
|
||||
if status == 401 || status == 403 {
|
||||
ApiError::AuthError(err_str)
|
||||
} else if status == 429 {
|
||||
ApiError::RateLimitError
|
||||
} else if (400..500).contains(&status) {
|
||||
// Client errors (400, 404, 422, etc.) are NOT retryable
|
||||
ApiError::HttpError {
|
||||
status,
|
||||
message: err_str,
|
||||
}
|
||||
} else {
|
||||
// 5xx errors are retryable via NetworkError mapping
|
||||
ApiError::NetworkError(err_str)
|
||||
}
|
||||
} else {
|
||||
ApiError::NetworkError(err_str)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use open_agent_sdk::api::provider::ProviderRequest;
|
||||
|
||||
fn param_overrides() -> ModelParamOverrides {
|
||||
ModelParamOverrides {
|
||||
temperature: None,
|
||||
max_tokens: Some(2048),
|
||||
top_p: None,
|
||||
frequency_penalty: None,
|
||||
use_max_completion_tokens: Some(false),
|
||||
no_system_role: None,
|
||||
force_max_tokens: None,
|
||||
thinking_param_style: Some("enable_thinking".to_string()),
|
||||
reasoning_profile: Some("siliconflow_enable_thinking".to_string()),
|
||||
reasoning_options: None,
|
||||
reasoning_default: None,
|
||||
extra_body: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_request_applies_non_token_model_param_overrides() {
|
||||
let messages = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
}];
|
||||
let request = ProviderRequest {
|
||||
model: "deepseek-reasoner",
|
||||
max_tokens: 0,
|
||||
messages: &messages,
|
||||
system: None,
|
||||
tools: None,
|
||||
thinking: None,
|
||||
};
|
||||
|
||||
let converted = convert_request(request, Some(¶m_overrides()));
|
||||
|
||||
assert_eq!(converted.max_tokens, None);
|
||||
assert_eq!(
|
||||
converted.reasoning_profile.as_deref(),
|
||||
Some("siliconflow_enable_thinking")
|
||||
);
|
||||
assert_eq!(converted.use_max_completion_tokens, Some(false));
|
||||
assert_eq!(
|
||||
converted.thinking_param_style.as_deref(),
|
||||
Some("enable_thinking")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_request_uses_model_max_tokens_only_when_forced() {
|
||||
let messages = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
}];
|
||||
let request = ProviderRequest {
|
||||
model: "deepseek-reasoner",
|
||||
max_tokens: 0,
|
||||
messages: &messages,
|
||||
system: None,
|
||||
tools: None,
|
||||
thinking: None,
|
||||
};
|
||||
let mut overrides = param_overrides();
|
||||
overrides.force_max_tokens = Some(true);
|
||||
|
||||
let converted = convert_request(request, Some(&overrides));
|
||||
|
||||
assert_eq!(converted.max_tokens, Some(2048));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_request_prefers_sdk_max_tokens_over_model_default() {
|
||||
let messages = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
}];
|
||||
let request = ProviderRequest {
|
||||
model: "deepseek-reasoner",
|
||||
max_tokens: 8192,
|
||||
messages: &messages,
|
||||
system: None,
|
||||
tools: None,
|
||||
thinking: None,
|
||||
};
|
||||
|
||||
let converted = convert_request(request, Some(¶m_overrides()));
|
||||
|
||||
assert_eq!(converted.max_tokens, Some(8192));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_thinking_block_becomes_reasoning_content() {
|
||||
let message = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Thinking {
|
||||
thinking: "hidden reasoning".to_string(),
|
||||
signature: None,
|
||||
},
|
||||
ContentBlock::Text {
|
||||
text: "final answer".to_string(),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "call-1".to_string(),
|
||||
name: "read_file".to_string(),
|
||||
input: serde_json::json!({"path": "README.md"}),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
let converted = convert_sdk_message_to_chat_messages(&message);
|
||||
|
||||
assert_eq!(converted.len(), 1);
|
||||
assert_eq!(
|
||||
converted[0].reasoning_content.as_deref(),
|
||||
Some("hidden reasoning")
|
||||
);
|
||||
assert!(converted[0].tool_calls.is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod bridge;
|
||||
pub mod permission;
|
||||
pub mod security;
|
||||
@@ -0,0 +1,108 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Tool risk classification
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum RiskLevel {
|
||||
ReadOnly,
|
||||
Write,
|
||||
Execute,
|
||||
}
|
||||
|
||||
/// Permission mode for the agent session
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum PermissionMode {
|
||||
Default,
|
||||
AcceptEdits,
|
||||
FullAccess,
|
||||
}
|
||||
|
||||
impl PermissionMode {
|
||||
pub fn from_str(s: &str) -> Self {
|
||||
match s {
|
||||
"accept_edits" => Self::AcceptEdits,
|
||||
"full_access" => Self::FullAccess,
|
||||
_ => Self::Default,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the permission system decides
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PermissionAction {
|
||||
AutoAllow,
|
||||
RequireApproval,
|
||||
HardDeny,
|
||||
}
|
||||
|
||||
/// Classify a tool's risk level based on its name
|
||||
pub fn classify_tool_risk(tool_name: &str) -> RiskLevel {
|
||||
let name_lower = tool_name.to_lowercase();
|
||||
|
||||
// Execute-level tools
|
||||
if matches!(
|
||||
name_lower.as_str(),
|
||||
"bash" | "shell" | "run_command" | "execute"
|
||||
) || name_lower.contains("exec")
|
||||
|| name_lower.contains("run")
|
||||
|| name_lower.contains("bash")
|
||||
|| name_lower.contains("shell")
|
||||
{
|
||||
return RiskLevel::Execute;
|
||||
}
|
||||
|
||||
// Write-level tools
|
||||
if matches!(
|
||||
name_lower.as_str(),
|
||||
"write"
|
||||
| "edit"
|
||||
| "create"
|
||||
| "delete"
|
||||
| "rename"
|
||||
| "patch"
|
||||
| "write_file"
|
||||
| "edit_file"
|
||||
| "create_file"
|
||||
| "delete_file"
|
||||
| "move"
|
||||
| "mkdir"
|
||||
| "remove"
|
||||
) || name_lower.contains("write")
|
||||
|| name_lower.contains("edit")
|
||||
|| name_lower.contains("create")
|
||||
|| name_lower.contains("delete")
|
||||
|| name_lower.contains("patch")
|
||||
{
|
||||
return RiskLevel::Write;
|
||||
}
|
||||
|
||||
// Everything else is read-only
|
||||
RiskLevel::ReadOnly
|
||||
}
|
||||
|
||||
/// Decision matrix: given permission mode, risk level, and whether the tool
|
||||
/// is in the "always allowed" set, return the action to take.
|
||||
pub fn decide_permission(
|
||||
mode: PermissionMode,
|
||||
risk: RiskLevel,
|
||||
is_always_allowed: bool,
|
||||
) -> PermissionAction {
|
||||
// If tool was previously approved with "always allow", auto-allow
|
||||
if is_always_allowed {
|
||||
return PermissionAction::AutoAllow;
|
||||
}
|
||||
|
||||
match (mode, risk) {
|
||||
// Default mode: only read is auto-allowed
|
||||
(PermissionMode::Default, RiskLevel::ReadOnly) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::Default, RiskLevel::Write) => PermissionAction::RequireApproval,
|
||||
(PermissionMode::Default, RiskLevel::Execute) => PermissionAction::RequireApproval,
|
||||
|
||||
// Accept edits: read + write auto-allowed
|
||||
(PermissionMode::AcceptEdits, RiskLevel::ReadOnly) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Write) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Execute) => PermissionAction::RequireApproval,
|
||||
|
||||
// Full access: everything auto-allowed
|
||||
(PermissionMode::FullAccess, _) => PermissionAction::AutoAllow,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
use open_agent_sdk::PermissionDecision;
|
||||
use serde_json::Value;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Validate that a path resolves to within the cwd.
|
||||
/// For existing paths, uses canonicalize (resolves symlinks).
|
||||
/// For non-existing paths, validates the parent directory.
|
||||
pub fn validate_path_within_cwd(path: &str, cwd: &str) -> Result<PathBuf, String> {
|
||||
let target = Path::new(path);
|
||||
let cwd_path = Path::new(cwd);
|
||||
|
||||
// Canonicalize cwd (must exist)
|
||||
let cwd_canonical = cwd_path
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Cannot resolve working directory '{}': {}", cwd, e))?;
|
||||
|
||||
// Try to resolve the target path
|
||||
let resolved = if target.is_absolute() {
|
||||
target.to_path_buf()
|
||||
} else {
|
||||
cwd_path.join(target)
|
||||
};
|
||||
|
||||
// If path exists, canonicalize it (resolves symlinks)
|
||||
if resolved.exists() {
|
||||
let canonical = resolved
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Cannot resolve path '{}': {}", path, e))?;
|
||||
if canonical.starts_with(&cwd_canonical) {
|
||||
Ok(canonical)
|
||||
} else {
|
||||
Err(format!(
|
||||
"Path '{}' resolves to '{}' which is outside working directory '{}'",
|
||||
path,
|
||||
canonical.display(),
|
||||
cwd_canonical.display()
|
||||
))
|
||||
}
|
||||
} else {
|
||||
// For non-existing paths, check parent directory
|
||||
if let Some(parent) = resolved.parent() {
|
||||
if parent.exists() {
|
||||
let parent_canonical = parent
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Cannot resolve parent of '{}': {}", path, e))?;
|
||||
if parent_canonical.starts_with(&cwd_canonical) {
|
||||
Ok(resolved)
|
||||
} else {
|
||||
Err(format!(
|
||||
"Parent of '{}' resolves outside working directory '{}'",
|
||||
path,
|
||||
cwd_canonical.display()
|
||||
))
|
||||
}
|
||||
} else {
|
||||
// Parent doesn't exist — walk up until we find an existing ancestor
|
||||
let mut ancestor = parent.to_path_buf();
|
||||
loop {
|
||||
if ancestor.exists() {
|
||||
let anc_canonical = ancestor
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Cannot resolve ancestor: {}", e))?;
|
||||
if anc_canonical.starts_with(&cwd_canonical) {
|
||||
return Ok(resolved);
|
||||
} else {
|
||||
return Err(format!(
|
||||
"Path '{}' is outside working directory '{}'",
|
||||
path,
|
||||
cwd_canonical.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
if !ancestor.pop() {
|
||||
return Err(format!(
|
||||
"Cannot verify path '{}' against working directory",
|
||||
path
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Err(format!("Invalid path: '{}'", path))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a tool's path arguments are safe (within cwd).
|
||||
/// Returns Some(PermissionDecision::Deny(reason)) if unsafe, None if safe or not applicable.
|
||||
pub fn check_path_safety(tool_name: &str, input: &Value, cwd: &str) -> Option<PermissionDecision> {
|
||||
let name_lower = tool_name.to_lowercase();
|
||||
|
||||
match name_lower.as_str() {
|
||||
// Single-file path tools
|
||||
"read" | "read_file" | "write" | "write_file" | "edit" | "edit_file" | "create"
|
||||
| "create_file" | "delete" | "delete_file" | "rename" | "list_dir" | "listdir" => {
|
||||
check_single_path(input, cwd)
|
||||
}
|
||||
|
||||
// Glob tool — check base directory of pattern
|
||||
"glob" | "glob_search" => check_glob_path(input, cwd),
|
||||
|
||||
// Grep tool — optional path
|
||||
"grep" | "search" | "ripgrep" => check_grep_path(input, cwd),
|
||||
|
||||
// Bash/shell — no path checking (cwd is injected, but cd escape is Phase 4)
|
||||
"bash" | "shell" | "run_command" | "execute" => None,
|
||||
|
||||
// Unknown tools — don't block
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn check_single_path(input: &Value, cwd: &str) -> Option<PermissionDecision> {
|
||||
// Try common field names for path
|
||||
let path_str = input
|
||||
.get("path")
|
||||
.or_else(|| input.get("file_path"))
|
||||
.or_else(|| input.get("file"))
|
||||
.and_then(|v| v.as_str());
|
||||
|
||||
if let Some(path) = path_str {
|
||||
if let Err(reason) = validate_path_within_cwd(path, cwd) {
|
||||
return Some(PermissionDecision::Deny(reason));
|
||||
}
|
||||
}
|
||||
|
||||
// Also check "new_path" for rename operations
|
||||
if let Some(new_path) = input.get("new_path").and_then(|v| v.as_str()) {
|
||||
if let Err(reason) = validate_path_within_cwd(new_path, cwd) {
|
||||
return Some(PermissionDecision::Deny(reason));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn check_glob_path(input: &Value, cwd: &str) -> Option<PermissionDecision> {
|
||||
if let Some(path) = input.get("path").and_then(|v| v.as_str()) {
|
||||
if let Err(reason) = validate_path_within_cwd(path, cwd) {
|
||||
return Some(PermissionDecision::Deny(reason));
|
||||
}
|
||||
}
|
||||
|
||||
let pattern = input
|
||||
.get("pattern")
|
||||
.or_else(|| input.get("glob"))
|
||||
.and_then(|v| v.as_str())?;
|
||||
|
||||
// Extract base directory from glob pattern (non-wildcard prefix)
|
||||
let base_dir = extract_glob_base(pattern);
|
||||
|
||||
if !base_dir.is_empty() {
|
||||
if let Err(reason) = validate_path_within_cwd(&base_dir, cwd) {
|
||||
return Some(PermissionDecision::Deny(reason));
|
||||
}
|
||||
}
|
||||
// Empty base_dir means pattern starts with wildcard → relative to cwd → safe
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn check_grep_path(input: &Value, cwd: &str) -> Option<PermissionDecision> {
|
||||
// Grep path is optional — if not specified, defaults to cwd (safe)
|
||||
if let Some(path) = input.get("path").and_then(|v| v.as_str()) {
|
||||
if let Err(reason) = validate_path_within_cwd(path, cwd) {
|
||||
return Some(PermissionDecision::Deny(reason));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract the non-wildcard prefix from a glob pattern as a base directory.
|
||||
/// e.g., "src/components/**/*.tsx" → "src/components"
|
||||
/// "**/*.rs" → ""
|
||||
/// "/absolute/path/to/*.txt" → "/absolute/path/to"
|
||||
fn extract_glob_base(pattern: &str) -> String {
|
||||
let mut parts = Vec::new();
|
||||
for segment in pattern.split('/') {
|
||||
if segment.contains('*')
|
||||
|| segment.contains('?')
|
||||
|| segment.contains('[')
|
||||
|| segment.contains('{')
|
||||
{
|
||||
break;
|
||||
}
|
||||
parts.push(segment);
|
||||
}
|
||||
parts.join("/")
|
||||
}
|
||||
@@ -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
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user