Skip to main content

mars_agents/models/probes/
opencode.rs

1use 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/// Result of probing OpenCode's runtime availability.
12#[derive(Debug, Clone, Default, Serialize, Deserialize)]
13pub struct OpenCodeProbeResult {
14    /// Provider availability: OpenCode provider ID -> has credentials.
15    pub providers: HashMap<String, bool>,
16    /// Full model slug list, e.g. `openai/gpt-5.4`.
17    pub model_slugs: Vec<String>,
18    /// Whether the provider probe succeeded.
19    pub provider_probe_success: bool,
20    /// Whether the model list probe succeeded.
21    pub model_probe_success: bool,
22    /// Redacted error message if either probe failed.
23    pub error: Option<String>,
24}
25
26const DEFAULT_PROBE_TIMEOUT_SECS: u64 = 5;
27
28/// Probe OpenCode with the configured timeout.
29pub fn probe() -> OpenCodeProbeResult {
30    probe_with_timeout(probe_timeout())
31}
32
33/// Probe OpenCode with a specific timeout.
34pub 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
164/// Parse `opencode providers list` output.
165///
166/// SECURITY: The raw input may reference credential paths. Do not log, store,
167/// or include it in diagnostics.
168fn 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
184/// Parse `opencode models` output into slug list.
185fn 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]
211212●  OpenAI oauth
213214●  Google api
215216●  OpenRouter api
217218└  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}