ai_proxy/
cli.rs

1use clap::{Arg, Command};
2use tokio::process::Command as TokioCommand;
3use tracing::error;
4
5use crate::{ProxyServer, Result};
6
7pub async fn run_cli() -> Result<()> {
8    let matches = Command::new("ai-proxy")
9        .version(env!("CARGO_PKG_VERSION"))
10        .about("AI proxy server for secure API routing")
11        .arg(
12            Arg::new("port")
13                .short('p')
14                .long("port")
15                .value_name("PORT")
16                .help("Port to run on")
17                .default_value("8080"),
18        )
19        .arg(
20            Arg::new("config")
21                .short('c')
22                .long("config")
23                .value_name("PATH")
24                .help("Path to vibekit.yaml file")
25                .default_value("vibekit.yaml"),
26        )
27        .subcommand(
28            Command::new("start")
29                .about("Start the proxy server")
30                .arg(
31                    Arg::new("port")
32                        .short('p')
33                        .long("port")
34                        .value_name("PORT")
35                        .help("Port to run on")
36                        .default_value("8080"),
37                )
38                .arg(
39                    Arg::new("config")
40                        .short('c')
41                        .long("config")
42                        .value_name("PATH")
43                        .help("Path to vibekit.yaml file")
44                        .default_value("vibekit.yaml"),
45                )
46                .arg(
47                    Arg::new("daemon")
48                        .short('d')
49                        .long("daemon")
50                        .help("Run in background")
51                        .action(clap::ArgAction::SetTrue),
52                )
53                .arg(
54                    Arg::new("redaction_api_url")
55                        .long("redaction-api-url")
56                        .value_name("URL")
57                        .help("URL for redaction API to screen user messages")
58                        .required(false),
59                ),
60        )
61        .subcommand(
62            Command::new("stop")
63                .about("Stop the proxy server")
64                .arg(
65                    Arg::new("port")
66                        .short('p')
67                        .long("port")
68                        .value_name("PORT")
69                        .help("Port to stop")
70                        .default_value("8080"),
71                ),
72        )
73        .subcommand(
74            Command::new("status")
75                .about("Show proxy server status")
76                .arg(
77                    Arg::new("port")
78                        .short('p')
79                        .long("port")
80                        .value_name("PORT")
81                        .help("Port to check")
82                        .default_value("8080"),
83                ),
84        )
85        .get_matches();
86
87    match matches.subcommand() {
88        Some(("start", sub_matches)) => {
89            let port: u16 = sub_matches
90                .get_one::<String>("port")
91                .unwrap()
92                .parse()
93                .unwrap_or_else(|_| {
94                    // Check environment variable
95                    std::env::var("PORT")
96                        .ok()
97                        .and_then(|p| p.parse().ok())
98                        .unwrap_or(8080)
99                });
100
101            let config_path = sub_matches
102                .get_one::<String>("config")
103                .unwrap()
104                .clone();
105
106            let redaction_api_url = sub_matches
107                .get_one::<String>("redaction_api_url")
108                .cloned()
109                .or_else(|| std::env::var("VIBEKIT_REDACTION_API_URL").ok());
110
111            let _daemon = sub_matches.get_flag("daemon");
112
113            // Daemon mode (simplified)
114
115            start_proxy_with_check(port, Some(config_path), redaction_api_url).await?;
116        }
117        Some(("stop", sub_matches)) => {
118            let port: u16 = sub_matches
119                .get_one::<String>("port")
120                .unwrap()
121                .parse()
122                .unwrap_or(8080);
123
124            stop_proxy_on_port(port).await;
125        }
126        Some(("status", sub_matches)) => {
127            let port: u16 = sub_matches
128                .get_one::<String>("port")
129                .unwrap()
130                .parse()
131                .unwrap_or(8080);
132
133            show_proxy_status(port).await;
134        }
135        _ => {
136            // Default command - start the proxy server
137            let port: u16 = matches
138                .get_one::<String>("port")
139                .unwrap()
140                .parse()
141                .unwrap_or(8080);
142
143            let config_path = matches
144                .get_one::<String>("config")
145                .unwrap()
146                .clone();
147
148            let redaction_api_url = matches
149                .get_one::<String>("redaction_api_url")
150                .cloned()
151                .or_else(|| std::env::var("VIBEKIT_REDACTION_API_URL").ok());
152
153            start_proxy_with_check(port, Some(config_path), redaction_api_url).await?;
154        }
155    }
156
157    Ok(())
158}
159
160async fn start_proxy_with_check(port: u16, config_path: Option<String>, redaction_api_url: Option<String>) -> Result<()> {
161    use std::sync::Arc;
162    use tokio::signal;
163    
164    let server = Arc::new(ProxyServer::new(port, config_path, redaction_api_url).await?);
165    let server_clone = Arc::clone(&server);
166
167    // Handle graceful shutdown
168    tokio::spawn(async move {
169        let _ = signal::ctrl_c().await;
170        server_clone.stop().await;
171    });
172
173    server.start().await
174}
175
176async fn stop_proxy_on_port(port: u16) {
177    // Kill processes using the port (Unix/Linux/macOS only)
178    if let Ok(output) = TokioCommand::new("lsof")
179        .arg("-ti")
180        .arg(format!(":{}", port))
181        .output()
182        .await
183    {
184        let pids_str = String::from_utf8_lossy(&output.stdout);
185        let pids: Vec<&str> = pids_str.trim().split('\n').filter(|s| !s.is_empty()).collect();
186
187        if pids.is_empty() {
188            return;
189        }
190
191        for pid_str in &pids {
192            if let Ok(pid) = pid_str.parse::<i32>() {
193                let _ = TokioCommand::new("kill")
194                    .arg("-TERM")
195                    .arg(pid.to_string())
196                    .output()
197                    .await;
198            }
199        }
200    } else {
201        error!("Failed to find processes on port {}", port);
202    }
203}
204
205async fn show_proxy_status(port: u16) {
206    println!("🌐 Proxy Server Status");
207    println!("{}", "─".repeat(30));
208
209    let running = is_port_in_use(port).await;
210    println!("Port {}: {}", port, if running { "✅ RUNNING" } else { "❌ NOT RUNNING" });
211
212    if running {
213        // Try to get health check
214        let health_url = format!("http://localhost:{}/health", port);
215        match reqwest::get(&health_url).await {
216            Ok(response) => {
217                if response.status().is_success() {
218                    if let Ok(data) = response.json::<serde_json::Value>().await {
219                        if let Some(uptime) = data.get("uptime").and_then(|u| u.as_u64()) {
220                            println!("Uptime: {}s", uptime);
221                        }
222                        if let Some(request_count) = data.get("requestCount").and_then(|r| r.as_u64()) {
223                            println!("Requests: {}", request_count);
224                        }
225                    }
226                } else {
227                    println!("Health check: ❌ Failed");
228                }
229            }
230            Err(_) => {
231                println!("Health check: ❌ Failed");
232            }
233        }
234
235        // Show process info
236        if let Ok(output) = TokioCommand::new("lsof")
237            .arg("-ti")
238            .arg(format!(":{}", port))
239            .output()
240            .await
241        {
242            let pids_str = String::from_utf8_lossy(&output.stdout);
243            let pids: Vec<&str> = pids_str.trim().split('\n').filter(|s| !s.is_empty()).collect();
244            if !pids.is_empty() {
245                println!("PIDs: {}", pids.join(", "));
246            }
247        }
248    }
249}
250
251async fn is_port_in_use(port: u16) -> bool {
252    let output = TokioCommand::new("lsof")
253        .arg("-ti")
254        .arg(format!(":{}", port))
255        .output()
256        .await;
257
258    if let Ok(output) = output {
259        !String::from_utf8_lossy(&output.stdout).trim().is_empty()
260    } else {
261        false
262    }
263}
ai_proxy/cli.rs

ai_proxy/
cli.rs
