PicoBot/src/agent/definition.rs

use std::path::{Path, PathBuf};
use std::sync::Arc;

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::config::LLMProviderConfig;

pub const MAX_DEFINITION_FILE_BYTES: u64 = 256 * 1024;
const MAX_DESCRIPTION_CHARS: usize = 4_096;
const MAX_ROLE_CHARS: usize = 65_536;

#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct AgentLimits {
    pub timeout_secs: u64,
    pub max_iterations: usize,
    pub max_children: usize,
    pub max_depth: u16,
    pub max_concurrent_runs: usize,
    pub max_concurrent_provider_steps: usize,
    pub max_concurrent_tool_steps: usize,
    pub max_result_chars: usize,
}

impl Default for AgentLimits {
    fn default() -> Self {
        Self {
            timeout_secs: 900,
            max_iterations: 24,
            max_children: 4,
            max_depth: 3,
            max_concurrent_runs: 2,
            max_concurrent_provider_steps: 1,
            max_concurrent_tool_steps: 4,
            max_result_chars: 16_000,
        }
    }
}

/// Signal delivery lane for background Agent events.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum SignalDelivery {
    #[default]
    Queue,
    Steer,
}

impl SignalDelivery {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Queue => "queue",
            Self::Steer => "steer",
        }
    }
}

/// Durable emit_signal contract.  An Agent only sees the `emit_signal` tool
/// when this block is present; every limit below is enforced by the tool and
/// the Coordinator, not by the model.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct SignalContract {
    pub delivery: SignalDelivery,
    /// Total number of signals one run may emit.
    pub max_total: u32,
    /// Upper bound for the serialized `details` JSON payload of one signal.
    pub max_details_bytes: usize,
    /// Minimum wall time between two signals of the same run.
    pub min_interval_ms: u64,
    /// Maximum signals emitted within `burst_window_ms`.
    pub max_burst: u32,
    pub burst_window_ms: u64,
    /// Allowlisted severities; anything else is rejected.
    pub severity_allowlist: Vec<String>,
    /// Dedupe key cooldown window; repeated keys inside the window collapse
    /// to the same event, keys outside it emit again.
    pub dedupe_cooldown_ms: u64,
    /// Maximum JSON nesting depth of `details`.
    pub max_payload_depth: usize,
}

impl Default for SignalContract {
    fn default() -> Self {
        Self {
            delivery: SignalDelivery::Queue,
            max_total: 64,
            max_details_bytes: 8 * 1024,
            min_interval_ms: 500,
            max_burst: 5,
            burst_window_ms: 10_000,
            severity_allowlist: vec![
                "info".to_string(),
                "warning".to_string(),
                "critical".to_string(),
            ],
            dedupe_cooldown_ms: 60_000,
            max_payload_depth: 16,
        }
    }
}

impl SignalContract {
    fn validate(&self) -> Result<(), AgentDefinitionError> {
        if self.max_total == 0 || self.max_total > 1024 {
            return Err(AgentDefinitionError::Invalid(
                "signal.max_total must be between 1 and 1024".to_string(),
            ));
        }
        if self.max_details_bytes == 0 || self.max_details_bytes > 64 * 1024 {
            return Err(AgentDefinitionError::Invalid(
                "signal.max_details_bytes must be between 1 and 65536".to_string(),
            ));
        }
        if self.max_burst == 0 || self.max_burst > self.max_total {
            return Err(AgentDefinitionError::Invalid(
                "signal.max_burst must be between 1 and max_total".to_string(),
            ));
        }
        if self.max_payload_depth == 0 || self.max_payload_depth > 64 {
            return Err(AgentDefinitionError::Invalid(
                "signal.max_payload_depth must be between 1 and 64".to_string(),
            ));
        }
        if self.severity_allowlist.is_empty() || self.severity_allowlist.len() > 16 {
            return Err(AgentDefinitionError::Invalid(
                "signal.severity_allowlist must contain 1..=16 severities".to_string(),
            ));
        }
        let mut seen = std::collections::HashSet::new();
        if let Some(duplicate) = self
            .severity_allowlist
            .iter()
            .find(|severity| !seen.insert(severity.as_str()))
        {
            return Err(AgentDefinitionError::Invalid(format!(
                "duplicate signal severity '{duplicate}'"
            )));
        }
        if self
            .severity_allowlist
            .iter()
            .any(|severity| severity.trim().is_empty() || severity.len() > 64)
        {
            return Err(AgentDefinitionError::Invalid(
                "signal severities must be non-empty and at most 64 characters".to_string(),
            ));
        }
        Ok(())
    }
}

