Skip to main content

twine_core/
harness.rs

1//! Built-in harness definitions and finding their binaries.
2
3use std::collections::HashSet;
4use std::ffi::{OsStr, OsString};
5use std::io::{Read, Seek};
6use std::os::unix::ffi::OsStrExt;
7use std::path::{Path, PathBuf};
8use std::process::{Command, Stdio};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::{Arc, Mutex};
11use std::thread::{self, JoinHandle};
12use std::time::{Duration, Instant};
13
14use thiserror::Error;
15use tracing::debug;
16
17use crate::process::CommandChild;
18
19pub(crate) mod antigravity;
20pub(crate) mod claude;
21pub(crate) mod codex;
22pub(crate) mod history;
23pub(crate) mod instructions;
24pub(crate) mod launch;
25pub(crate) mod models;
26pub(crate) mod opencode;
27pub(crate) mod pi;
28pub(crate) mod resume;
29pub(crate) mod steps;
30
31/// How long to wait for the login shell to report its `PATH`.
32const LOGIN_SHELL_TIMEOUT: Duration = Duration::from_secs(5);
33const MAX_LOGIN_SHELL_OUTPUT: u64 = 64 * 1024;
34const PATH_START: &[u8] = b"__TWINE_PATH__";
35const PATH_END: &[u8] = b"__TWINE_END__";
36
37#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, serde::Serialize, serde::Deserialize)]
38#[serde(rename_all = "camelCase")]
39pub enum HarnessId {
40    Codex,
41    ClaudeCode,
42    Pi,
43    Antigravity,
44    Omp,
45    Opencode,
46}
47
48/// How to launch one harness. Adding a harness means adding a definition here.
49#[derive(Debug)]
50pub(crate) struct HarnessDefinition {
51    pub name: &'static str,
52    pub binary: &'static str,
53    /// Whether the harness reads a prompt that starts with `@` as a file path.
54    pub reads_at_prefix_as_file: bool,
55    /// A flag for an initial prompt that keeps the harness interactive, when it needs one.
56    pub prompt_option: Option<&'static str>,
57    /// An explicit interactive subcommand, when positional prompts could select another command.
58    pub subcommand: Option<&'static str>,
59}
60
61static CODEX: HarnessDefinition = HarnessDefinition {
62    name: "Codex",
63    binary: "codex",
64    reads_at_prefix_as_file: false,
65    prompt_option: None,
66    subcommand: None,
67};
68static CLAUDE_CODE: HarnessDefinition = HarnessDefinition {
69    name: "Claude Code",
70    binary: "claude",
71    reads_at_prefix_as_file: false,
72    prompt_option: None,
73    subcommand: None,
74};
75static PI: HarnessDefinition = HarnessDefinition {
76    name: "pi",
77    binary: "pi",
78    reads_at_prefix_as_file: true,
79    prompt_option: None,
80    subcommand: None,
81};
82static ANTIGRAVITY: HarnessDefinition = HarnessDefinition {
83    name: "Antigravity",
84    binary: "agy",
85    reads_at_prefix_as_file: false,
86    prompt_option: Some("--prompt-interactive"),
87    subcommand: None,
88};
89static OMP: HarnessDefinition = HarnessDefinition {
90    name: "OMP",
91    binary: "omp",
92    // OMP treats everything after `--` as literal text, including `@` prefixes.
93    reads_at_prefix_as_file: false,
94    prompt_option: None,
95    subcommand: Some("launch"),
96};
97static OPENCODE: HarnessDefinition = HarnessDefinition {
98    name: "OpenCode",
99    binary: "opencode",
100    reads_at_prefix_as_file: false,
101    prompt_option: Some("--prompt"),
102    subcommand: Some("mini"),
103};
104
105impl HarnessId {
106    pub(crate) fn definition(self) -> &'static HarnessDefinition {
107        match self {
108            Self::Codex => &CODEX,
109            Self::ClaudeCode => &CLAUDE_CODE,
110            Self::Pi => &PI,
111            Self::Antigravity => &ANTIGRAVITY,
112            Self::Omp => &OMP,
113            Self::Opencode => &OPENCODE,
114        }
115    }
116}
117
118impl HarnessDefinition {
119    /// An interactive prompt, protected from being read as a flag even with a leading dash.
120    pub(crate) fn arguments(&self, prompt: &str) -> Vec<OsString> {
121        let prompt = if self.reads_at_prefix_as_file && prompt.starts_with('@') {
122            format!(" {prompt}")
123        } else {
124            prompt.to_owned()
125        };
126        match self.prompt_option {
127            Some(option) => vec![OsString::from(format!("{option}={prompt}"))],
128            None => vec![OsString::from("--"), OsString::from(prompt)],
129        }
130    }
131}
132
133/// A harness binary and the `PATH` it should run with.
134#[derive(Debug, Eq, PartialEq)]
135pub(crate) struct LocatedHarness {
136    pub program: PathBuf,
137    pub path: OsString,
138}
139
140#[derive(Debug, Error)]
141pub enum HarnessError {
142    #[error(
143        "{name} wasn't found on your PATH. Install it or add `{binary}` to your PATH, then try again."
144    )]
145    NotFound {
146        name: &'static str,
147        binary: &'static str,
148    },
149}
150
151/// Finds `definition`'s binary on the login shell's `PATH`, then on `process_path`.
152///
153/// A GUI app inherits a minimal `PATH`, so the login shell's is searched first.
154pub(crate) fn locate_in(
155    definition: &HarnessDefinition,
156    login_path: Option<&OsStr>,
157    process_path: Option<&OsStr>,
158) -> Result<LocatedHarness, HarnessError> {
159    let directories = unique_directories(login_path.into_iter().chain(process_path));
160    let program = directories
161        .iter()
162        .map(|directory| directory.join(definition.binary))
163        .find(|candidate| is_executable(candidate))
164        .ok_or(HarnessError::NotFound {
165            name: definition.name,
166            binary: definition.binary,
167        })?;
168    let path = std::env::join_paths(&directories).unwrap_or_default();
169    Ok(LocatedHarness { program, path })
170}
171
172fn unique_directories<'a>(paths: impl Iterator<Item = &'a OsStr>) -> Vec<PathBuf> {
173    let mut seen = HashSet::new();
174    paths
175        .flat_map(|path| std::env::split_paths(path))
176        .filter(|directory| !directory.as_os_str().is_empty() && seen.insert(directory.clone()))
177        .collect()
178}
179
180fn is_executable(path: &Path) -> bool {
181    use std::os::unix::fs::PermissionsExt;
182    std::fs::metadata(path)
183        .is_ok_and(|metadata| metadata.is_file() && metadata.permissions().mode() & 0o111 != 0)
184}
185
186/// The login shell's `PATH`, looked up once in the background so starting an agent never waits on
187/// the user's shell startup files.
188pub(crate) struct LoginPath {
189    state: Mutex<LoginPathState>,
190    cancelled: Arc<AtomicBool>,
191}
192
193enum LoginPathState {
194    Pending(JoinHandle<Option<OsString>>),
195    Ready(Option<OsString>),
196}
197
198impl LoginPath {
199    pub(crate) fn spawn(shell: PathBuf) -> Self {
200        let cancelled = Arc::new(AtomicBool::new(false));
201        let flag = Arc::clone(&cancelled);
202        match crate::blocking_worker::spawn("login-path".into(), move || {
203            login_shell_path(&shell, &flag)
204        }) {
205            Ok(handle) => Self {
206                state: Mutex::new(LoginPathState::Pending(handle)),
207                cancelled,
208            },
209            Err(error) => {
210                debug!(%error, "failed to start the login PATH lookup");
211                Self::ready(None)
212            }
213        }
214    }
215
216    pub(crate) fn ready(path: Option<OsString>) -> Self {
217        Self {
218            state: Mutex::new(LoginPathState::Ready(path)),
219            cancelled: Arc::new(AtomicBool::new(false)),
220        }
221    }
222
223    pub(crate) fn get(&self) -> Option<OsString> {
224        let mut state = self
225            .state
226            .lock()
227            .unwrap_or_else(std::sync::PoisonError::into_inner);
228        if let LoginPathState::Pending(_) = &*state {
229            let LoginPathState::Pending(handle) =
230                std::mem::replace(&mut *state, LoginPathState::Ready(None))
231            else {
232                unreachable!("state was just checked to be pending");
233            };
234            *state = LoginPathState::Ready(handle.join().unwrap_or(None));
235        }
236        if self.cancelled.load(Ordering::Relaxed) {
237            return None;
238        }
239        match &*state {
240            LoginPathState::Ready(path) => path.clone(),
241            LoginPathState::Pending(_) => None,
242        }
243    }
244
245    pub(crate) fn cancel(&self) {
246        self.cancelled.store(true, Ordering::Relaxed);
247    }
248
249    pub(crate) fn shutdown(&self) {
250        self.cancel();
251        // `get` joins even a cancelled lookup, so every caller waits for process cleanup.
252        let _ = self.get();
253    }
254}
255
256impl Drop for LoginPath {
257    fn drop(&mut self) {
258        self.shutdown();
259    }
260}
261
262/// Asks an interactive login shell for its `PATH`, where version managers such as nvm and fnm are
263/// usually set up. Only `PATH` is read, between markers so startup banners can't be mistaken for it.
264fn login_shell_path(shell: &Path, cancelled: &AtomicBool) -> Option<OsString> {
265    if cancelled.load(Ordering::Relaxed) {
266        return None;
267    }
268    // A file avoids blocking on a full pipe while the shell is being supervised.
269    let mut output = tempfile::tempfile().ok()?;
270    let mut command = Command::new(shell);
271    command
272        .args(["-i", "-l", "-c"])
273        .arg("printf '\\n__TWINE_PATH__%s__TWINE_END__\\n' \"$(printenv PATH)\"")
274        .stdin(Stdio::null())
275        .stdout(output.try_clone().ok()?)
276        .stderr(Stdio::null());
277    let mut child = CommandChild::spawn(&mut command)
278        .map_err(|error| debug!(%error, "failed to start the login shell for PATH lookup"))
279        .ok()?;
280    let deadline = Instant::now() + LOGIN_SHELL_TIMEOUT;
281    loop {
282        if cancelled.load(Ordering::Relaxed) {
283            return None;
284        }
285        if output.metadata().ok()?.len() > MAX_LOGIN_SHELL_OUTPUT {
286            return None;
287        }
288        if child.has_exited().ok()? {
289            // Stop background jobs before reaping the shell and releasing its process identity.
290            child.finish().ok()?;
291            break;
292        }
293        if Instant::now() >= deadline {
294            debug!("login shell PATH lookup timed out");
295            return None;
296        }
297        thread::sleep(Duration::from_millis(5));
298    }
299    output.rewind().ok()?;
300    let mut bytes = Vec::new();
301    output
302        .take(MAX_LOGIN_SHELL_OUTPUT)
303        .read_to_end(&mut bytes)
304        .ok()?;
305    parse_path(&bytes)
306}
307
308fn parse_path(output: &[u8]) -> Option<OsString> {
309    let start = find(output, PATH_START)? + PATH_START.len();
310    let end = start + find(&output[start..], PATH_END)?;
311    let path = &output[start..end];
312    (!path.is_empty()).then(|| OsStr::from_bytes(path).to_owned())
313}
314
315fn find(haystack: &[u8], needle: &[u8]) -> Option<usize> {
316    haystack
317        .windows(needle.len())
318        .position(|window| window == needle)
319}
320
321#[cfg(test)]
322mod tests {
323    use std::os::unix::fs::PermissionsExt;
324
325    use super::*;
326    use crate::process::tests::FixtureProcess;
327
328    #[test]
329    fn prompt_follows_a_double_dash() {
330        let arguments = HarnessId::ClaudeCode.definition().arguments("-fix it");
331        assert_eq!(arguments, ["--", "-fix it"].map(OsString::from));
332    }
333
334    #[test]
335    fn pi_keeps_an_at_prompt_from_being_read_as_a_file() {
336        let pi = HarnessId::Pi.definition().arguments("@README.md summarize");
337        assert_eq!(pi, ["--", " @README.md summarize"].map(OsString::from));
338        let claude = HarnessId::ClaudeCode.definition().arguments("@README.md");
339        assert_eq!(claude, ["--", "@README.md"].map(OsString::from));
340    }
341
342    #[test]
343    fn antigravity_keeps_initial_prompts_interactive_and_literal() {
344        let definition = HarnessId::Antigravity.definition();
345        assert_eq!(definition.binary, "agy");
346        for prompt in ["-fix it", "@README.md summarize", "first line\nsecond line"] {
347            assert_eq!(
348                definition.arguments(prompt),
349                [OsString::from(format!("--prompt-interactive={prompt}"))]
350            );
351        }
352    }
353
354    #[test]
355    fn omp_treats_flag_file_and_command_shaped_prompts_as_literal_text() {
356        let definition = HarnessId::Omp.definition();
357        assert_eq!(definition.binary, "omp");
358        assert_eq!(definition.subcommand, Some("launch"));
359        for prompt in ["-fix it", "@README.md", "models", "line one\nline two"] {
360            assert_eq!(
361                definition.arguments(prompt),
362                ["--", prompt].map(OsString::from)
363            );
364        }
365    }
366
367    #[test]
368    fn opencode_keeps_prompts_out_of_its_directory_and_command_arguments() {
369        let definition = HarnessId::Opencode.definition();
370        assert_eq!(definition.binary, "opencode");
371        assert_eq!(definition.subcommand, Some("mini"));
372        for prompt in ["-fix it", "@README.md", "models", "line one\nline two"] {
373            assert_eq!(
374                definition.arguments(prompt),
375                [OsString::from(format!("--prompt={prompt}"))]
376            );
377        }
378    }
379
380    #[test]
381    fn parses_path_between_markers_and_ignores_banners() {
382        assert_eq!(
383            parse_path(b"PATH=/wrong\nbanner\n__TWINE_PATH__/a:/b__TWINE_END__\n"),
384            Some(OsString::from("/a:/b"))
385        );
386        assert_eq!(parse_path(b"PATH=/wrong\n"), None);
387        assert_eq!(parse_path(b"__TWINE_PATH____TWINE_END__"), None);
388    }
389
390    fn fake_shell(directory: &Path, body: &str) -> PathBuf {
391        let shell = directory.join("shell");
392        std::fs::write(&shell, format!("#!/bin/sh\n{body}\n")).unwrap();
393        std::fs::set_permissions(&shell, std::fs::Permissions::from_mode(0o755)).unwrap();
394        shell
395    }
396
397    #[test]
398    fn login_path_comes_from_an_interactive_login_shell() {
399        let directory = tempfile::tempdir().unwrap();
400        // Like a zsh that only sets up its version manager in the interactive rc file.
401        let shell = fake_shell(
402            directory.path(),
403            "echo banner PATH=/wrong\n[ \"$1 $2\" = '-i -l' ] && printf '\\n__TWINE_PATH__/interactive__TWINE_END__\\n'",
404        );
405
406        assert_eq!(
407            LoginPath::spawn(shell).get(),
408            Some(OsString::from("/interactive"))
409        );
410    }
411
412    #[test]
413    fn a_background_job_holding_the_output_does_not_delay_the_lookup() {
414        let directory = tempfile::tempdir().unwrap();
415        let pid_path = directory.path().join("child.pid");
416        let shell = fake_shell(
417            directory.path(),
418            &format!(
419                "sleep 60 & echo $! > '{}'; printf '\\n__TWINE_PATH__/fast__TWINE_END__\\n'",
420                pid_path.display()
421            ),
422        );
423
424        let started = Instant::now();
425        let path = login_shell_path(&shell, &AtomicBool::new(false));
426
427        assert_eq!(path, Some(OsString::from("/fast")));
428        assert!(started.elapsed() < Duration::from_secs(4));
429        FixtureProcess::read(&pid_path).assert_stopped();
430    }
431
432    #[test]
433    fn dropping_an_unused_login_path_stops_and_joins_its_shell() {
434        let directory = tempfile::tempdir().unwrap();
435        let pid_path = directory.path().join("shell.pid");
436        let shell = fake_shell(
437            directory.path(),
438            &format!("echo $$ > '{}'; exec sleep 60", pid_path.display()),
439        );
440        let path = LoginPath::spawn(shell);
441        let shell = FixtureProcess::read(&pid_path);
442        let started = Instant::now();
443        drop(path);
444        assert!(started.elapsed() < Duration::from_secs(2));
445        assert!(
446            !shell.is_running(),
447            "shutdown reaps the login shell before returning"
448        );
449    }
450
451    #[test]
452    fn locates_an_executable_and_reports_a_missing_one() {
453        let directory = tempfile::tempdir().expect("a directory should be available");
454        let binary = directory.path().join("twine-test-harness");
455        std::fs::write(&binary, "#!/bin/sh\n").expect("binary should be written");
456        std::fs::set_permissions(&binary, std::fs::Permissions::from_mode(0o755))
457            .expect("permissions should be set");
458        let found = HarnessDefinition {
459            name: "Test",
460            binary: "twine-test-harness",
461            reads_at_prefix_as_file: false,
462            prompt_option: None,
463            subcommand: None,
464        };
465        let missing = HarnessDefinition {
466            binary: "twine-no-such-harness",
467            ..found
468        };
469        let path = Some(directory.path().as_os_str());
470
471        let located = locate_in(&found, None, path).expect("binary is on PATH");
472        let error = locate_in(&missing, None, path).expect_err("a missing binary is not found");
473
474        assert_eq!(located.program, binary);
475        assert!(error.to_string().contains("`twine-no-such-harness`"));
476    }
477}