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