use std::collections::{HashMap, HashSet};
use awaken_server_contract::contract::storage::StorageError;
use awaken_server_contract::{
BuiltinSeedSet, BuiltinSpec, ConfigRecord, ConfigStore, RecordMeta, RecordSource, SkillSpec,
validate_model_pool_spec_struct,
};
const SEED_LIST_PAGE_SIZE: usize = 256;
const BUILTIN_SEED_NAMESPACES: [&str; 7] = [
"agents",
"providers",
"models",
"model-pools",
"mcp-servers",
"tools",
"skills",
];
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SeedReport {
pub created: Vec<RecordRef>,
pub updated: Vec<RecordRef>,
pub unchanged: Vec<RecordRef>,
pub deleted: Vec<RecordRef>,
pub preserved_user: Vec<RecordRef>,
pub preserved_overridden: Vec<RecordRef>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordRef {
pub namespace: String,
pub id: String,
}
impl RecordRef {
fn new(namespace: &str, id: &str) -> Self {
Self {
namespace: namespace.to_owned(),
id: id.to_owned(),
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum SeedError {
#[error("storage error: {0}")]
Storage(#[from] StorageError),
#[error("serialization error: {0}")]
Serde(#[from] serde_json::Error),
#[error("agent spec '{id}' has invalid tool catalog: {errors}")]
InvalidAgentCatalog { id: String, errors: String },
#[error("skill spec '{id}' is invalid: {errors}")]
InvalidSkillSpec { id: String, errors: String },
#[error("model pool spec '{id}' is invalid: {errors}")]
InvalidModelPoolSpec { id: String, errors: String },
}
pub async fn apply_builtin_seed(
store: &dyn ConfigStore,
seed: &BuiltinSeedSet,
) -> Result<SeedReport, SeedError> {
for spec in &seed.specs {
match spec {
BuiltinSpec::Agent(agent) => validate_agent_spec_catalog(agent)?,
BuiltinSpec::ModelPool(pool) => {
validate_model_pool_spec_struct(pool).map_err(|error| {
SeedError::InvalidModelPoolSpec {
id: pool.id.clone(),
errors: error.to_string(),
}
})?;
}
BuiltinSpec::Skill(skill) => validate_builtin_skill_spec(skill)?,
_ => {}
}
}
let mut report = SeedReport::default();
let mut seeded: HashMap<&str, HashSet<String>> = HashMap::new();
for ns in BUILTIN_SEED_NAMESPACES {
seeded.insert(ns, HashSet::new());
}
for spec in &seed.specs {
let namespace = spec.namespace();
let id = spec.id();
let new_spec_value = builtin_spec_to_value(spec)?;
seeded.entry(namespace).or_default().insert(id.to_owned());
let existing_raw = store.get(namespace, id).await?;
match existing_raw {
None => {
let mut record = ConfigRecord {
spec: new_spec_value,
meta: RecordMeta::new_builtin(&seed.binary_version),
};
record.meta.revision = 1;
store
.put_if_absent(namespace, id, &record.to_value()?)
.await?;
report.created.push(RecordRef::new(namespace, id));
}
Some(raw) => {
let existing: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw)?;
match &existing.meta.source {
RecordSource::User => {
report.preserved_user.push(RecordRef::new(namespace, id));
}
RecordSource::Builtin {
binary_version: stored_version,
} => {
let same_version = stored_version == &seed.binary_version;
let same_spec = existing.spec == new_spec_value;
if same_version && same_spec {
report.unchanged.push(RecordRef::new(namespace, id));
} else {
let now = awaken_server_contract::time::now_ms();
let expected_revision = existing.meta.revision;
let record = ConfigRecord {
spec: new_spec_value,
meta: RecordMeta {
source: RecordSource::Builtin {
binary_version: seed.binary_version.clone(),
},
hidden: false,
user_overrides: existing.meta.user_overrides,
created_at: existing.meta.created_at,
updated_at: now,
revision: expected_revision + 1,
},
};
store
.put_if_revision(
namespace,
id,
&record.to_value()?,
expected_revision,
)
.await?;
report.updated.push(RecordRef::new(namespace, id));
}
}
}
}
}
}
for namespace in BUILTIN_SEED_NAMESPACES {
let empty = HashSet::new();
let seeded_ids: &HashSet<String> = seeded.get(namespace).unwrap_or(&empty);
let mut candidates: Vec<String> = Vec::new();
let mut offset = 0usize;
loop {
let page = store.list(namespace, offset, SEED_LIST_PAGE_SIZE).await?;
let page_len = page.len();
for (id, raw) in page {
if seeded_ids.contains(&id) {
continue;
}
let record: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw)?;
if matches!(record.meta.source, RecordSource::Builtin { .. }) {
candidates.push(id);
}
}
if page_len < SEED_LIST_PAGE_SIZE {
break;
}
offset += page_len;
}
for id in candidates {
let Some(raw) = store.get(namespace, &id).await? else {
continue;
};
let mut record: ConfigRecord<serde_json::Value> = ConfigRecord::from_value(raw)?;
let expected_revision = record.meta.revision;
if record.meta.user_overrides.is_some() {
record.meta.hidden = true;
record.meta.updated_at = awaken_server_contract::time::now_ms();
record.meta.revision = expected_revision + 1;
store
.put_if_revision(namespace, &id, &record.to_value()?, expected_revision)
.await?;
report
.preserved_overridden
.push(RecordRef::new(namespace, &id));
} else {
store
.delete_if_revision(namespace, &id, expected_revision)
.await?;
report.deleted.push(RecordRef::new(namespace, &id));
}
}
}
Ok(report)
}
fn builtin_spec_to_value(spec: &BuiltinSpec) -> Result<serde_json::Value, serde_json::Error> {
match spec {
BuiltinSpec::Agent(s) => serde_json::to_value(s.as_ref()),
BuiltinSpec::Provider(s) => serde_json::to_value(s),
BuiltinSpec::Model(s) => serde_json::to_value(s),
BuiltinSpec::ModelPool(s) => serde_json::to_value(s),
BuiltinSpec::A2aServer(s) => serde_json::to_value(s),
BuiltinSpec::McpServer(s) => serde_json::to_value(s),
BuiltinSpec::Tool(s) => serde_json::to_value(s),
BuiltinSpec::Skill(s) => serde_json::to_value(s),
}
}
fn validate_agent_spec_catalog(spec: &awaken_server_contract::AgentSpec) -> Result<(), SeedError> {
let errors = crate::services::agent_catalog::collect_catalog_errors(spec);
if errors.is_empty() {
Ok(())
} else {
Err(SeedError::InvalidAgentCatalog {
id: spec.id.clone(),
errors: errors.join("; "),
})
}
}
fn validate_builtin_skill_spec(spec: &SkillSpec) -> Result<(), SeedError> {
let value = serde_json::to_value(spec)?;
awaken_server_contract::validate_skill_spec(value).map_err(|error| {
SeedError::InvalidSkillSpec {
id: spec.id.clone(),
errors: error.to_string(),
}
})?;
Ok(())
}
#[cfg(test)]
mod tests;