mars_agents/models/probes/
opencode.rs1use std::collections::HashMap;
2use std::io::Read;
3use std::path::PathBuf;
4use std::process::{Command, Stdio};
5use std::thread;
6use std::time::{Duration, Instant};
7
8use serde::{Deserialize, Serialize};
9use wait_timeout::ChildExt;
10
11#[derive(Debug, Clone, Default, Serialize, Deserialize)]
13pub struct OpenCodeProbeResult {
14 pub providers: HashMap<String, bool>,
16 pub model_slugs: Vec<String>,
18 pub provider_probe_success: bool,
20 pub model_probe_success: bool,
22 pub error: Option<String>,
24}
25
26const DEFAULT_PROBE_TIMEOUT_SECS: u64 = 5;
27
28pub fn probe() -> OpenCodeProbeResult {
30 probe_with_timeout(probe_timeout())
31}
32
33pub fn probe_with_timeout(timeout: Duration) -> OpenCodeProbeResult {
35 let deadline = Instant::now() + timeout;
36 let mut result = OpenCodeProbeResult::default();
37
38 match run_command("opencode", &["providers", "list"], timeout) {
39 Ok(stdout) => {
40 result.providers = parse_providers_output(&stdout);
41 result.provider_probe_success = true;
42 }
43 Err(error) => {
44 result.error = Some(format!("provider probe failed: {error}"));
45 result.provider_probe_success = false;
46 }
47 }
48
49 let remaining = deadline.saturating_duration_since(Instant::now());
50 if remaining.is_zero() {
51 result.model_probe_success = false;
52 if result.error.is_none() {
53 result.error = Some("timeout before model probe".to_string());
54 }
55 return result;
56 }
57
58 match run_command("opencode", &["models"], remaining) {
59 Ok(stdout) => {
60 result.model_slugs = parse_models_output(&stdout);
61 result.model_probe_success = true;
62 }
63 Err(error) => {
64 result.model_probe_success = false;
65 if result.error.is_none() {
66 result.error = Some(format!("model probe failed: {error}"));
67 }
68 }
69 }
70
71 result
72}
73
74fn probe_timeout() -> Duration {
75 std::env::var("MARS_PROBE_TIMEOUT_SECS")
76 .ok()
77 .and_then(|value| value.parse::<u64>().ok())
78 .map(Duration::from_secs)
79 .unwrap_or(Duration::from_secs(DEFAULT_PROBE_TIMEOUT_SECS))
80}
81
82fn run_command(cmd: &str, args: &[&str], timeout: Duration) -> Result<String, String> {
83 let program = resolve_command(cmd);
84 let mut child = Command::new(&program)
85 .args(args)
86 .stdout(Stdio::piped())
87 .stderr(Stdio::null())
88 .spawn()
89 .map_err(|error| format!("spawn failed: {error}"))?;
90
91 let stdout = child
92 .stdout
93 .take()
94 .ok_or_else(|| "stdout capture unavailable".to_string())?;
95 let stdout_reader = thread::spawn(move || {
96 let mut stdout = stdout;
97 let mut output = Vec::new();
98 stdout
99 .read_to_end(&mut output)
100 .map(|_| output)
101 .map_err(|error| format!("stdout read failed: {error}"))
102 });
103
104 match child
105 .wait_timeout(timeout)
106 .map_err(|error| format!("wait failed: {error}"))?
107 {
108 Some(status) if status.success() => {
109 let stdout = stdout_reader
110 .join()
111 .map_err(|_| "stdout reader panicked".to_string())??;
112 String::from_utf8(stdout).map_err(|error| format!("invalid utf8: {error}"))
113 }
114 Some(status) => {
115 let _ = stdout_reader.join();
116 Err(format!("exit code {}", status.code().unwrap_or(-1)))
117 }
118 None => {
119 let _ = child.kill();
120 let _ = child.wait();
121 let _ = stdout_reader.join();
122 Err("timeout".to_string())
123 }
124 }
125}
126
127fn resolve_command(cmd: &str) -> PathBuf {
128 if let Ok(path) = which::which(cmd) {
129 return path;
130 }
131
132 #[cfg(windows)]
133 {
134 for ext in ["exe", "cmd", "bat"] {
135 if let Ok(path) = which::which(format!("{cmd}.{ext}")) {
136 return path;
137 }
138 }
139 }
140
141 PathBuf::from(cmd)
142}
143
144fn strip_ansi(s: &str) -> String {
145 let mut result = String::with_capacity(s.len());
146 let mut chars = s.chars().peekable();
147
148 while let Some(ch) = chars.next() {
149 if ch == '\x1b' {
150 while let Some(&next) = chars.peek() {
151 chars.next();
152 if next.is_ascii_alphabetic() {
153 break;
154 }
155 }
156 } else {
157 result.push(ch);
158 }
159 }
160
161 result
162}
163
164fn parse_providers_output(stdout: &str) -> HashMap<String, bool> {
169 let mut providers = HashMap::new();
170
171 for line in stdout.lines() {
172 let clean = strip_ansi(line.trim());
173 if let Some(rest) = clean.strip_prefix('●').or_else(|| clean.strip_prefix('*')) {
174 let parts: Vec<&str> = rest.split_whitespace().collect();
175 if parts.len() >= 2 {
176 providers.insert(parts[0].to_lowercase(), true);
177 }
178 }
179 }
180
181 providers
182}
183
184fn parse_models_output(stdout: &str) -> Vec<String> {
186 stdout
187 .lines()
188 .map(|line| strip_ansi(line.trim()))
189 .filter(|line| !line.is_empty() && line.contains('/'))
190 .collect()
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196
197 #[test]
198 fn test_strip_ansi_basic() {
199 let input = "\x1b[32mGreen\x1b[0m";
200 assert_eq!(strip_ansi(input), "Green");
201 }
202
203 #[test]
204 fn test_strip_ansi_no_escapes() {
205 assert_eq!(strip_ansi("Plain text"), "Plain text");
206 }
207
208 #[test]
209 fn test_parse_providers_bullet() {
210 let output = r#"┌ Credentials [path redacted]
211│
212● OpenAI oauth
213│
214● Google api
215│
216● OpenRouter api
217│
218└ 3 credentials"#;
219
220 let providers = parse_providers_output(output);
221
222 assert!(providers.get("openai").copied().unwrap_or(false));
223 assert!(providers.get("google").copied().unwrap_or(false));
224 assert!(providers.get("openrouter").copied().unwrap_or(false));
225 assert!(!providers.contains_key("credentials"));
226 }
227
228 #[test]
229 fn test_parse_providers_empty() {
230 assert!(parse_providers_output("").is_empty());
231 }
232
233 #[test]
234 fn test_parse_models_basic() {
235 let output = r#"opencode/big-pickle
236google/gemini-2.5-pro
237openai/gpt-5.4
238openrouter/anthropic/claude-opus-4.7"#;
239
240 let slugs = parse_models_output(output);
241
242 assert_eq!(slugs.len(), 4);
243 assert!(slugs.contains(&"openai/gpt-5.4".to_string()));
244 assert!(slugs.contains(&"openrouter/anthropic/claude-opus-4.7".to_string()));
245 }
246
247 #[test]
248 fn test_parse_models_filters_invalid() {
249 let slugs = parse_models_output("some-invalid-line\nopenai/gpt-5.4\n\n");
250 assert_eq!(slugs, vec!["openai/gpt-5.4"]);
251 }
252
253 #[test]
254 fn test_parse_models_strips_ansi() {
255 let slugs = parse_models_output("\x1b[32mopenai/gpt-5.4\x1b[0m");
256 assert_eq!(slugs, vec!["openai/gpt-5.4"]);
257 }
258
259 #[test]
260 fn test_probe_result_round_trip() {
261 let result = OpenCodeProbeResult {
262 providers: HashMap::from([("openai".to_string(), true)]),
263 model_slugs: vec!["openai/gpt-5.4".to_string()],
264 provider_probe_success: true,
265 model_probe_success: true,
266 error: None,
267 };
268 let json = serde_json::to_string(&result).unwrap();
269 let back: OpenCodeProbeResult = serde_json::from_str(&json).unwrap();
270 assert_eq!(back.providers.get("openai"), Some(&true));
271 assert_eq!(back.model_slugs, result.model_slugs);
272 assert!(back.provider_probe_success);
273 assert!(back.model_probe_success);
274 assert_eq!(back.error, None);
275 }
276}