mars_agents/models/probes/
opencode.rs

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

mars_agents/models/probes/
opencode.rs
