Skip to main content

mars_agents/models/probes/
opencode.rs

1use std::io::Read;
2use std::path::PathBuf;
3use std::process::{Command, Stdio};
4use std::thread;
5use std::time::Duration;
6
7use serde::{Deserialize, Serialize};
8use wait_timeout::ChildExt;
9
10/// Result of probing OpenCode's runtime availability.
11#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12pub struct OpenCodeProbeResult {
13    /// Full model slug list, e.g. `openai/gpt-5.4`.
14    pub model_slugs: Vec<String>,
15    /// Whether the model list probe succeeded.
16    pub model_probe_success: bool,
17    /// Redacted error message if probing failed.
18    pub error: Option<String>,
19}
20
21const DEFAULT_PROBE_TIMEOUT_SECS: u64 = 5;
22
23/// Probe OpenCode with the configured timeout.
24pub fn probe() -> OpenCodeProbeResult {
25    probe_with_timeout(probe_timeout())
26}
27
28/// Probe OpenCode with a specific timeout.
29pub fn probe_with_timeout(timeout: Duration) -> OpenCodeProbeResult {
30    let mut result = OpenCodeProbeResult::default();
31
32    match run_command("opencode", &["models"], timeout) {
33        Ok(stdout) => {
34            result.model_slugs = parse_models_output(&stdout);
35            result.model_probe_success = true;
36        }
37        Err(error) => {
38            result.model_probe_success = false;
39            result.error = Some(format!("model probe failed: {error}"));
40        }
41    }
42
43    result
44}
45
46fn probe_timeout() -> Duration {
47    std::env::var("MARS_PROBE_TIMEOUT_SECS")
48        .ok()
49        .and_then(|value| value.parse::<u64>().ok())
50        .map(Duration::from_secs)
51        .unwrap_or(Duration::from_secs(DEFAULT_PROBE_TIMEOUT_SECS))
52}
53
54fn run_command(cmd: &str, args: &[&str], timeout: Duration) -> Result<String, String> {
55    let program = resolve_command(cmd);
56    let mut child = Command::new(&program)
57        .args(args)
58        .stdout(Stdio::piped())
59        .stderr(Stdio::null())
60        .spawn()
61        .map_err(|error| format!("spawn failed: {error}"))?;
62
63    let stdout = child
64        .stdout
65        .take()
66        .ok_or_else(|| "stdout capture unavailable".to_string())?;
67    let stdout_reader = thread::spawn(move || {
68        let mut stdout = stdout;
69        let mut output = Vec::new();
70        stdout
71            .read_to_end(&mut output)
72            .map(|_| output)
73            .map_err(|error| format!("stdout read failed: {error}"))
74    });
75
76    match child
77        .wait_timeout(timeout)
78        .map_err(|error| format!("wait failed: {error}"))?
79    {
80        Some(status) if status.success() => {
81            let stdout = stdout_reader
82                .join()
83                .map_err(|_| "stdout reader panicked".to_string())??;
84            String::from_utf8(stdout).map_err(|error| format!("invalid utf8: {error}"))
85        }
86        Some(status) => {
87            let _ = stdout_reader.join();
88            Err(format!("exit code {}", status.code().unwrap_or(-1)))
89        }
90        None => {
91            let _ = child.kill();
92            let _ = child.wait();
93            let _ = stdout_reader.join();
94            Err("timeout".to_string())
95        }
96    }
97}
98
99fn resolve_command(cmd: &str) -> PathBuf {
100    let resolver = crate::harness::host::PathExecutableResolver;
101    crate::harness::host::resolve_binary_path(cmd, &resolver).unwrap_or_else(|| cmd.into())
102}
103
104fn strip_ansi(s: &str) -> String {
105    let mut result = String::with_capacity(s.len());
106    let mut chars = s.chars().peekable();
107
108    while let Some(ch) = chars.next() {
109        if ch == '\x1b' {
110            while let Some(&next) = chars.peek() {
111                chars.next();
112                if next.is_ascii_alphabetic() {
113                    break;
114                }
115            }
116        } else {
117            result.push(ch);
118        }
119    }
120
121    result
122}
123
124/// Parse `opencode models` output into slug list.
125fn parse_models_output(stdout: &str) -> Vec<String> {
126    stdout
127        .lines()
128        .map(|line| strip_ansi(line.trim()))
129        .filter(|line| !line.is_empty() && line.contains('/'))
130        .collect()
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn test_parse_models_basic() {
139        let output = r#"opencode/big-pickle
140google/gemini-2.5-pro
141openai/gpt-5.4
142openrouter/anthropic/claude-opus-4.7"#;
143
144        let slugs = parse_models_output(output);
145
146        assert_eq!(slugs.len(), 4);
147        assert!(slugs.contains(&"openai/gpt-5.4".to_string()));
148        assert!(slugs.contains(&"openrouter/anthropic/claude-opus-4.7".to_string()));
149    }
150
151    #[test]
152    fn test_parse_models_filters_invalid() {
153        let slugs = parse_models_output("some-invalid-line\nopenai/gpt-5.4\n\n");
154        assert_eq!(slugs, vec!["openai/gpt-5.4"]);
155    }
156
157    #[test]
158    fn test_parse_models_strips_ansi() {
159        let slugs = parse_models_output("\x1b[32mopenai/gpt-5.4\x1b[0m");
160        assert_eq!(slugs, vec!["openai/gpt-5.4"]);
161    }
162
163    #[test]
164    fn test_probe_result_round_trip() {
165        let result = OpenCodeProbeResult {
166            model_slugs: vec!["openai/gpt-5.4".to_string()],
167            model_probe_success: true,
168            error: None,
169        };
170        let json = serde_json::to_string(&result).unwrap();
171        let back: OpenCodeProbeResult = serde_json::from_str(&json).unwrap();
172        assert_eq!(back.model_slugs, result.model_slugs);
173        assert!(back.model_probe_success);
174        assert_eq!(back.error, None);
175    }
176}