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::ProbeRefreshMode;
10use crate::models::probes::cursor_cache::{self, CachedCursorProbeOutcome};
11use crate::models::probes::opencode_cache::{self, CachedProbeOutcome};
12use crate::models::probes::pi_cache::{self, CachedPiProbeOutcome};
13
14#[derive(Debug, Clone)]
15pub struct CapabilityCollectionOptions {
16    /// `MARS_OFFLINE` — skip network/catalog assumptions; probes treat env as offline.
17    pub offline: bool,
18    pub probe_refresh: ProbeRefreshMode,
19}
20
21impl Default for CapabilityCollectionOptions {
22    fn default() -> Self {
23        Self {
24            offline: false,
25            probe_refresh: ProbeRefreshMode::Background,
26        }
27    }
28}
29
30#[derive(Debug, Clone)]
31pub struct CapabilitySnapshot {
32    pub executable: BTreeMap<HarnessId, ExecutableState>,
33    pub auth: BTreeMap<HarnessId, AuthState>,
34    pub opencode: CachedProbeOutcome,
35    pub pi: CachedPiProbeOutcome,
36    pub cursor: CachedCursorProbeOutcome,
37    pub offline: bool,
38}
39
40impl CapabilitySnapshot {
41    pub fn installed_harnesses(&self) -> HashSet<String> {
42        self.executable
43            .iter()
44            .filter(|(_, state)| matches!(state, ExecutableState::Found { .. }))
45            .map(|(id, _)| id)
46            .map(|id| id.as_str().to_string())
47            .collect()
48    }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
52pub enum ExecutableState {
53    Found { path: PathBuf },
54    Missing,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub enum AuthState {
59    NotApplicable,
60    Authenticated,
61    Unauthenticated,
62    Unknown { reason: String },
63}
64
65pub trait ExecutableResolver {
66    fn resolve(&self, binary: &str) -> ExecutableState;
67}
68
69#[derive(Debug, Default, Clone, Copy)]
70pub struct PathExecutableResolver;
71
72impl ExecutableResolver for PathExecutableResolver {
73    fn resolve(&self, binary: &str) -> ExecutableState {
74        if let Ok(path) = which::which(binary) {
75            return ExecutableState::Found { path };
76        }
77
78        #[cfg(windows)]
79        {
80            for ext in ["exe", "cmd", "bat"] {
81                if let Ok(path) = which::which(format!("{binary}.{ext}")) {
82                    return ExecutableState::Found { path };
83                }
84            }
85        }
86
87        ExecutableState::Missing
88    }
89}
90
91pub fn collect_capability_snapshot(options: &CapabilityCollectionOptions) -> CapabilitySnapshot {
92    collect_capability_snapshot_with_resolver(options, &PathExecutableResolver)
93}
94
95pub fn collect_capability_snapshot_with_resolver(
96    options: &CapabilityCollectionOptions,
97    resolver: &dyn ExecutableResolver,
98) -> CapabilitySnapshot {
99    let mut executable = BTreeMap::new();
100    let mut auth = BTreeMap::new();
101
102    for descriptor in registry::descriptors() {
103        let state = resolver.resolve(descriptor.binary);
104        executable.insert(descriptor.id, state.clone());
105        auth.insert(
106            descriptor.id,
107            native_auth_state(descriptor.id, &state, resolver, auth_probe_timeout()),
108        );
109    }
110
111    let installed = executable
112        .iter()
113        .filter(|(_, state)| matches!(state, ExecutableState::Found { .. }))
114        .map(|(id, _)| id)
115        .map(|id| id.as_str().to_string())
116        .collect::<HashSet<_>>();
117
118    let mars_offline = options.offline;
119
120    CapabilitySnapshot {
121        executable,
122        auth,
123        opencode: opencode_cache::probe_cached(&installed, mars_offline, options.probe_refresh),
124        pi: pi_cache::probe_cached(&installed, mars_offline, options.probe_refresh),
125        cursor: cursor_cache::probe_cached(&installed, mars_offline, options.probe_refresh),
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            probe_refresh: ProbeRefreshMode::Skip,
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