impl AgentLimits {
    fn validate(&self) -> Result<(), AgentDefinitionError> {
        if self.timeout_secs == 0 || self.timeout_secs > 86_400 {
            return Err(AgentDefinitionError::Invalid(
                "limits.timeout_secs must be between 1 and 86400".to_string(),
            ));
        }
        if self.max_iterations == 0 || self.max_iterations > 256 {
            return Err(AgentDefinitionError::Invalid(
                "limits.max_iterations must be between 1 and 256".to_string(),
            ));
        }
        if self.max_children == 0 || self.max_children > 128 {
            return Err(AgentDefinitionError::Invalid(
                "limits.max_children must be between 1 and 128".to_string(),
            ));
        }
        if self.max_depth == 0 || self.max_depth > 32 {
            return Err(AgentDefinitionError::Invalid(
                "limits.max_depth must be between 1 and 32".to_string(),
            ));
        }
        for (name, value, hard_max) in [
            ("max_concurrent_runs", self.max_concurrent_runs, 128),
            (
                "max_concurrent_provider_steps",
                self.max_concurrent_provider_steps,
                128,
            ),
            (
                "max_concurrent_tool_steps",
                self.max_concurrent_tool_steps,
                512,
            ),
        ] {
            if value == 0 || value > hard_max {
                return Err(AgentDefinitionError::Invalid(format!(
                    "limits.{name} must be between 1 and {hard_max}"
                )));
            }
        }
        if self.max_result_chars == 0 || self.max_result_chars > 1_000_000 {
            return Err(AgentDefinitionError::Invalid(
                "limits.max_result_chars must be between 1 and 1000000".to_string(),
            ));
        }
        Ok(())
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AgentFrontmatter {
    pub id: String,
    pub description: String,
    /// Either this (a key in `config.json`'s `agents` map) or the inline
    /// `provider` + `model` pair must be present.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub llm_profile: Option<String>,
    /// Inline provider/model selection; the preferred way to author an Agent
    /// from the WebUI.  Overrides `llm_profile` when both are present.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub provider: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub token_limit: Option<usize>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_tool_iterations: Option<usize>,
    /// Disabled definitions stay on disk but never load into the catalog.
    #[serde(default = "default_true")]
    pub enabled: bool,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tools: Vec<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub delegates: Vec<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub skills: Vec<String>,
    #[serde(default)]
    pub limits: AgentLimits,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub signal: Option<SignalContract>,
}

pub fn default_true() -> bool {
    true
}

#[derive(Debug, Clone)]
pub struct AgentDefinition {
    pub id: String,
    pub description: String,
    pub llm_profile: Option<String>,
    pub provider: Option<String>,
    pub model: Option<String>,
    pub provider_config: Arc<LLMProviderConfig>,
    pub enabled: bool,
    pub tools: Vec<String>,
    pub delegates: Vec<String>,
    pub skills: Vec<String>,
    pub limits: AgentLimits,
    pub signal_contract: Option<SignalContract>,
    pub role_prompt: String,
    pub definition_hash: String,
    pub source_path: PathBuf,
}

#[derive(Debug, thiserror::Error)]
pub enum AgentDefinitionError {
    #[error("failed to read Agent definition {path}: {source}")]
    Io {
        path: String,
        #[source]
        source: std::io::Error,
    },
    #[error("invalid Agent definition: {0}")]
    Invalid(String),
    #[error("invalid YAML frontmatter: {0}")]
    Yaml(#[from] serde_yaml::Error),
}

pub(crate) fn parse_definition(
    path: &Path,
    provider_config: Arc<LLMProviderConfig>,
) -> Result<AgentDefinition, AgentDefinitionError> {
    let (frontmatter, role_prompt) = read_frontmatter(path)?;
    let canonical_frontmatter = serde_json::to_vec(&frontmatter)
        .map_err(|error| AgentDefinitionError::Invalid(error.to_string()))?;
    let mut hasher = Sha256::new();
    hasher.update(canonical_frontmatter);
    hasher.update(b"\n---\n");
    hasher.update(role_prompt.as_bytes());
    let definition_hash = hasher
        .finalize()
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect();

    Ok(AgentDefinition {
        id: frontmatter.id,
        description: frontmatter.description.trim().to_string(),
        llm_profile: frontmatter.llm_profile,
        provider: frontmatter.provider,
        model: frontmatter.model,
        provider_config,
        enabled: frontmatter.enabled,
        tools: frontmatter.tools,
        delegates: frontmatter.delegates,
        skills: frontmatter.skills,
        limits: frontmatter.limits,
        signal_contract: frontmatter.signal,
        role_prompt,
        definition_hash,
        source_path: path.to_path_buf(),
    })
}

/// Provider-agnostic view of a definition file, used by the management UI.
/// It carries the frontmatter plus the role body but no resolved
/// `provider_config`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentDefinitionInfo {
    #[serde(flatten)]
    pub frontmatter: AgentFrontmatter,
    pub role_prompt: String,
}

/// Read and validate a definition file without resolving its provider
/// config.  Used by the management API to list/validate definitions; the
/// catalog performs the full provider/tool/delegate resolution on load.
pub fn parse_definition_info(path: &Path) -> Result<AgentDefinitionInfo, AgentDefinitionError> {
    let (frontmatter, role_prompt) = read_frontmatter(path)?;
    Ok(AgentDefinitionInfo {
        frontmatter,
        role_prompt,
    })
}

/// Serialize a definition back to the Markdown file format.
pub fn serialize_definition(info: &AgentDefinitionInfo) -> String {
    let yaml = serde_yaml::to_string(&info.frontmatter).unwrap_or_default();
    format!("---\n{yaml}---\n{}\n", info.role_prompt.trim_end())
}

/// Read + validate the frontmatter and role body of a definition file,
/// shared by `parse_definition` and `parse_definition_info`.
fn read_frontmatter(path: &Path) -> Result<(AgentFrontmatter, String), AgentDefinitionError> {
    let metadata = std::fs::symlink_metadata(path).map_err(|source| AgentDefinitionError::Io {
        path: path.display().to_string(),
        source,
    })?;
    if !metadata.file_type().is_file() || metadata.file_type().is_symlink() {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} must be a regular non-symlink file",
            path.display()
        )));
    }
    if metadata.len() > MAX_DEFINITION_FILE_BYTES {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} exceeds the {} byte limit",
            path.display(),
            MAX_DEFINITION_FILE_BYTES
        )));
    }
    let content = std::fs::read_to_string(path).map_err(|source| AgentDefinitionError::Io {
        path: path.display().to_string(),
        source,
    })?;
    let normalized = content.replace("\r\n", "\n");
    let mut lines = normalized.lines();
    if lines.next() != Some("---") {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} must start with a standalone --- line",
            path.display()
        )));
    }
    let mut yaml_lines = Vec::new();
    let mut found_end = false;
    for line in &mut lines {
        if line == "---" {
            found_end = true;
            break;
        }
        yaml_lines.push(line);
    }
    if !found_end {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} has no closing frontmatter delimiter",
            path.display()
        )));
    }
    let role_prompt = lines.collect::<Vec<_>>().join("\n").trim().to_string();
    if role_prompt.is_empty() {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} has an empty role body",
            path.display()
        )));
    }
    if role_prompt.chars().count() > MAX_ROLE_CHARS {
        return Err(AgentDefinitionError::Invalid(format!(
            "{} role body exceeds {MAX_ROLE_CHARS} characters",
            path.display()
        )));
    }
    let frontmatter: AgentFrontmatter = serde_yaml::from_str(&yaml_lines.join("\n"))?;
    validate_agent_id(&frontmatter.id)?;
    if frontmatter.description.trim().is_empty()
        || frontmatter.description.chars().count() > MAX_DESCRIPTION_CHARS
    {
        return Err(AgentDefinitionError::Invalid(format!(
            "Agent '{}' description must contain 1..={MAX_DESCRIPTION_CHARS} characters",
            frontmatter.id
        )));
    }
    let has_profile = frontmatter
        .llm_profile
        .as_deref()
        .is_some_and(|p| !p.trim().is_empty());
    let has_inline = frontmatter.provider.is_some() || frontmatter.model.is_some();
    if !has_profile && !has_inline {
        return Err(AgentDefinitionError::Invalid(format!(
            "Agent '{}' must declare either llm_profile or provider+model",
            frontmatter.id
        )));
    }
    if frontmatter.provider.is_some() != frontmatter.model.is_some() {
        return Err(AgentDefinitionError::Invalid(format!(
            "Agent '{}' must declare provider and model together",
            frontmatter.id
        )));
    }
    frontmatter.limits.validate()?;
    if let Some(signal) = frontmatter.signal.as_ref() {
        signal.validate()?;
    }
    reject_duplicates("tools", &frontmatter.tools)?;
    reject_duplicates("delegates", &frontmatter.delegates)?;
    reject_duplicates("skills", &frontmatter.skills)?;
    let file_stem = path.file_stem().and_then(|value| value.to_str());
    if file_stem != Some(frontmatter.id.as_str()) {
        return Err(AgentDefinitionError::Invalid(format!(
            "Agent id '{}' must match file name {}",
            frontmatter.id,
            path.display()
        )));
    }
    Ok((frontmatter, role_prompt))
}

