Genarrative/server-rs/crates/module-runtime/src/agc_models.rs

use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// 目录 revision 乐观锁冲突。
pub const AGC_MODEL_CATALOG_CONFLICT: &str = "AGC_MODEL_CATALOG_CONFLICT";
/// 目录尚未初始化：SpacetimeDB 缺行，或存量内容与当前定义不符。
pub const AGC_MODEL_CATALOG_NOT_INITIALIZED: &str = "AGC_MODEL_CATALOG_NOT_INITIALIZED";
/// 客户端未显式选择模型时使用的占位标识。
pub const AGC_MODEL_PLATFORM_DEFAULT: &str = "platform-default";
/// 模型标识的长度上限，与客户端 `select_game_creator_model` 的校验保持一致。
pub const AGC_MODEL_ID_MAX_BYTES: usize = 64;
/// 目录项数上限，与后台「AGC 模型」页的新增上限保持一致。
pub const AGC_MODEL_CATALOG_MAX_MODELS: usize = 32;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AgcAgentMode {
    Codex,
    Cc,
}

impl Default for AgcAgentMode {
    fn default() -> Self {
        Self::Codex
    }
}

impl AgcAgentMode {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Codex => "codex",
            Self::Cc => "cc",
        }
    }

    pub fn parse(value: &str) -> Result<Self, String> {
        match value {
            "codex" => Ok(Self::Codex),
            "cc" => Ok(Self::Cc),
            _ => Err("Agent 模式只允许 codex 或 cc".into()),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AgcModel {
    pub id: String,
    pub alias: String,
    pub model_id: String,
    pub enabled: bool,
    /// 绑定 AGC 客户端实际启动的 Agent，存量目录缺失时兼容为 Codex。
    #[serde(default)]
    pub agent_mode: AgcAgentMode,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AgcModelCatalog {
    pub revision: u64,
    pub default_model_id: String,
    pub models: Vec<AgcModel>,
}

impl AgcModelCatalog {
    /// 按上游模型列表生成目录：`modelId` 是上游原始模型名，`alias` 也直接用原名
    /// （不再填「高质量/快速」这类人工别名），`id` 是模型名的稳定 slug。
    ///
    /// 上游返回顺序不稳定，所以先按原始模型名排序再生成，重复同步得到一致的目录与默认项。
    pub fn from_upstream_models(
        models: impl IntoIterator<Item = String>,
        revision: u64,
    ) -> Result<Self, String> {
        let mut model_names = models
            .into_iter()
            .map(|model| model.trim().to_string())
            .filter(|model| !model.is_empty())
            .collect::<Vec<_>>();
        model_names.sort();
        model_names.dedup();
        if model_names.is_empty() {
            return Err("上游模型列表为空".into());
        }

        let mut used_ids = HashSet::new();
        let mut entries = Vec::with_capacity(model_names.len());
        for model_id in model_names {
            let id = unique_model_id(&model_id, &mut used_ids);
            entries.push(AgcModel {
                id,
                alias: model_id.clone(),
                model_id,
                enabled: true,
                agent_mode: AgcAgentMode::default(),
            });
        }
        let default_model_id = entries
            .first()
            .map(|entry| entry.id.clone())
            .ok_or_else(|| "上游模型列表为空".to_string())?;
        let catalog = Self {
            revision,
            default_model_id,
            models: entries,
        };
        catalog.validate()?;
        Ok(catalog)
    }

    pub fn validate(&self) -> Result<(), String> {
        if self.models.is_empty() || self.models.len() > AGC_MODEL_CATALOG_MAX_MODELS {
            return Err(format!(
                "模型列表必须包含 1 至 {AGC_MODEL_CATALOG_MAX_MODELS} 项"
            ));
        }
        let mut ids = HashSet::new();
        let mut aliases = HashSet::new();
        for model in &self.models {
            if model.id.is_empty()
                || model.id == AGC_MODEL_PLATFORM_DEFAULT
                || model.id.len() > AGC_MODEL_ID_MAX_BYTES
                || !model
                    .id
                    .bytes()
                    .all(|c| c.is_ascii_alphanumeric() || c == b'-' || c == b'_')
                || !ids.insert(&model.id)
            {
                return Err("模型标识须唯一且只含字母、数字、连字符或下划线".into());
            }
            if model.alias.trim().is_empty()
                || model.alias != model.alias.trim()
                || model.alias.chars().count() > 40
                || model.alias.chars().any(char::is_control)
                || !aliases.insert(&model.alias)
            {
                return Err("模型别名须唯一，长度为 1 至 40 字符".into());
            }
            if model.model_id.trim().is_empty()
                || model.model_id != model.model_id.trim()
                || model.model_id.len() > 200
                || model.model_id.chars().any(char::is_control)
            {
                return Err("实际模型名不能为空，长度不能超过 200 字节".into());
            }
        }
        self.resolve(&self.default_model_id).map(|_| ())
    }

    pub fn resolve(&self, id: &str) -> Result<&str, String> {
        self.models
            .iter()
            .find(|m| m.enabled && m.id == id)
            .map(|m| m.model_id.as_str())
            .ok_or_else(|| "所选模型不可用，请刷新模型列表".into())
    }

    /// 请求侧解析：显式选择的标识按目录校验，未选择或占位标识使用默认项。
    pub fn resolve_requested(&self, requested: Option<&str>) -> Result<&str, String> {
        let requested = requested
            .map(str::trim)
            .filter(|id| !id.is_empty() && *id != AGC_MODEL_PLATFORM_DEFAULT);
        match requested {
            Some(id) => self.resolve(id),
            None => self.resolve(&self.default_model_id),
        }
    }
}

/// 由上游模型名生成稳定标识：只保留小写字母、数字、连字符与下划线，其余字符折叠成 `-`。
fn agc_model_id_from_name(model_name: &str) -> String {
    let mut id = String::new();
    let mut separator_pending = false;
    for value in model_name.chars() {
        let lowered = value.to_ascii_lowercase();
        if lowered.is_ascii_alphanumeric() || lowered == '_' {
            if separator_pending && !id.is_empty() {
                id.push('-');
            }
            separator_pending = false;
            id.push(lowered);
        } else {
            separator_pending = true;
        }
    }
    id
}

/// 生成在本次目录内唯一的标识：同名 slug 追加 `-2`/`-3`，并保证不超过长度上限。
fn unique_model_id(model_name: &str, used_ids: &mut HashSet<String>) -> String {
    let slug = agc_model_id_from_name(model_name);
    let slug = if slug.is_empty() {
        "model".to_string()
    } else {
        slug
    };
    // 预留后缀空间（`-` 加最多两位序号）后截断，保证候选标识仍在长度上限内。
    let base = slug
        .char_indices()
        .take_while(|(index, _)| *index < AGC_MODEL_ID_MAX_BYTES - 3)
        .map(|(_, value)| value)
        .collect::<String>();
    let base = base.trim_end_matches('-').to_string();
    let base = if base.is_empty() {
        "model".to_string()
    } else {
        base
    };

    let mut candidate = base.clone();
    let mut suffix = 2;
    while !used_ids.insert(candidate.clone()) {
        candidate = format!("{base}-{suffix}");
        suffix += 1;
    }
    candidate
}

#[cfg(test)]
mod tests {
    use super::*;

    fn upstream(models: &[&str]) -> Vec<String> {
        models.iter().map(|model| (*model).to_string()).collect()
    }

    fn model(id: &str, alias: &str, model_id: &str) -> AgcModel {
        AgcModel {
            id: id.into(),
            alias: alias.into(),
            model_id: model_id.into(),
            enabled: true,
            agent_mode: AgcAgentMode::default(),
        }
    }

    #[test]
    fn catalog_builds_from_upstream_models_with_stable_ids() {
        let catalog = AgcModelCatalog::from_upstream_models(
            upstream(&[
                " qwen3.8-flash ",
                "glm-5.3",
                "qwen3.8-flash",
                "deepseek-v4-pro",
                "",
                "vendor/model.v1:latest",
            ]),
            3,
        )
        .unwrap();

        assert_eq!(catalog.revision, 3);
        // 默认项是排序后第一项，与上游返回顺序无关。
        assert_eq!(catalog.default_model_id, "deepseek-v4-pro");
        assert_eq!(
            catalog.models,
            vec![
                model("deepseek-v4-pro", "deepseek-v4-pro", "deepseek-v4-pro"),
                model("glm-5-3", "glm-5.3", "glm-5.3"),
                model("qwen3-8-flash", "qwen3.8-flash", "qwen3.8-flash"),
                model(
                    "vendor-model-v1-latest",
                    "vendor/model.v1:latest",
                    "vendor/model.v1:latest"
                ),
            ]
        );
        assert!(catalog.validate().is_ok());
        // 同一模型集合重复生成结果一致。
        assert_eq!(
            AgcModelCatalog::from_upstream_models(
                upstream(&[
                    "vendor/model.v1:latest",
                    "deepseek-v4-pro",
                    "glm-5.3",
                    "qwen3.8-flash",
                ]),
                3
            )
            .unwrap(),
            catalog
        );
        assert!(AgcModelCatalog::from_upstream_models(upstream(&["", "   "]), 0).is_err());
    }

    #[test]
    fn catalog_keeps_ids_unique_and_within_client_contract() {
        // 不同模型名折叠成同一个 slug 时按排序追加序号，且标识始终符合客户端校验。
        let catalog = AgcModelCatalog::from_upstream_models(
            upstream(&["GLM-5.3", "glm/5.3", "glm_5.3", "模型名"]),
            0,
        )
        .unwrap();
        let ids = catalog
            .models
            .iter()
            .map(|entry| entry.id.as_str())
            .collect::<Vec<_>>();
        assert_eq!(ids, vec!["glm-5-3", "glm-5-3-2", "glm_5-3", "model"]);
        for entry in &catalog.models {
            assert!(entry.id.len() <= AGC_MODEL_ID_MAX_BYTES);
            assert!(
                entry
                    .id
                    .bytes()
                    .all(|c| c.is_ascii_alphanumeric() || c == b'-' || c == b'_')
            );
        }
        assert!(catalog.validate().is_ok());
    }

    #[test]
    fn catalog_maps_only_enabled_ids() {
        let mut catalog =
            AgcModelCatalog::from_upstream_models(upstream(&["model-a", "model-b"]), 0).unwrap();
        assert!(catalog.validate().is_ok());
        assert_eq!(catalog.resolve("model-a").unwrap(), "model-a");
        // 客户端不能直接指定实际模型名，只能回传目录标识。
        assert!(catalog.resolve("model-c").is_err());
        assert_eq!(catalog.resolve_requested(None).unwrap(), "model-a");
        assert_eq!(
            catalog
                .resolve_requested(Some(AGC_MODEL_PLATFORM_DEFAULT))
                .unwrap(),
            "model-a"
        );
        catalog.models[0].enabled = false;
        assert!(catalog.resolve("model-a").is_err());
        assert!(catalog.validate().is_err());
    }

    #[test]
    fn legacy_catalog_defaults_missing_agent_mode_to_codex() {
        let catalog: AgcModelCatalog = serde_json::from_str(
            r#"{"revision":4,"defaultModelId":"quality","models":[{"id":"quality","alias":"高质量","modelId":"gpt-6-astra","enabled":true}]}"#,
        )
        .expect("legacy catalog should remain readable");

        assert_eq!(catalog.models[0].agent_mode, AgcAgentMode::Codex);
        assert_eq!(catalog.models[0].agent_mode.as_str(), "codex");
    }

    #[test]
    fn catalog_agent_mode_accepts_codex_and_cc_only() {
        let mut catalog =
            AgcModelCatalog::from_upstream_models(upstream(&["model-a", "model-b"]), 0).unwrap();
        catalog.models[0].agent_mode = AgcAgentMode::Cc;
        assert!(catalog.validate().is_ok());
        assert_eq!(catalog.models[0].agent_mode.as_str(), "cc");

        let invalid: Result<AgcModelCatalog, _> = serde_json::from_str(
            r#"{"revision":0,"defaultModelId":"model-a","models":[{"id":"model-a","alias":"model-a","modelId":"model-a","enabled":true,"agentMode":"provider"}]}"#,
        );
        assert!(invalid.is_err());
    }

    #[test]
    fn catalog_rejects_duplicate_aliases_and_ids() {
        let mut catalog =
            AgcModelCatalog::from_upstream_models(upstream(&["model-a", "model-b"]), 0).unwrap();
        catalog.models[1].alias = catalog.models[0].alias.clone();
        assert!(catalog.validate().is_err());
        catalog.models[1].alias = "model-b".into();
        catalog.models[1].id = catalog.models[0].id.clone();
        assert!(catalog.validate().is_err());
        catalog.models[1].id = "model-b".into();
        catalog.models[1].id = AGC_MODEL_PLATFORM_DEFAULT.into();
        assert!(catalog.validate().is_err());
        catalog.models[1].id = "model-b".into();
        catalog.models[1].model_id = "".into();
        assert!(catalog.validate().is_err());

        let too_many = (0..AGC_MODEL_CATALOG_MAX_MODELS + 1)
            .map(|index| format!("model-{index}"))
            .collect::<Vec<_>>();
        assert_eq!(
            AgcModelCatalog::from_upstream_models(too_many, 0).unwrap_err(),
            format!("模型列表必须包含 1 至 {AGC_MODEL_CATALOG_MAX_MODELS} 项")
        );
    }
}