mars_agents/harness/
host.rs

1use std::collections::{BTreeMap, HashSet};
2use std::path::PathBuf;
3use std::process::{Command, Stdio};
4use std::time::Duration;
5
6use wait_timeout::ChildExt;
7
8use crate::harness::registry::{self, HarnessId};
9use crate::models::probes::cursor_cache::{self, CachedCursorProbeOutcome};
10use crate::models::probes::opencode_cache::{self, CachedProbeOutcome};
11use crate::models::probes::pi_cache::{self, CachedPiProbeOutcome};
12
13#[derive(Debug, Clone)]
14pub struct CapabilityCollectionOptions {
15    pub offline: bool,
16    pub allow_probe_refresh: bool,
17}
18
19impl Default for CapabilityCollectionOptions {
20    fn default() -> Self {
21        Self {
22            offline: false,
23            allow_probe_refresh: true,
24        }
25    }
26}
27
28#[derive(Debug, Clone)]
29pub struct CapabilitySnapshot {
30    pub executable: BTreeMap<HarnessId, ExecutableState>,
31    pub auth: BTreeMap<HarnessId, AuthState>,
32    pub opencode: CachedProbeOutcome,
33    pub pi: CachedPiProbeOutcome,
34    pub cursor: CachedCursorProbeOutcome,
35    pub offline: bool,
36}
37
38impl CapabilitySnapshot {
39    pub fn installed_harnesses(&self) -> HashSet<String> {
40        self.executable
41            .iter()
42            .filter(|(_, state)| matches!(state, ExecutableState::Found { .. }))
43            .map(|(id, _)| id)
44            .map(|id| id.as_str().to_string())
45            .collect()
46    }
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum ExecutableState {
51    Found { path: PathBuf },
52    Missing,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub enum AuthState {
57    NotApplicable,
58    Authenticated,
59    Unauthenticated,
60    Unknown { reason: String },
61}
62
63pub trait ExecutableResolver {
64    fn resolve(&self, binary: &str) -> ExecutableState;
65}
66
67#[derive(Debug, Default, Clone, Copy)]
68pub struct PathExecutableResolver;
69
70impl ExecutableResolver for PathExecutableResolver {
71    fn resolve(&self, binary: &str) -> ExecutableState {
72        if let Ok(path) = which::which(binary) {
73            return ExecutableState::Found { path };
74        }
75
76        #[cfg(windows)]
77        {
78            for ext in ["exe", "cmd", "bat"] {
79                if let Ok(path) = which::which(format!("{binary}.{ext}")) {
80                    return ExecutableState::Found { path };
81                }
82            }
83        }
84
85        ExecutableState::Missing
86    }
87}
88
89pub fn collect_capability_snapshot(options: &CapabilityCollectionOptions) -> CapabilitySnapshot {
90    collect_capability_snapshot_with_resolver(options, &PathExecutableResolver)
91}
92
93pub fn collect_capability_snapshot_with_resolver(
94    options: &CapabilityCollectionOptions,
95    resolver: &dyn ExecutableResolver,
96) -> CapabilitySnapshot {
97    let mut executable = BTreeMap::new();
98    let mut auth = BTreeMap::new();
99
100    for descriptor in registry::descriptors() {
101        let state = resolver.resolve(descriptor.binary);
102        executable.insert(descriptor.id, state.clone());
103        auth.insert(
104            descriptor.id,
105            native_auth_state(descriptor.id, &state, resolver, auth_probe_timeout()),
106        );
107    }
108
109    let installed = executable
110        .iter()
111        .filter(|(_, state)| matches!(state, ExecutableState::Found { .. }))
112        .map(|(id, _)| id)
113        .map(|id| id.as_str().to_string())
114        .collect::<HashSet<_>>();
115
116    let opencode_offline = options.offline || !options.allow_probe_refresh;
117    let pi_offline = options.offline || !options.allow_probe_refresh;
118    let cursor_offline = options.offline || !options.allow_probe_refresh;
119
120    CapabilitySnapshot {
121        executable,
122        auth,
123        opencode: opencode_cache::probe_cached(&installed, opencode_offline),
124        pi: pi_cache::probe_cached(&installed, pi_offline),
125        cursor: cursor_cache::probe_cached(&installed, cursor_offline),
126        offline: options.offline,
127    }
128}
129
130pub fn native_harness_authenticated(harness: &str) -> bool {
131    native_auth_state_for_name(harness) == AuthState::Authenticated
132}
133
134pub fn native_auth_state_for_name(harness: &str) -> AuthState {
135    let Some(id) = registry::parse(harness) else {
136        return AuthState::Unknown {
137            reason: "unknown harness".to_string(),
138        };
139    };
140
141    let resolver = PathExecutableResolver;
142    let state = resolver.resolve(registry::descriptor(id).binary);
143    native_auth_state(id, &state, &resolver, auth_probe_timeout())
144}
145
146fn native_auth_state(
147    id: HarnessId,
148    executable: &ExecutableState,
149    resolver: &dyn ExecutableResolver,
150    timeout: Duration,
151) -> AuthState {
152    let (binary, args) = match id {
153        HarnessId::Codex => ("codex", &["login", "status"][..]),
154        HarnessId::Claude => ("claude", &["auth", "status"][..]),
155        _ => return AuthState::NotApplicable,
156    };
157
158    if !matches!(executable, ExecutableState::Found { .. }) {
159        return AuthState::Unauthenticated;
160    }
161
162    run_status_command(binary, args, timeout, resolver)
163}
164
165pub fn auth_probe_timeout() -> Duration {
166    std::env::var("MARS_NATIVE_HARNESS_AUTH_TIMEOUT_SECS")
167        .ok()
168        .and_then(|value| value.parse::<u64>().ok())
169        .map(Duration::from_secs)
170        .unwrap_or(Duration::from_secs(2))
171}
172
173fn run_status_command(
174    command: &str,
175    args: &[&str],
176    timeout: Duration,
177    resolver: &dyn ExecutableResolver,
178) -> AuthState {
179    let program = resolve_binary_path(command, resolver).unwrap_or_else(|| PathBuf::from(command));
180
181    let mut child = match Command::new(program)
182        .args(args)
183        .stdin(Stdio::null())
184        .stdout(Stdio::null())
185        .stderr(Stdio::null())
186        .spawn()
187    {
188        Ok(child) => child,
189        Err(error) => {
190            return AuthState::Unknown {
191                reason: format!("spawn failed: {error}"),
192            };
193        }
194    };
195
196    match child.wait_timeout(timeout) {
197        Ok(Some(status)) if status.success() => AuthState::Authenticated,
198        Ok(Some(_)) => AuthState::Unauthenticated,
199        Ok(None) => {
200            let _ = child.kill();
201            let _ = child.wait();
202            AuthState::Unknown {
203                reason: "auth probe timeout".to_string(),
204            }
205        }
206        Err(error) => AuthState::Unknown {
207            reason: format!("auth probe wait failed: {error}"),
208        },
209    }
210}
211
212pub fn resolve_binary_path(binary: &str, resolver: &dyn ExecutableResolver) -> Option<PathBuf> {
213    match resolver.resolve(binary) {
214        ExecutableState::Found { path } => Some(path),
215        ExecutableState::Missing => None,
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use std::collections::HashMap;
223
224    #[derive(Default)]
225    struct FakeResolver {
226        map: HashMap<String, ExecutableState>,
227    }
228
229    impl ExecutableResolver for FakeResolver {
230        fn resolve(&self, binary: &str) -> ExecutableState {
231            self.map
232                .get(binary)
233                .cloned()
234                .unwrap_or(ExecutableState::Missing)
235        }
236    }
237
238    #[test]
239    fn snapshot_marks_installed_harnesses_from_resolver() {
240        let mut resolver = FakeResolver::default();
241        resolver.map.insert(
242            "pi".to_string(),
243            ExecutableState::Found {
244                path: PathBuf::from("/tmp/pi"),
245            },
246        );
247
248        let options = CapabilityCollectionOptions {
249            offline: true,
250            allow_probe_refresh: false,
251        };
252        let snapshot = collect_capability_snapshot_with_resolver(&options, &resolver);
253
254        let installed = snapshot.installed_harnesses();
255        assert!(installed.contains("pi"));
256        assert!(!installed.contains("codex"));
257    }
258
259    #[test]
260    fn native_auth_for_non_native_harness_is_not_applicable() {
261        let resolver = FakeResolver::default();
262        let state = native_auth_state(
263            HarnessId::Pi,
264            &ExecutableState::Found {
265                path: PathBuf::from("/tmp/pi"),
266            },
267            &resolver,
268            Duration::from_secs(1),
269        );
270
271        assert_eq!(state, AuthState::NotApplicable);
272    }
273
274    #[test]
275    fn resolve_binary_path_returns_none_when_missing() {
276        let resolver = FakeResolver::default();
277        assert_eq!(resolve_binary_path("codex", &resolver), None);
278    }
279}
mars_agents/harness/host.rs

mars_agents/harness/
host.rs
