Genarrative/server-rs/crates/agent-runtime-core/src/profile.rs

use std::collections::{BTreeMap, BTreeSet};

use serde_json::{Map, Value};

use crate::capability::CapabilityRegistry;
use crate::catalog::collect_unique_capability_ids;
use crate::contract::{ContractError, ContractErrorKind, validate_identifier, validate_metadata};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RunProfileDefinition {
    id: String,
    capability_ids: Vec<String>,
    completion_policy_id: String,
    metadata: Value,
}

impl RunProfileDefinition {
    pub fn try_new(
        id: impl Into<String>,
        capability_ids: impl IntoIterator<Item = impl Into<String>>,
        completion_policy_id: impl Into<String>,
    ) -> Result<Self, ContractError> {
        let id = id.into();
        let completion_policy_id = completion_policy_id.into();
        validate_identifier(&id, "run profile id")?;
        validate_identifier(&completion_policy_id, "completion policy id")?;
        let capability_ids = collect_unique_capability_ids(capability_ids, "run profile")?;
        Ok(Self {
            id,
            capability_ids,
            completion_policy_id,
            metadata: Value::Object(Map::new()),
        })
    }

    pub fn with_metadata(mut self, metadata: Value) -> Result<Self, ContractError> {
        validate_metadata(&metadata, "run profile metadata")?;
        self.metadata = metadata;
        Ok(self)
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn capability_ids(&self) -> &[String] {
        &self.capability_ids
    }

    pub fn completion_policy_id(&self) -> &str {
        &self.completion_policy_id
    }

    pub fn metadata(&self) -> &Value {
        &self.metadata
    }
}

#[derive(Clone, Debug)]
pub struct RunProfileCatalog {
    profiles: Vec<RunProfileDefinition>,
    by_id: BTreeMap<String, usize>,
}

impl RunProfileCatalog {
    pub fn try_new(
        profiles: impl IntoIterator<Item = RunProfileDefinition>,
    ) -> Result<Self, ContractError> {
        let profiles = profiles.into_iter().collect::<Vec<_>>();
        let mut by_id = BTreeMap::new();
        for (index, profile) in profiles.iter().enumerate() {
            if by_id.insert(profile.id.clone(), index).is_some() {
                return Err(ContractError::new(
                    ContractErrorKind::DuplicateId,
                    format!("run profile id 重复：{}", profile.id),
                ));
            }
        }
        Ok(Self { profiles, by_id })
    }

    pub fn get(&self, id: &str) -> Option<&RunProfileDefinition> {
        self.by_id
            .get(id)
            .and_then(|index| self.profiles.get(*index))
    }

    pub fn iter(&self) -> impl ExactSizeIterator<Item = &RunProfileDefinition> {
        self.profiles.iter()
    }

    pub fn validate_capabilities<D>(
        &self,
        capabilities: &CapabilityRegistry<D>,
    ) -> Result<(), ContractError> {
        for profile in &self.profiles {
            for capability_id in &profile.capability_ids {
                if capabilities.get(capability_id).is_none() {
                    return Err(ContractError::new(
                        ContractErrorKind::UnknownCapability,
                        format!(
                            "run profile {} 引用了未知 capability：{capability_id}",
                            profile.id
                        ),
                    ));
                }
            }
        }
        Ok(())
    }

    pub fn validate_completion_policy_ids<'a>(
        &self,
        policy_ids: impl IntoIterator<Item = &'a str>,
    ) -> Result<(), ContractError> {
        let mut known_policy_ids = BTreeSet::new();
        for policy_id in policy_ids {
            validate_identifier(policy_id, "completion policy id")?;
            known_policy_ids.insert(policy_id);
        }
        for profile in &self.profiles {
            if !known_policy_ids.contains(profile.completion_policy_id.as_str()) {
                return Err(ContractError::new(
                    ContractErrorKind::UnknownCompletionPolicy,
                    format!(
                        "run profile {} 引用了未知 completion policy：{}",
                        profile.id, profile.completion_policy_id
                    ),
                ));
            }
        }
        Ok(())
    }
}