Skip to content
Clone
194c44c
421 lines 14.3 KB Rust
use std::collections::HashMap;
use models::agent::agent_item as agent;
use models::agent::agents;
use models::dir::dirs::{self, DirKind};
use sea_orm::{
ActiveModelTrait, ColumnTrait, EntityTrait, PaginatorTrait, QueryFilter, QueryOrder,
QuerySelect, Set, TransactionTrait, Unchanged,
};
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use utoipa::{IntoParams, ToSchema};
use uuid::Uuid;
use crate::errors::AgentError;
use crate::services::service::AgentsService;
use crate::share::authz::require_collection_manage;
use crate::share::lookup::{collection_context, lookup_collection, space_context};
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct AgentCollectionDto {
pub id: Uuid,
pub dir_id: Uuid,
pub name: String,
pub icon: Option<String>,
pub description: Option<String>,
pub agent_count: i64,
/// Caller may attach the platform MCP to this set's agent runs (space
/// management permission plus a runtime-reachable MCP endpoint).
pub can_enable_mcp: bool,
/// The sandbox runtime bridge is configured for this deployment.
pub runtime_online: bool,
pub created_at: OffsetDateTime,
pub updated_at: Option<OffsetDateTime>,
}
pub type AgentCollectionPage = export::pagination::Page<AgentCollectionDto>;
#[derive(Debug, Clone, Deserialize, ToSchema)]
pub struct CreateAgentCollectionCmd {
pub name: String,
#[serde(default, alias = "dir_id")]
pub parent_dir_id: Option<Uuid>,
#[serde(default)]
pub icon: Option<String>,
#[serde(default)]
pub description: Option<String>,
}
#[derive(Debug, Clone, Deserialize, ToSchema)]
pub struct UpdateAgentCollectionCmd {
#[serde(default)]
pub name: Option<String>,
#[serde(default, deserialize_with = "export::serde_patch::deserialize_opt_opt")]
pub icon: Option<Option<String>>,
#[serde(default, deserialize_with = "export::serde_patch::deserialize_opt_opt")]
pub description: Option<Option<String>>,
}
#[derive(Debug, Clone, Deserialize, IntoParams, ToSchema)]
pub struct ListAgentCollectionQuery {
pub dir_id: Option<Uuid>,
pub q: Option<String>,
pub limit: Option<u64>,
pub offset: Option<u64>,
}
#[derive(Debug, Clone, Deserialize, IntoParams, ToSchema)]
pub struct AgentCollectionPath {
pub space_name: String,
pub collection_id: Uuid,
}
impl AgentsService {
fn collection_dto(
collection: agents::Model,
agent_count: i64,
can_enable_mcp: bool,
runtime_online: bool,
) -> AgentCollectionDto {
AgentCollectionDto {
id: collection.id,
dir_id: collection.dir_id,
name: collection.name,
icon: collection.icon,
description: collection.description,
agent_count,
can_enable_mcp,
runtime_online,
created_at: collection.created_at,
updated_at: collection.updated_at,
}
}
async fn agent_counts(
&self,
collection_ids: &[Uuid],
) -> Result<HashMap<Uuid, i64>, AgentError> {
if collection_ids.is_empty() {
return Ok(HashMap::new());
}
let rows: Vec<(Uuid, i64)> = agent::Entity::find()
.select_only()
.column(agent::Column::CollectionId)
.column_as(agent::Column::CollectionId.count(), "cnt")
.filter(agent::Column::CollectionId.is_in(collection_ids.to_vec()))
.group_by(agent::Column::CollectionId)
.into_tuple()
.all(self.db().read())
.await?;
Ok(rows.into_iter().collect())
}
#[tracing::instrument(
name = "agents.create_agent_collection",
skip_all,
fields(caller = %caller, space_name = %space_name)
)]
pub async fn create_agent_collection(
&self,
caller: i64,
space_name: &str,
cmd: CreateAgentCollectionCmd,
) -> Result<AgentCollectionDto, AgentError> {
let space = space_context(self, caller, space_name).await?;
require_collection_manage(self, &space, caller).await?;
let name = cmd.name.trim();
if name.is_empty() || name.chars().count() > self.config().agents.max_collection_name_len {
return Err(AgentError::BadRequest(
"collection name is empty or too long".to_string(),
));
}
if let Some(p) = cmd.parent_dir_id {
let parent = dirs::Entity::find_by_id(p)
.filter(dirs::Column::SpaceId.eq(Some(space.id)))
.one(self.db().read())
.await?
.ok_or(AgentError::NotFound)?;
if parent.kind != DirKind::Dir {
return Err(AgentError::BadRequest(
"parent node must be a dir".to_string(),
));
}
}
let (dir, collection) = {
let txn = self.db().write().begin().await?;
let position = {
let max_node = dirs::Entity::find()
.filter(dirs::Column::SpaceId.eq(Some(space.id)))
.filter(dirs::Column::ParentDirId.eq(cmd.parent_dir_id))
.order_by_desc(dirs::Column::Position)
.lock(sea_orm::sea_query::LockType::Update)
.one(&txn)
.await?;
max_node.map(|n| n.position + 1).unwrap_or(0)
};
let now = OffsetDateTime::now_utc();
let dir_row = dirs::ActiveModel {
id: Set(Uuid::now_v7()),
space_id: Set(Some(space.id)),
name: Set(name.to_string()),
kind: Set(DirKind::Object),
parent_dir_id: Set(cmd.parent_dir_id),
color: Set(None),
position: Set(position),
icon: Set(cmd.icon.clone()),
topic: Set(None),
purpose: Set(cmd.description.clone()),
is_archived: Set(false),
archived_at: Set(None),
archived_by: Set(None),
created_at: Set(now),
updated_at: Set(now),
created_by_uid: Set(caller),
};
let dir = dir_row.insert(&txn).await?;
let collection_row = agents::ActiveModel {
id: Set(Uuid::now_v7()),
dir_id: Set(dir.id),
name: Set(name.to_string()),
icon: Set(cmd.icon),
description: Set(cmd.description),
created_at: Set(now),
updated_at: Set(None),
};
let collection = collection_row.insert(&txn).await?;
txn.commit().await?;
(dir, collection)
};
let _ = dir;
let can_enable_mcp = crate::share::authz::can_enable_mcp(self, &space, caller).await;
Ok(Self::collection_dto(
collection,
0,
can_enable_mcp,
self.rpc_opt().is_some(),
))
}
pub async fn list_agent_collections(
&self,
caller: i64,
space_name: &str,
q: ListAgentCollectionQuery,
) -> Result<AgentCollectionPage, AgentError> {
let space = space_context(self, caller, space_name).await?;
let (limit, offset) = self.paginate(q.limit, q.offset);
let dirs_in_space = dirs::Entity::find()
.filter(dirs::Column::SpaceId.eq(Some(space.id)))
.filter(dirs::Column::Kind.eq(DirKind::Object))
.all(self.db().read())
.await?;
let dir_ids: Vec<Uuid> = if let Some(target_dir) = q.dir_id {
dirs_in_space
.into_iter()
.filter(|d| d.parent_dir_id == Some(target_dir) || d.id == target_dir)
.map(|d| d.id)
.collect()
} else {
dirs_in_space.into_iter().map(|d| d.id).collect()
};
if dir_ids.is_empty() {
return Ok(AgentCollectionPage {
items: Vec::new(),
total: 0,
limit,
offset,
});
}
let mut select = agents::Entity::find().filter(agents::Column::DirId.is_in(dir_ids));
if let Some(ref search) = q.q {
let kw = search.trim();
if !kw.is_empty() {
select = select.filter(agents::Column::Name.contains(kw));
}
}
select = select.order_by_desc(agents::Column::CreatedAt);
let total = select.clone().count(self.db().read()).await? as i64;
let rows = select
.limit(limit)
.offset(offset)
.all(self.db().read())
.await?;
let can_enable_mcp = crate::share::authz::can_enable_mcp(self, &space, caller).await;
let runtime_online = self.rpc_opt().is_some();
let ids: Vec<Uuid> = rows.iter().map(|row| row.id).collect();
let counts = self.agent_counts(&ids).await?;
let items = rows
.into_iter()
.map(|row| {
let count = counts.get(&row.id).copied().unwrap_or(0);
Self::collection_dto(row, count, can_enable_mcp, runtime_online)
})
.collect();
Ok(AgentCollectionPage {
items,
total,
limit,
offset,
})
}
pub async fn get_agent_collection(
&self,
caller: i64,
space_name: &str,
collection_id: Uuid,
) -> Result<AgentCollectionDto, AgentError> {
let (space, collection) =
collection_context(self, caller, space_name, collection_id).await?;
let counts = self.agent_counts(&[collection.id]).await?;
let count = counts.get(&collection.id).copied().unwrap_or(0);
let can_enable_mcp = crate::share::authz::can_enable_mcp(self, &space, caller).await;
Ok(Self::collection_dto(
collection,
count,
can_enable_mcp,
self.rpc_opt().is_some(),
))
}
#[tracing::instrument(
name = "agents.update_agent_collection",
skip_all,
fields(caller = %caller, space_name = %space_name, collection_id = %collection_id)
)]
pub async fn update_agent_collection(
&self,
caller: i64,
space_name: &str,
collection_id: Uuid,
cmd: UpdateAgentCollectionCmd,
) -> Result<AgentCollectionDto, AgentError> {
let space = space_context(self, caller, space_name).await?;
require_collection_manage(self, &space, caller).await?;
let collection = lookup_collection(self, &space, collection_id, caller).await?;
let now = OffsetDateTime::now_utc();
let mut am: agents::ActiveModel = collection.clone().into();
let mut dir_am = dirs::ActiveModel {
id: Unchanged(collection.dir_id),
..Default::default()
};
let mut dir_changed = false;
if let Some(new_name) = cmd.name {
let name = new_name.trim();
if name.is_empty()
|| name.chars().count() > self.config().agents.max_collection_name_len
{
return Err(AgentError::BadRequest(
"collection name is empty or too long".to_string(),
));
}
am.name = Set(name.to_string());
dir_am.name = Set(name.to_string());
dir_changed = true;
}
if let Some(icon) = cmd.icon {
am.icon = Set(icon.clone());
dir_am.icon = Set(icon);
dir_changed = true;
}
if let Some(description) = cmd.description {
am.description = Set(description.clone());
dir_am.purpose = Set(description);
dir_changed = true;
}
am.updated_at = Set(Some(now));
let updated = {
let txn = self.db().write().begin().await?;
if dir_changed {
dir_am.updated_at = Set(now);
dir_am.update(&txn).await?;
}
let updated = am.update(&txn).await?;
txn.commit().await?;
updated
};
let counts = self.agent_counts(&[updated.id]).await?;
let count = counts.get(&updated.id).copied().unwrap_or(0);
let can_enable_mcp = crate::share::authz::can_enable_mcp(self, &space, caller).await;
Ok(Self::collection_dto(
updated,
count,
can_enable_mcp,
self.rpc_opt().is_some(),
))
}
#[tracing::instrument(
name = "agents.delete_agent_collection",
skip_all,
fields(caller = %caller, space_name = %space_name, collection_id = %collection_id)
)]
pub async fn delete_agent_collection(
&self,
caller: i64,
space_name: &str,
collection_id: Uuid,
) -> Result<(), AgentError> {
let space = space_context(self, caller, space_name).await?;
require_collection_manage(self, &space, caller).await?;
let collection = lookup_collection(self, &space, collection_id, caller).await?;
let rows = agent::Entity::find()
.filter(agent::Column::CollectionId.eq(collection.id))
.all(self.db().read())
.await?;
if let Some(rpc) = self.rpc_opt() {
let space_ref = space.id.to_string();
let user_ref = caller.to_string();
for row in &rows {
if let Some(conversation_id) = row.conversation_id.as_deref()
&& let Err(error) = rpc
.delete_session(&space_ref, &user_ref, conversation_id)
.await
{
tracing::warn!(%error, agent_id = %row.id, "sandbox session cleanup failed");
}
}
}
self.db()
.primary()
.transaction::<_, (), AgentError>(|txn| {
Box::pin(async move {
agent::Entity::delete_many()
.filter(agent::Column::CollectionId.eq(collection.id))
.exec(txn)
.await?;
agents::Entity::delete_by_id(collection.id)
.exec(txn)
.await?;
dirs::Entity::delete_by_id(collection.dir_id)
.exec(txn)
.await?;
Ok(())
})
})
.await?;
Ok(())
}
}