ai-dispatch 8.99.9

//! CLI adapter for user-defined agents loaded from TOML.
//! Exports: CustomAgentConfig, CustomAgent, parse_config.
//! Deps: serde, toml, serde_json, chrono, std::process::Command, crate::types.
#![allow(dead_code)]

use anyhow::Result;
use chrono::Local;
use serde::Deserialize;
use serde_json::Value;
use std::process::Command;

use super::RunOpts;
use crate::types::*;

#[derive(Debug, Clone, Deserialize)]
pub struct CustomAgentFile {
    pub agent: CustomAgentConfig,
}

#[derive(Debug, Clone, Deserialize)]
pub struct CustomAgentConfig {
    pub id: String,
    pub display_name: String,
    pub command: String,
    #[serde(default = "default_prompt_mode")]
    pub prompt_mode: String,
    #[serde(default)]
    pub prompt_flag: String,
    #[serde(default)]
    pub dir_flag: String,
    #[serde(default)]
    pub model_flag: String,
    #[serde(default)]
    pub output_flag: String,
    #[serde(default)]
    pub fixed_args: Vec<String>,
    #[serde(default)]
    pub streaming: bool,
    #[serde(default = "default_output_format")]
    pub output_format: String,
    #[serde(default)]
    pub capabilities: CapabilityScores,
    /// Trust tier: "local" (runs locally), "api" (sends prompts to third-party API)
    #[serde(default = "default_trust_tier")]
    pub trust_tier: String,
    #[serde(default)]
    pub strengths: Vec<String>,
    /// When set (currently only "opencode"), the registry returns an
    /// adapter overlay around the named agent instead of the bash-wrapper
    /// CustomAgent. Required `forced_model` accompanies it.
    #[serde(default)]
    pub delegate_to: Option<String>,
    #[serde(default)]
    pub forced_model: Option<String>,
}

fn default_trust_tier() -> String {
    "api".to_string()
}

#[derive(Debug, Clone, Default, Deserialize)]
pub struct CapabilityScores {
    #[serde(default)]
    pub research: i32,
    #[serde(default = "default_score")]
    pub simple_edit: i32,
    #[serde(default = "default_score")]
    pub complex_impl: i32,
    #[serde(default)]
    pub frontend: i32,
    #[serde(default = "default_score")]
    pub debugging: i32,
    #[serde(default = "default_score")]
    pub testing: i32,
    #[serde(default = "default_score")]
    pub refactoring: i32,
    #[serde(default)]
    pub documentation: i32,
}

fn default_prompt_mode() -> String {
    "arg".to_string()
}

fn default_output_format() -> String {
    "text".to_string()
}

fn default_score() -> i32 {
    3
}

pub struct CustomAgent {
    pub config: CustomAgentConfig,
}

impl super::Agent for CustomAgent {
    fn kind(&self) -> AgentKind {
        AgentKind::Custom
    }

    fn streaming(&self) -> bool {
        self.config.streaming
    }

    fn build_command(&self, prompt: &str, opts: &RunOpts) -> Result<Command> {
        let effective_prompt = apply_read_only_prefix(prompt, opts);
        if opts.read_only {
            aid_warn!("[aid] ⚠ Custom agent read-only is prompt-level only, not enforced. Use --worktree for isolation.");
        }
        let mut cmd = Command::new(&self.config.command);

        for arg in &self.config.fixed_args {
            cmd.arg(arg);
        }

        if let Some(ref dir) = opts.dir {
            if !self.config.dir_flag.is_empty() {
                cmd.args([&self.config.dir_flag, dir]);
            }
            cmd.current_dir(dir);
        }

        if let Some(ref model) = opts.model
            && !self.config.model_flag.is_empty()
        {
            cmd.args([&self.config.model_flag, model]);
        }

        if let Some(ref output) = opts.output
            && !self.config.output_flag.is_empty()
        {
            cmd.args([&self.config.output_flag, output]);
        }

        match self.config.prompt_mode.as_str() {
            "flag" if !self.config.prompt_flag.is_empty() => {
                cmd.args([&self.config.prompt_flag, &effective_prompt]);
            }
            "flag" => {
                cmd.arg(&effective_prompt);
            }
            "stdin" => {}
            _ => {
                cmd.arg(&effective_prompt);
            }
        }

        Ok(cmd)
    }

