crank/crates/crank-core/src/agent.rs

use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use time::OffsetDateTime;

use crate::ids::{AgentId, OperationId, WorkspaceId};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentStatus {
    Draft,
    Published,
    Archived,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Agent {
    pub id: AgentId,
    pub workspace_id: WorkspaceId,
    pub slug: String,
    pub display_name: String,
    pub description: String,
    pub status: AgentStatus,
    pub current_draft_version: u32,
    pub latest_published_version: Option<u32>,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339")]
    pub updated_at: OffsetDateTime,
    #[serde(default, with = "time::serde::rfc3339::option")]
    pub published_at: Option<OffsetDateTime>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentVersion {
    pub agent_id: AgentId,
    pub version: u32,
    pub status: AgentStatus,
    pub instructions: Value,
    #[serde(default)]
    pub tool_selection_policy: ToolSelectionPolicy,
    #[serde(with = "time::serde::rfc3339")]
    pub created_at: OffsetDateTime,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolAccessMode {
    #[default]
    #[serde(alias = "allow_list")]
    Direct,
    Search,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolGroup {
    pub id: String,
    pub name: String,
    pub description: String,
    #[serde(default)]
    pub tool_names: Vec<String>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolSearchSettings {
    #[serde(default = "default_search_result_limit")]
    pub max_results: usize,
}

impl Default for ToolSearchSettings {
    fn default() -> Self {
        Self {
            max_results: default_search_result_limit(),
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolSelectionPolicy {
    #[serde(default)]
    pub mode: ToolAccessMode,
    #[serde(default)]
    pub groups: Vec<ToolGroup>,
    #[serde(default)]
    pub search: ToolSearchSettings,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum ToolSelectionPolicyError {
    #[error("search result limit must be between 1 and 20")]
    InvalidSearchResultLimit,
    #[error("tool groups are only allowed in search mode")]
    GroupsRequireSearchMode,
    #[error("tool group id {group_id} is invalid")]
    InvalidGroupId { group_id: String },
    #[error("tool group id {group_id} is duplicated")]
    DuplicateGroupId { group_id: String },
    #[error("tool group {group_id} must have a name and description")]
    IncompleteGroup { group_id: String },
    #[error("tool group {group_id} references unknown tool {tool_name}")]
    UnknownTool { group_id: String, tool_name: String },
    #[error("tool group {group_id} contains duplicate tool {tool_name}")]
    DuplicateTool { group_id: String, tool_name: String },
    #[error("tool name {tool_name} is reserved by search mode")]
    ReservedToolName { tool_name: String },
}

impl ToolSelectionPolicy {
    pub fn validate_for_tools<'a>(
        &self,
        tool_names: impl IntoIterator<Item = &'a str>,
    ) -> Result<(), ToolSelectionPolicyError> {
        use std::collections::BTreeSet;

        if !(1..=20).contains(&self.search.max_results) {
            return Err(ToolSelectionPolicyError::InvalidSearchResultLimit);
        }
        if self.mode == ToolAccessMode::Direct && !self.groups.is_empty() {
            return Err(ToolSelectionPolicyError::GroupsRequireSearchMode);
        }

        let known_tools = tool_names.into_iter().collect::<BTreeSet<_>>();
        if self.mode == ToolAccessMode::Search
            && let Some(tool_name) = known_tools
                .iter()
                .find(|tool_name| matches!(**tool_name, "search_tools" | "call_tool"))
        {
            return Err(ToolSelectionPolicyError::ReservedToolName {
                tool_name: (*tool_name).to_owned(),
            });
        }
        let mut group_ids = BTreeSet::new();
        for group in &self.groups {
            if !valid_group_id(&group.id) {
                return Err(ToolSelectionPolicyError::InvalidGroupId {
                    group_id: group.id.clone(),
                });
            }
            if !group_ids.insert(group.id.as_str()) {
                return Err(ToolSelectionPolicyError::DuplicateGroupId {
                    group_id: group.id.clone(),
                });
            }
            if group.name.trim().is_empty() || group.description.trim().is_empty() {
                return Err(ToolSelectionPolicyError::IncompleteGroup {
                    group_id: group.id.clone(),
                });
            }

            let mut grouped_tools = BTreeSet::new();
            for tool_name in &group.tool_names {
                if !known_tools.contains(tool_name.as_str()) {
                    return Err(ToolSelectionPolicyError::UnknownTool {
                        group_id: group.id.clone(),
                        tool_name: tool_name.clone(),
                    });
                }
                if !grouped_tools.insert(tool_name.as_str()) {
                    return Err(ToolSelectionPolicyError::DuplicateTool {
                        group_id: group.id.clone(),
                        tool_name: tool_name.clone(),
                    });
                }
            }
        }

        Ok(())
    }
}

const fn default_search_result_limit() -> usize {
    8
}

fn valid_group_id(value: &str) -> bool {
    !value.is_empty()
        && value.len() <= 64
        && value
            .bytes()
            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentOperationBinding {
    pub operation_id: OperationId,
    pub operation_version: u32,
    pub tool_name: String,
    pub tool_title: String,
    pub tool_description_override: Option<String>,
    pub enabled: bool,
}

#[cfg(test)]
mod tests {
    use serde_json::json;
    use time::{OffsetDateTime, format_description::well_known::Rfc3339};

    use super::{
        Agent, AgentStatus, AgentVersion, ToolAccessMode, ToolGroup, ToolSelectionPolicy,
        ToolSelectionPolicyError,
    };
    use crate::ids::{AgentId, WorkspaceId};

    #[test]
    fn agent_serializes_timestamps_as_rfc3339() {
        let agent = Agent {
            id: AgentId::new("agent_01"),
            workspace_id: WorkspaceId::new("ws_01"),
            slug: "triage".to_owned(),
            display_name: "Triage".to_owned(),
            description: "Triage agent".to_owned(),
            status: AgentStatus::Published,
            current_draft_version: 2,
            latest_published_version: Some(2),
            created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
            updated_at: OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap(),
            published_at: Some(OffsetDateTime::parse("2026-03-25T12:10:00Z", &Rfc3339).unwrap()),
        };

        let value = serde_json::to_value(&agent).unwrap();

        assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
        assert_eq!(value["updated_at"], json!("2026-03-25T12:05:00Z"));
        assert_eq!(value["published_at"], json!("2026-03-25T12:10:00Z"));
    }

    #[test]
    fn agent_version_deserializes_created_at_from_rfc3339() {
        let version: AgentVersion = serde_json::from_value(json!({
            "agent_id": "agent_01",
            "version": 1,
            "status": "draft",
            "instructions": {"system": "triage"},
            "tool_selection_policy": {"mode": "allow_list"},
            "created_at": "2026-03-25T12:00:00Z"
        }))
        .unwrap();

        assert_eq!(
            version.created_at,
            OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap()
        );
        assert_eq!(version.tool_selection_policy.mode, ToolAccessMode::Direct);
    }

    #[test]
    fn search_policy_validates_groups_against_published_tools() {
        let policy = ToolSelectionPolicy {
            mode: ToolAccessMode::Search,
            groups: vec![ToolGroup {
                id: "finance".to_owned(),
                name: "Finance".to_owned(),
                description: "Invoices and payments".to_owned(),
                tool_names: vec!["create_invoice".to_owned()],
            }],
            ..ToolSelectionPolicy::default()
        };

        assert_eq!(
            policy.validate_for_tools(["list_invoices"]),
            Err(ToolSelectionPolicyError::UnknownTool {
                group_id: "finance".to_owned(),
                tool_name: "create_invoice".to_owned(),
            })
        );
    }
}