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,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(())
|
||||
}
|
||||
Reference in New Issue
Block a user