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

use std::collections::BTreeMap;
use std::fmt;

use serde_json::Value;

use crate::contract::{
    ContractError, ContractErrorKind, validate_description, validate_function_name,
    validate_identifier,
};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CapabilityRegistryErrorKind {
    InvalidDefinition,
    DuplicateId,
    DuplicateFunctionName,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CapabilityRegistryError {
    kind: CapabilityRegistryErrorKind,
    detail: String,
}

impl CapabilityRegistryError {
    pub fn kind(&self) -> CapabilityRegistryErrorKind {
        self.kind
    }

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

impl fmt::Display for CapabilityRegistryError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.detail)
    }
}

impl std::error::Error for CapabilityRegistryError {}

impl From<ContractError> for CapabilityRegistryError {
    fn from(error: ContractError) -> Self {
        let kind = match error.kind() {
            ContractErrorKind::DuplicateId => CapabilityRegistryErrorKind::DuplicateId,
            ContractErrorKind::InvalidDefinition
            | ContractErrorKind::DuplicateReference
            | ContractErrorKind::UnknownCapability
            | ContractErrorKind::UnknownCompletionPolicy => {
                CapabilityRegistryErrorKind::InvalidDefinition
            }
        };
        Self {
            kind,
            detail: error.to_string(),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CapabilityDefinition<D> {
    id: String,
    function_name: String,
    description: String,
    input_schema: Value,
    dispatch: D,
}

impl<D> CapabilityDefinition<D> {
    pub fn try_new(
        id: impl Into<String>,
        function_name: impl Into<String>,
        description: impl Into<String>,
        input_schema: Value,
        dispatch: D,
    ) -> Result<Self, CapabilityRegistryError> {
        let id = id.into();
        let function_name = function_name.into();
        let description = description.into();
        validate_identifier(&id, "capability id")?;
        validate_function_name(&function_name)?;
        validate_description(&description, "capability description")?;
        validate_input_schema(&id, &input_schema)?;
        Ok(Self {
            id,
            function_name,
            description,
            input_schema,
            dispatch,
        })
    }

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

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

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

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

    pub fn dispatch(&self) -> &D {
        &self.dispatch
    }
}

fn validate_input_schema(id: &str, schema: &Value) -> Result<(), CapabilityRegistryError> {
    let Some(object) = schema.as_object() else {
        return Err(CapabilityRegistryError {
            kind: CapabilityRegistryErrorKind::InvalidDefinition,
            detail: format!("capability {id} 的 inputSchema 必须是 JSON object"),
        });
    };
    if object.get("type").and_then(Value::as_str) != Some("object") {
        return Err(CapabilityRegistryError {
            kind: CapabilityRegistryErrorKind::InvalidDefinition,
            detail: format!("capability {id} 的 inputSchema.type 必须为 object"),
        });
    }
    Ok(())
}

#[derive(Clone, Debug)]
pub struct CapabilityRegistry<D> {
    definitions: Vec<CapabilityDefinition<D>>,
    by_id: BTreeMap<String, usize>,
    by_function_name: BTreeMap<String, usize>,
}

impl<D> CapabilityRegistry<D> {
    pub fn try_new(
        definitions: impl IntoIterator<Item = CapabilityDefinition<D>>,
    ) -> Result<Self, CapabilityRegistryError> {
        let definitions = definitions.into_iter().collect::<Vec<_>>();
        let mut by_id = BTreeMap::new();
        let mut by_function_name = BTreeMap::new();
        for (index, definition) in definitions.iter().enumerate() {
            if by_id.insert(definition.id.clone(), index).is_some() {
                return Err(CapabilityRegistryError {
                    kind: CapabilityRegistryErrorKind::DuplicateId,
                    detail: format!("capability id 重复：{}", definition.id),
                });
            }
            if by_function_name
                .insert(definition.function_name.clone(), index)
                .is_some()
            {
                return Err(CapabilityRegistryError {
                    kind: CapabilityRegistryErrorKind::DuplicateFunctionName,
                    detail: format!("capability functionName 重复：{}", definition.function_name),
                });
            }
        }
        Ok(Self {
            definitions,
            by_id,
            by_function_name,
        })
    }

    pub fn len(&self) -> usize {
        self.definitions.len()
    }

    pub fn is_empty(&self) -> bool {
        self.definitions.is_empty()
    }

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

    pub fn get_by_function_name(&self, function_name: &str) -> Option<&CapabilityDefinition<D>> {
        self.by_function_name
            .get(function_name)
            .and_then(|index| self.definitions.get(*index))
    }

    pub fn iter(&self) -> impl ExactSizeIterator<Item = &CapabilityDefinition<D>> {
        self.definitions.iter()
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    fn definition(id: &str, function_name: &str) -> CapabilityDefinition<&'static str> {
        CapabilityDefinition::try_new(
            id,
            function_name,
            "测试能力",
            json!({"type": "object", "properties": {}, "additionalProperties": false}),
            "host",
        )
        .expect("definition")
    }

    #[test]
    fn registry_preserves_order_and_supports_both_lookups() {
        let registry = CapabilityRegistry::try_new([
            definition("document.read", "runtime_tool_document_read"),
            definition("document.review", "runtime_tool_document_review"),
        ])
        .expect("registry");
        assert_eq!(
            registry.iter().map(|item| item.id()).collect::<Vec<_>>(),
            vec!["document.read", "document.review"]
        );
        assert_eq!(
            registry
                .get_by_function_name("runtime_tool_document_review")
                .map(CapabilityDefinition::id),
            Some("document.review")
        );
        assert_eq!(
            registry.get("document.read").map(|item| *item.dispatch()),
            Some("host")
        );
    }

    #[test]
    fn registry_rejects_duplicate_id_and_function_name() {
        let duplicate_id = CapabilityRegistry::try_new([
            definition("document.read", "runtime_tool_document_read"),
            definition("document.read", "runtime_tool_document_search"),
        ])
        .expect_err("duplicate id");
        assert_eq!(
            duplicate_id.kind(),
            CapabilityRegistryErrorKind::DuplicateId
        );

        let duplicate_function = CapabilityRegistry::try_new([
            definition("document.read", "runtime_tool_document_read"),
            definition("document.search", "runtime_tool_document_read"),
        ])
        .expect_err("duplicate function");
        assert_eq!(
            duplicate_function.kind(),
            CapabilityRegistryErrorKind::DuplicateFunctionName
        );
    }

    #[test]
    fn definition_rejects_non_object_schema() {
        let error = CapabilityDefinition::try_new(
            "document.read",
            "runtime_tool_document_read",
            "读取文档",
            json!({"type": "string"}),
            (),
        )
        .expect_err("schema must be object");
        assert_eq!(error.kind(), CapabilityRegistryErrorKind::InvalidDefinition);
    }

    #[test]
    fn definition_rejects_invalid_identity_and_empty_description() {
        for (id, function_name, description) in [
            ("../document.read", "runtime_tool_document_read", "读取文档"),
            ("document.read", "runtime-tool-document-read", "读取文档"),
            ("document.read", "runtime_tool_document_read", "  "),
        ] {
            let error = CapabilityDefinition::try_new(
                id,
                function_name,
                description,
                json!({"type": "object"}),
                (),
            )
            .expect_err("invalid definition must fail closed");
            assert_eq!(error.kind(), CapabilityRegistryErrorKind::InvalidDefinition);
        }
    }
}