Skip to main content

aether_cli/show_prompt/
run.rs

1use std::collections::BTreeMap;
2
3use super::PromptArgs;
4use crate::error::CliError;
5use crate::resolve::resolve_agent_spec;
6use crate::runtime::RuntimeBuilder;
7use aether_core::core::Prompt;
8use llm::ToolDefinition;
9use serde_json::Value;
10
11pub async fn run_prompt(args: PromptArgs) -> Result<(), CliError> {
12    let cwd = args.cwd.canonicalize().map_err(CliError::IoError)?;
13    let catalog = args.settings_source.load_agent_catalog(&cwd).map_err(|e| CliError::AgentError(e.to_string()))?;
14    let spec = resolve_agent_spec(&catalog, args.agent.as_deref())?;
15
16    let info = RuntimeBuilder::from_spec(cwd.clone(), spec)
17        .mcp_sources(args.mcp_config.sources(&cwd))
18        .build_prompt_info()
19        .await?;
20
21    let system_prompt = build_prompt(&info.spec.prompts, args.system_prompt.as_deref()).await?;
22    let tools_output = build_tools(&info.tool_definitions);
23
24    println!("{system_prompt}");
25
26    if !tools_output.is_empty() {
27        println!();
28        println!("--- Tools ({} tools) ---", info.tool_definitions.len());
29        println!();
30        println!("{tools_output}");
31    }
32
33    println!();
34    println!("{}", format_stats(system_prompt.len(), tools_output.len(), info.tool_definitions.len()));
35
36    Ok(())
37}
38
39pub async fn build_prompt(prompts: &[Prompt], custom: Option<&str>) -> Result<String, CliError> {
40    let mut prompts = prompts.to_vec();
41    if let Some(custom) = custom {
42        prompts.push(Prompt::text(custom));
43    }
44    Prompt::build_all(&prompts).await.map_err(|e| CliError::AgentError(e.to_string()))
45}
46
47pub fn build_tools(tools: &[ToolDefinition]) -> String {
48    if tools.is_empty() {
49        return String::new();
50    }
51
52    let mut grouped: BTreeMap<&str, Vec<Value>> = BTreeMap::new();
53    for tool in tools {
54        let server = tool.server.as_deref().unwrap_or("(built-in)");
55        let input_schema = serde_json::from_str::<Value>(&tool.parameters).unwrap_or(Value::Null);
56        let entry = serde_json::json!({
57            "name": tool.name,
58            "description": tool.description,
59            "input_schema": input_schema,
60        });
61        grouped.entry(server).or_default().push(entry);
62    }
63
64    let mut sections = Vec::new();
65    for (server, entries) in &grouped {
66        let json = serde_json::to_string_pretty(entries).unwrap_or_default();
67        sections.push(format!("Server: {server}\n{json}"));
68    }
69
70    sections.join("\n\n")
71}
72
73pub fn format_stats(prompt_chars: usize, tool_schema_chars: usize, tool_count: usize) -> String {
74    let est_tokens = (prompt_chars + tool_schema_chars) / 4;
75    format!(
76        "---\n\
77         Prompt chars:     {prompt_chars:>8}\n\
78         Tool schema chars:{tool_schema_chars:>8}\n\
79         Est. tokens:     ~{est_tokens:>8}\n\
80         MCP tools:        {tool_count:>8}"
81    )
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    fn tool(name: &str, desc: &str, params: &str, server: Option<&str>) -> ToolDefinition {
89        let mut tool = ToolDefinition::new(name, desc, params);
90        tool.server = server.map(String::from);
91        tool
92    }
93
94    #[test]
95    fn format_stats_computes_token_estimate() {
96        let output = format_stats(12000, 8500, 14);
97        assert_eq!(
98            output,
99            "---\n\
100             Prompt chars:        12000\n\
101             Tool schema chars:    8500\n\
102             Est. tokens:     ~    5125\n\
103             MCP tools:              14"
104        );
105    }
106
107    #[test]
108    fn format_stats_handles_zero() {
109        let output = format_stats(0, 0, 0);
110        assert_eq!(
111            output,
112            "---\n\
113             Prompt chars:            0\n\
114             Tool schema chars:       0\n\
115             Est. tokens:     ~       0\n\
116             MCP tools:               0"
117        );
118    }
119
120    #[test]
121    fn format_stats_handles_small_values() {
122        let output = format_stats(3, 0, 1);
123        assert_eq!(
124            output,
125            "---\n\
126             Prompt chars:            3\n\
127             Tool schema chars:       0\n\
128             Est. tokens:     ~       0\n\
129             MCP tools:               1"
130        );
131    }
132
133    #[test]
134    fn build_tools_groups_by_server() {
135        let tools = vec![
136            tool("fs_read", "Read a file", r#"{"type":"object"}"#, Some("filesystem")),
137            tool("git_log", "Show log", r#"{"type":"object"}"#, Some("git")),
138            tool("fs_write", "Write a file", r#"{"type":"object"}"#, Some("filesystem")),
139        ];
140        let output = build_tools(&tools);
141        // BTreeMap sorts: filesystem < git
142        let fs_pos = output.find("Server: filesystem").unwrap();
143        let git_pos = output.find("Server: git").unwrap();
144        assert!(fs_pos < git_pos);
145        // filesystem group has both tools
146        assert!(output.contains("fs_read"));
147        assert!(output.contains("fs_write"));
148    }
149
150    #[test]
151    fn build_tools_handles_no_server() {
152        let tools = vec![tool("builtin_tool", "A built-in", r#"{"type":"object"}"#, None)];
153        let output = build_tools(&tools);
154        assert!(output.contains("Server: (built-in)"));
155        assert!(output.contains("builtin_tool"));
156    }
157
158    #[test]
159    fn build_tools_produces_api_format() {
160        let tools = vec![tool("my_tool", "Does stuff", r#"{"type":"object","properties":{}}"#, Some("test"))];
161        let output = build_tools(&tools);
162        // Strip "Server: test\n" prefix to get the JSON
163        let json_start = output.find('[').unwrap();
164        let parsed: Vec<Value> = serde_json::from_str(&output[json_start..]).unwrap();
165        assert_eq!(parsed.len(), 1);
166        let entry = &parsed[0];
167        assert_eq!(entry["name"], "my_tool");
168        assert_eq!(entry["description"], "Does stuff");
169        assert!(entry["input_schema"].is_object());
170    }
171
172    #[test]
173    fn build_tools_empty() {
174        assert_eq!(build_tools(&[]), "");
175    }
176
177    #[test]
178    fn build_tools_malformed_params() {
179        let tools = vec![tool("bad_tool", "Broken params", "not valid json", Some("srv"))];
180        let output = build_tools(&tools);
181        assert!(output.contains("bad_tool"));
182        assert!(output.contains("null"));
183    }
184}