Skip to content
Clone
eba4fd1
102 lines 3.3 KB Rust
use std::collections::HashSet;
use std::time::Duration;
use cache::CacheClient;
use export::AppState;
use models::system::system_settings;
use sea_orm::sea_query::OnConflict;
use sea_orm::{DatabaseConnection, EntityTrait, Set};
use serde::{Deserialize, Serialize};
pub const DB_KEY: &str = "mcp.tools";
pub const CACHE_KEY: &str = "system:setting:mcp.tools";
const CACHE_TTL: Duration = Duration::from_secs(30);
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct McpToolSettings {
#[serde(default)]
pub disabled: Vec<String>,
#[serde(default)]
pub updated_by: Option<String>,
#[serde(default)]
pub updated_at: Option<String>,
}
pub async fn read(state: &AppState) -> McpToolSettings {
let cache = CacheClient::new(state.redis.clone());
if let Ok(Some(settings)) = cache.get::<McpToolSettings>(CACHE_KEY).await {
return settings;
}
match system_settings::Entity::find_by_id(DB_KEY.to_string())
.one(state.database.read())
.await
{
Ok(Some(row)) => match serde_json::from_str::<McpToolSettings>(&row.value) {
Ok(settings) => {
let _ = cache.set(CACHE_KEY, &settings, Some(CACHE_TTL)).await;
settings
}
Err(error) => {
tracing::warn!(%error, "invalid mcp tool settings payload; treating all tools as enabled");
McpToolSettings::default()
}
},
Ok(None) => McpToolSettings::default(),
Err(error) => {
tracing::warn!(%error, "failed to load mcp tool settings; treating all tools as enabled");
McpToolSettings::default()
}
}
}
pub async fn read_disabled(state: &AppState) -> HashSet<String> {
read(state).await.disabled.into_iter().collect()
}
pub async fn update(
state: &AppState,
disabled: Vec<String>,
updated_by: Option<String>,
) -> Result<McpToolSettings, String> {
let mut disabled: Vec<String> = disabled
.into_iter()
.map(|name| name.trim().to_string())
.filter(|name| !name.is_empty())
.collect();
disabled.sort();
disabled.dedup();
let settings = McpToolSettings {
disabled,
updated_by,
updated_at: Some(time::OffsetDateTime::now_utc().to_string()),
};
let value = serde_json::to_string(&settings).map_err(|error| error.to_string())?;
save(state.database.write(), &value)
.await
.map_err(|error| error.to_string())?;
if let Err(error) = export::redis::delete(&state.redis, CACHE_KEY).await {
tracing::warn!(%error, "failed to invalidate mcp tool settings cache");
}
Ok(settings)
}
async fn save(db: &DatabaseConnection, value: &str) -> Result<(), sea_orm::DbErr> {
let now = time::OffsetDateTime::now_utc();
let model = system_settings::ActiveModel {
key: Set(DB_KEY.to_string()),
value: Set(value.to_string()),
updated_at: Set(now),
};
system_settings::Entity::insert(model)
.on_conflict(
OnConflict::column(system_settings::Column::Key)
.update_columns([
system_settings::Column::Value,
system_settings::Column::UpdatedAt,
])
.to_owned(),
)
.exec_without_returning(db)
.await?;
Ok(())
}