pub fn validate_agent_id(id: &str) -> Result<(), AgentDefinitionError> {
    let valid = !id.is_empty()
        && id.len() <= 64
        && id
            .bytes()
            .next()
            .is_some_and(|value| value.is_ascii_lowercase())
        && id.bytes().all(|value| {
            value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'_' || value == b'-'
        });
    if !valid || matches!(id, "root" | "main" | "default" | "general") {
        return Err(AgentDefinitionError::Invalid(format!(
            "invalid or reserved Agent id '{id}'"
        )));
    }
    Ok(())
}

fn reject_duplicates(field: &str, values: &[String]) -> Result<(), AgentDefinitionError> {
    let mut seen = std::collections::HashSet::new();
    if let Some(value) = values.iter().find(|value| !seen.insert(value.as_str())) {
        return Err(AgentDefinitionError::Invalid(format!(
            "duplicate {field} entry '{value}'"
        )));
    }
    Ok(())
}

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

    fn provider() -> Arc<LLMProviderConfig> {
        Arc::new(LLMProviderConfig {
            provider_type: "openai".to_string(),
            name: "test".to_string(),
            base_url: "https://example.invalid/v1".to_string(),
            api_key: "test".to_string(),
            extra_headers: HashMap::new(),
            model_id: "test-model".to_string(),
            temperature: None,
            max_tokens: None,
            model_extra: HashMap::new(),
            max_tool_iterations: 99,
            token_limit: 4096,
            workspace_dir: std::env::temp_dir(),
            input_types: vec!["text".to_string()],
            price_input_per_million: None,
            price_output_per_million: None,
        })
    }

    #[test]
    fn strict_definition_parses_and_hashes_stably() {
        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("researcher.md");
        std::fs::write(
            &path,
            concat!(
                "---\n",
                "id: researcher\n",
                "description: Research primary sources\n",
                "llm_profile: research\n",
                "tools:\n",
                "  - calculator\n",
                "limits:\n",
                "  max_iterations: 12\n",
                "---\n",
                "# Role\n\n",
                "Return evidence and uncertainty.\n"
            ),
        )
        .unwrap();

        let first = parse_definition(&path, provider()).unwrap();
        let second = parse_definition(&path, provider()).unwrap();

        assert_eq!(first.id, "researcher");
        assert_eq!(first.limits.max_iterations, 12);
        assert_eq!(first.definition_hash, second.definition_hash);
        assert_eq!(first.definition_hash.len(), 64);
    }

    #[test]
    fn unknown_frontmatter_key_is_rejected() {
        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("researcher.md");
        std::fs::write(
            &path,
            "---\nid: researcher\ndescription: test\nllm_profile: research\nunsafe_tools: true\n---\n# Role\n",
        )
        .unwrap();

        assert!(parse_definition(&path, provider()).is_err());
    }

    #[test]
    fn reserved_and_mismatched_ids_are_rejected() {
        assert!(validate_agent_id("root").is_err());
        assert!(validate_agent_id("Uppercase").is_err());

        let directory = tempfile::tempdir().unwrap();
        let path = directory.path().join("researcher.md");
        std::fs::write(
            &path,
            "---\nid: reviewer\ndescription: test\nllm_profile: research\n---\n# Role\n",
        )
        .unwrap();
        assert!(parse_definition(&path, provider()).is_err());
    }
}