    fn parse_event(&self, task_id: &TaskId, line: &str) -> Option<TaskEvent> {
        let trimmed = line.trim();
        if trimmed.is_empty() {
            return None;
        }

        if self.config.output_format == "jsonl"
            && let Ok(value) = serde_json::from_str::<Value>(trimmed)
        {
            let event_type = value
                .get("type")
                .or_else(|| value.get("event"))
                .or_else(|| value.get("kind"))
                .and_then(|v| v.as_str());
            if let Some(et) = event_type {
                let detail = value
                    .get("message")
                    .or_else(|| value.get("text"))
                    .or_else(|| value.get("detail"))
                    .and_then(|v| v.as_str())
                    .unwrap_or(et);
                let kind = match et {
                    t if t.contains("error") => EventKind::Error,
                    t if t.contains("tool") => EventKind::ToolCall,
                    t if t.contains("complet") => EventKind::Completion,
                    _ => EventKind::Reasoning,
                };
                return Some(TaskEvent {
                    task_id: task_id.clone(),
                    timestamp: Local::now(),
                    event_kind: kind,
                    detail: super::truncate::truncate_text(detail, 120),
                    metadata: None,
                });
            }
        }

        Some(TaskEvent {
            task_id: task_id.clone(),
            timestamp: Local::now(),
            event_kind: EventKind::Reasoning,
            detail: super::truncate::truncate_text(trimmed, 120),
            metadata: None,
        })
    }

    fn parse_completion(&self, _output: &str) -> CompletionInfo {
        CompletionInfo {
            tokens: None,
            status: TaskStatus::Done,
            model: None,
            cost_usd: None,
            exit_code: None,
        }
    }
}

pub fn parse_config(toml_content: &str) -> Result<CustomAgentConfig> {
    let file: CustomAgentFile = toml::from_str(toml_content)?;
    Ok(file.agent)
}

fn apply_read_only_prefix(prompt: &str, opts: &RunOpts) -> String {
    if !opts.read_only {
        return prompt.to_string();
    }
    if opts.result_file.is_some() {
        format!(
            "IMPORTANT: READ-ONLY MODE. Do NOT modify, create, or delete any files, EXCEPT the result file specified in this prompt. Only read, analyze, and write your findings to the designated result file.\n\n{}",
            prompt
        )
    } else {
        format!(
            "IMPORTANT: READ-ONLY MODE. Do NOT modify, create, or delete any files. Only read and analyze.\n\n{}",
            prompt
        )
    }
}

#[cfg(test)]
mod tests {
    use super::super::Agent;
    use super::*;
    use crate::types::{EventKind, TaskId};

    fn base_opts() -> RunOpts {
        RunOpts {
            dir: None,
            output: None,
            result_file: None,
            model: None,
            budget: false,
            read_only: false,
            context_files: Vec::new(),
            session_id: None,
            env: None,
            env_forward: None,
        }
    }

    fn base_config(command: &str) -> CustomAgentConfig {
        CustomAgentConfig {
            id: "custom".into(),
            display_name: "Custom Agent".into(),
            command: command.into(),
            prompt_mode: default_prompt_mode(),
            prompt_flag: String::new(),
            dir_flag: String::new(),
            model_flag: String::new(),
            output_flag: String::new(),
            fixed_args: Vec::new(),
            streaming: false,
            output_format: default_output_format(),
            capabilities: CapabilityScores::default(),
            trust_tier: default_trust_tier(),
            strengths: Vec::new(),
            delegate_to: None,
            forced_model: None,
        }
    }

    #[test]
    fn parse_minimal_config() {
        let toml_data = r#"
            [agent]
            id = "aider"
            display_name = "Aider"
            command = "aider"
        "#;
        let config = parse_config(toml_data).unwrap();
        assert_eq!(config.id, "aider");
        assert_eq!(config.prompt_mode, "arg");
        assert_eq!(config.output_format, "text");
        assert_eq!(config.capabilities.simple_edit, 0);
    }

