mars_agents/models/probes/
opencode.rs1use 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#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12pub struct OpenCodeProbeResult {
13 pub model_slugs: Vec<String>,
15 pub model_probe_success: bool,
17 pub error: Option<String>,
19}
20
21const DEFAULT_PROBE_TIMEOUT_SECS: u64 = 5;
22
23pub fn probe() -> OpenCodeProbeResult {
25 probe_with_timeout(probe_timeout())
26}
27
28pub 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
124fn 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}