    #[test]
    fn parse_full_config() {
        let toml_data = r#"
            [agent]
            id = "walker"
            display_name = "Walker"
            command = "walker-cli"
            prompt_mode = "flag"
            prompt_flag = "--input"
            dir_flag = "--dir"
            model_flag = "--model"
            output_flag = "--out"
            fixed_args = ["--yes", "--batch"]
            streaming = true
            output_format = "jsonl"

            [agent.capabilities]
            research = 2
            simple_edit = 4
            complex_impl = 8
            frontend = 1
            debugging = 5
            testing = 6
            refactoring = 7
            documentation = 2
        "#;
        let config = parse_config(toml_data).unwrap();
        assert!(config.streaming);
        assert_eq!(config.prompt_mode, "flag");
        assert_eq!(config.capabilities.complex_impl, 8);
        assert_eq!(config.fixed_args.len(), 2);
    }

    #[test]
    fn build_command_with_arg_mode() {
        let mut config = base_config("agent-cli");
        config.fixed_args.push("--yes".into());
        let agent = CustomAgent { config };
        let cmd = agent.build_command("ask", &base_opts()).unwrap();
        let args: Vec<_> = cmd
            .get_args()
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["--yes", "ask"]);
    }

    #[test]
    fn build_command_with_flag_mode() {
        let mut config = base_config("agent-cli");
        config.prompt_mode = "flag".into();
        config.prompt_flag = "--message".into();
        config.fixed_args.push("--ready".into());
        let agent = CustomAgent { config };
        let args: Vec<_> = agent
            .build_command("prompt", &base_opts())
            .unwrap()
            .get_args()
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["--ready", "--message", "prompt"]);
    }

    #[test]
    fn custom_agent_kind_is_custom() {
        let agent = CustomAgent {
            config: base_config("agent-cli"),
        };
        assert_eq!(agent.kind(), AgentKind::Custom);
    }

    #[test]
    fn build_command_read_only_prepends_warning_to_prompt() {
        let agent = CustomAgent {
            config: base_config("agent-cli"),
        };
        let mut opts = base_opts();
        opts.read_only = true;
        let cmd = agent.build_command("ask", &opts).unwrap();
        let args: Vec<_> = cmd
            .get_args()
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect();
        assert!(args.last().unwrap().contains("READ-ONLY MODE"));
        assert!(args.last().unwrap().ends_with("\n\nask"));
    }

    #[test]
    fn build_command_read_only_with_result_file_uses_exception_text() {
        let agent = CustomAgent {
            config: base_config("agent-cli"),
        };
        let mut opts = base_opts();
        opts.read_only = true;
        opts.result_file = Some("result.md".into());
        let cmd = agent.build_command("audit", &opts).unwrap();
        let last = cmd
            .get_args()
            .last()
            .map(|arg| arg.to_string_lossy().into_owned())
            .unwrap();
        assert!(last.contains("EXCEPT the result file"));
    }

    #[test]
    fn build_command_normal_mode_does_not_mutate_prompt() {
        let agent = CustomAgent {
            config: base_config("agent-cli"),
        };
        let cmd = agent.build_command("hello", &base_opts()).unwrap();
        let args: Vec<_> = cmd
            .get_args()
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["hello"]);
    }

    #[test]
    fn build_command_with_dir() {
        let mut config = base_config("agent-cli");
        config.dir_flag = "--dir".into();
        let mut opts = base_opts();
        opts.dir = Some("/tmp/work".into());
        let cmd = CustomAgent { config }
            .build_command("prompt", &opts)
            .unwrap();
        let args: Vec<_> = cmd
            .get_args()
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect();
        assert_eq!(&args[..2], ["--dir".to_string(), "/tmp/work".to_string()]);
        let current_dir = cmd
            .get_current_dir()
            .map(|p| p.to_string_lossy().into_owned());
        assert_eq!(current_dir, Some("/tmp/work".to_string()));
    }

    #[test]
    fn parse_event_jsonl() {
        let mut config = base_config("agent-cli");
        config.output_format = "jsonl".into();
        let agent = CustomAgent { config };
        let task_id = TaskId("t-0001".into());
        let line = r#"{"type":"completion","message":"done"}"#;
        let event = agent.parse_event(&task_id, line).unwrap();
        assert_eq!(event.event_kind, EventKind::Completion);
        assert_eq!(event.detail, "done");
    }

    #[test]
    fn parse_event_text() {
        let config = base_config("agent-cli");
        let agent = CustomAgent { config };
        let task_id = TaskId("t-0002".into());
        let event = agent.parse_event(&task_id, " step ").unwrap();
        assert_eq!(event.event_kind, EventKind::Reasoning);
        assert_eq!(event.detail, "step");
    }
}