From 464d216f6758ccf42e9583f3d7d5c018af9a698d Mon Sep 17 00:00:00 2001 From: Paolo Santos Date: Fri, 6 Nov 2026 23:32:01 +0800 Subject: [PATCH 1/2] exec: one-shot codex exec runs in a private CODEX_HOME; auth: may_refresh, refresh_lock exec::Exec::run copies the login into a fresh 0700 CODEX_HOME, runs `codex exec --json` with the prompt on stdin and the images attached, copies refreshed tokens back (serialised in-process), and removes the run dir when the Run is dropped. exec::failure reads why a run failed (a dead ChatGPT login, a 401, the turn's error); image_prompt and Run::generated_images serve image generation. auth::may_refresh says when a run may refresh the single-use refresh token, and refresh_lock lets such a run go alone. --- Cargo.lock | 32 +++ Cargo.toml | 7 +- README.md | 9 +- src/auth.rs | 70 ++++- src/exec.rs | 609 ++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + tests/exec.rs | 114 ++++++++ tests/fake_codex_exec | 31 +++ 8 files changed, 871 insertions(+), 3 deletions(-) create mode 100644 src/exec.rs create mode 100644 tests/exec.rs create mode 100755 tests/fake_codex_exec diff --git a/Cargo.lock b/Cargo.lock index 0efa7b8..1156b4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,18 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bitflags" version = "2.13.2" @@ -20,10 +32,21 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "num-traits", +] + [[package]] name = "codexwire" version = "0.1.0" dependencies = [ + "base64", + "chrono", "serde", "serde_json", "tempfile", @@ -92,6 +115,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" diff --git a/Cargo.toml b/Cargo.toml index caeb065..5cdb6da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,16 +18,21 @@ exclude = ["schema", "scripts"] [workspace] [dependencies] +# auth::may_refresh: the access token's claims and `last_refresh`. +base64 = "0.22" +chrono = { version = "0.4", default-features = false, features = ["alloc"] } serde = { version = "1", features = ["derive"] } serde_json = "1" # The client: the child process and its stdio (process, io-util), the -# reader task and response channels (rt, sync), auth.json (fs). +# reader task and response channels (rt, sync), auth.json (fs), exec's +# timeout (time). tokio = { version = "1", default-features = false, features = [ "fs", "io-util", "process", "rt", "sync", + "time", ] } [dev-dependencies] diff --git a/README.md b/README.md index a6819c5..68fcf3e 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,14 @@ Unix only (the child's process handling, `auth.json`'s file modes). `tokio::time::timeout` (a dropped call forgets its request). `redact` scrubs secrets out of its stderr. - `codexwire::auth`: a session's private copy of `auth.json`, copied back - only if Codex refreshed it and the real one didn't change meanwhile. + only if Codex refreshed it and the real one didn't change meanwhile; + `may_refresh` and `refresh_lock` keep two runs from spending the same + single-use refresh token. +- `codexwire::exec`: one-shot `codex exec --json` runs, each in a fresh + private `CODEX_HOME` holding only a copy of the login (copied back when + refreshed, the rest removed); `failure` says why a run failed (a dead + login, a refused one, the turn's error); `image_prompt` and + `Run::generated_images` for image generation. From `examples/initialize.rs`: diff --git a/src/auth.rs b/src/auth.rs index 1cbaf11..66a5f17 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -18,6 +18,49 @@ pub fn auth_to_copy_back(orig: &[u8], now_real: &[u8], after: &[u8]) -> bool { && serde_json::from_slice::(after).is_ok_and(|v| v.is_object()) } +/// Whether a run on login `auth` (an `auth.json`) may refresh its tokens +/// at `now` (unix seconds; pure): the access token ends within 10 +/// minutes, the last refresh is a week old (Codex refreshes ahead of +/// time), or either can't be read. +/// +/// Such a run should go alone ([`refresh_lock`]): `ChatGPT`'s refresh +/// tokens are single-use, and two runs refreshing the same one get it +/// revoked (`refresh_token_reused`) until the user signs in again. +pub fn may_refresh(auth: &[u8], now: i64) -> bool { + use base64::Engine as _; + let Ok(v) = serde_json::from_slice::(auth) else { + return true; + }; + let exp = v["tokens"]["access_token"] + .as_str() + .and_then(|t| t.split('.').nth(1)) + .and_then(|p| { + base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(p.trim_end_matches('=')) + .ok() + }) + .and_then(|b| serde_json::from_slice::(&b).ok()) + .and_then(|c| c["exp"].as_i64()); + let last = v["last_refresh"] + .as_str() + .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok()) + .map(|d| d.timestamp()); + match (exp, last) { + (Some(exp), Some(last)) => exp <= now + 600 || last <= now - 7 * 86_400, + _ => true, + } +} + +/// This process's lock for runs that [`may_refresh`] the login: hold it +/// (it is async, so across the run) from reading `auth.json` until the +/// refreshed one is copied back, and the next run reads the new tokens. +/// One lock for every login: a process rarely has more than one, and +/// runs that may refresh are rare. It doesn't reach other processes. +pub fn refresh_lock() -> &'static tokio::sync::Mutex<()> { + static L: std::sync::OnceLock> = std::sync::OnceLock::new(); + L.get_or_init(tokio::sync::Mutex::default) +} + /// What was copied into a session's home, to copy back on close. /// /// It holds the copied `auth.json` (its tokens), so its `Debug` shows only @@ -30,6 +73,10 @@ pub struct Copied { } impl Copied { + pub(crate) fn new(real: PathBuf, ours: PathBuf, orig: Vec) -> Self { + Self { real, ours, orig } + } + /// The length of the `auth.json` that was copied in. pub fn orig_len(&self) -> usize { self.orig.len() @@ -89,7 +136,7 @@ pub async fn mkdir_private(p: &Path) -> std::io::Result<()> { } /// Write `data` to `p` (0600) through a temp file and a rename. -async fn write_private_atomic(p: &Path, data: &[u8]) -> std::io::Result<()> { +pub(crate) async fn write_private_atomic(p: &Path, data: &[u8]) -> std::io::Result<()> { let name = p .file_name() .map(|n| n.to_string_lossy().into_owned()) @@ -138,6 +185,27 @@ mod tests { assert!(!auth_to_copy_back(orig, orig, b"[1]"), "not an object"); } + #[test] + fn may_refresh_only_near_expiry() { + use base64::Engine as _; + let now = 1_800_000_000; + let auth = |exp: i64, last: &str| { + let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(format!(r#"{{"exp":{exp}}}"#)); + format!(r#"{{"last_refresh":"{last}","tokens":{{"access_token":"h.{claims}.s"}}}}"#) + }; + let recent = "2027-01-14T00:00:00Z"; // a day before `now` + assert!(!may_refresh(auth(now + 86_400, recent).as_bytes(), now)); + assert!(may_refresh(auth(now + 60, recent).as_bytes(), now)); + assert!(may_refresh(auth(now - 60, recent).as_bytes(), now)); + assert!(may_refresh( + auth(now + 86_400, "2027-01-01T00:00:00Z").as_bytes(), + now + )); + assert!(may_refresh(b"{}", now)); + assert!(may_refresh(b"nope", now)); + } + #[test] fn debug_redacts_the_tokens() { let c = Copied { diff --git a/src/exec.rs b/src/exec.rs new file mode 100644 index 0000000..7a25e08 --- /dev/null +++ b/src/exec.rs @@ -0,0 +1,609 @@ +//! One-shot `codex exec` runs on a `ChatGPT` login, each in a fresh private +//! `CODEX_HOME`. +//! +//! [`Exec::run`] makes `//` (0700) with `home/` (the run's +//! `CODEX_HOME`, holding only a copy of the real `auth.json`), `work/` (its +//! working directory) and `refs/` (the input images), runs `codex exec` +//! there with the prompt on stdin and `--json` events on stdout, and copies +//! the login back when Codex refreshed it ([`crate::auth::copy_back`]'s +//! rule). A run that [`may_refresh`] the login goes alone +//! ([`crate::auth::refresh_lock`]). The [`Run`] it returns holds the run's +//! directory, removed when it is dropped (or [`Run::remove`]). +//! +//! [`failure`] reads why a run failed from its events and log; +//! [`image_prompt`] and [`Run::generated_images`] are for image generation +//! (Codex's image tool writes to `/generated_images/`). +//! +//! Unix only. + +use core::{ + fmt, + sync::atomic::{AtomicU64, Ordering}, + time::Duration, +}; +use std::{ + path::{Path, PathBuf}, + process::Stdio, +}; + +use crate::auth::{Copied, may_refresh, mkdir_private, refresh_lock, write_private_atomic}; + +/// How to run `codex exec`. +#[derive(Clone, Debug)] +pub struct Exec { + /// The `codex` binary. + pub command: String, + /// The real `CODEX_HOME`: only its `auth.json` is used (copied in, and + /// back when refreshed). + pub codex_home: PathBuf, + /// Where the per-run directories go (made 0700 if missing). + pub tmp_root: PathBuf, + /// `--model`, if any (else Codex's default). + pub model: Option, + /// `--sandbox` (`read-only`, `workspace-write`, `danger-full-access`). + pub sandbox: String, + /// Killed after this long. + pub timeout: Duration, + /// The child's environment on top of the parent's allowlisted variables + /// ([`crate::client::child_env`]), e.g. `PATH`. `HOME`, `CODEX_HOME` + /// and `TMPDIR` are always the run's; `LANG` is `C.UTF-8` unless given. + pub env: Vec<(String, String)>, + /// More `codex exec` arguments, before the prompt. + pub args: Vec, +} + +impl Exec { + /// `command` on the login in `codex_home`, runs under `tmp_root`: a + /// read-only sandbox, a 5 minute timeout. + pub fn new( + command: impl Into, + codex_home: impl Into, + tmp_root: impl Into, + ) -> Self { + Self { + command: command.into(), + codex_home: codex_home.into(), + tmp_root: tmp_root.into(), + model: None, + sandbox: "read-only".into(), + timeout: Duration::from_secs(300), + env: Vec::new(), + args: Vec::new(), + } + } +} + +/// What a run is given. +#[derive(Clone, Debug, Default)] +pub struct Input { + /// The prompt (on stdin). + pub prompt: String, + /// Images to attach (`--image`): a file name (e.g. `ref-1.png`; only + /// its last component is used) and the bytes. + pub images: Vec<(String, Vec)>, +} + +/// Why a run didn't start or didn't finish setting up. +#[derive(Debug)] +pub enum ExecError { + /// The real `auth.json` can't be read (no login, or no access). + Login(std::io::Error), + /// Making the run's directory or files (what, and why). + Setup(&'static str, std::io::Error), + /// The `codex` binary didn't start. + Spawn(std::io::Error), +} + +impl fmt::Display for ExecError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Login(e) => write!(f, "Codex's auth.json isn't readable: {e}"), + Self::Setup(what, e) => write!(f, "codex run {what}: {e}"), + Self::Spawn(e) => write!(f, "codex didn't start: {e}"), + } + } +} + +impl core::error::Error for ExecError { + fn source(&self) -> Option<&(dyn core::error::Error + 'static)> { + match self { + Self::Login(e) | Self::Setup(_, e) | Self::Spawn(e) => Some(e), + } + } +} + +/// A finished (or timed out) run. Its directory goes when it is dropped. +#[derive(Debug)] +pub struct Run { + dir: RunDir, + /// The exit code; `None` when it timed out and was killed (`-1` when it + /// died of a signal). + pub code: Option, + /// Its `--json` events. + pub stdout: Vec, + /// Its log. + pub stderr: Vec, + /// `--output-last-message`: the final answer, if it wrote one. + pub last_message: Vec, + /// Whether Codex refreshed the login and it was copied back (or why + /// copying it back failed). + pub copied_back: std::io::Result, +} + +impl Run { + /// The run's `CODEX_HOME`. + pub fn home(&self) -> PathBuf { + self.dir.0.join("home") + } + + /// Why it failed, from its events and log ([`failure`]). + pub fn failure(&self) -> Option { + failure(&self.stdout, &self.stderr) + } + + /// The images Codex made (regular files under + /// `/generated_images/`, no symlinks, at most `max_bytes` + /// each), oldest first. Not checked to be images. + pub async fn generated_images(&self, max_bytes: u64) -> Vec { + files_oldest_first(&self.home().join("generated_images"), max_bytes).await + } + + /// Remove the run's directory now. + pub async fn remove(mut self) -> std::io::Result<()> { + let p = core::mem::take(&mut self.dir.0); + match tokio::fs::remove_dir_all(&p).await { + Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e), + _ => Ok(()), + } + } +} + +/// Removes its directory when dropped (unless emptied by [`Run::remove`]). +#[derive(Debug)] +struct RunDir(PathBuf); + +impl Drop for RunDir { + fn drop(&mut self) { + if !self.0.as_os_str().is_empty() { + // Drop can't await; a run's directory is small. + let _ = std::fs::remove_dir_all(&self.0); + } + } +} + +/// Serializes copying refreshed logins back within this process. +fn copy_back_lock() -> &'static tokio::sync::Mutex<()> { + static L: std::sync::OnceLock> = std::sync::OnceLock::new(); + L.get_or_init(tokio::sync::Mutex::default) +} + +/// A name no other run of this process (or another, mostly) has. +fn run_name() -> String { + static N: AtomicU64 = AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()); + format!( + "run-{}-{nanos}-{}", + std::process::id(), + N.fetch_add(1, Ordering::Relaxed) + ) +} + +/// The last component of `name`, or `image-` if it has none. +fn safe_name(name: &str, i: usize) -> String { + Path::new(name) + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .filter(|n| !n.starts_with('.')) + .unwrap_or_else(|| format!("image-{}", i + 1)) +} + +impl Exec { + /// The `codex exec` arguments for a run in `dir` with `images`. + fn argv(&self, dir: &Path, images: &[PathBuf]) -> Vec { + let s = |p: &Path| p.to_string_lossy().into_owned(); + let mut argv: Vec = [ + "exec", + "--ephemeral", + "--ignore-user-config", + "--ignore-rules", + "--skip-git-repo-check", + "--sandbox", + &self.sandbox, + "--color", + "never", + "--json", + "--cd", + ] + .map(String::from) + .into(); + argv.push(s(&dir.join("work"))); + argv.push("--output-last-message".into()); + argv.push(s(&dir.join("last.txt"))); + if let Some(m) = &self.model { + argv.extend(["--model".into(), m.clone()]); + } + for p in images { + argv.push(format!("--image={}", s(p))); + } + argv.extend(self.args.iter().cloned()); + argv.extend(["--".into(), "-".into()]); + argv + } + + /// The child's environment for a run in `dir`. + fn child_env(&self, dir: &Path) -> Vec<(String, String)> { + let s = |p: &Path| p.to_string_lossy().into_owned(); + let mut extra: Vec<(String, String)> = Vec::new(); + if !self.env.iter().any(|(k, _)| k == "LANG") { + extra.push(("LANG".into(), "C.UTF-8".into())); + } + extra.extend( + self.env + .iter() + .filter(|(k, _)| !matches!(k.as_str(), "HOME" | "CODEX_HOME" | "TMPDIR")) + .cloned(), + ); + extra.extend([ + ("HOME".into(), s(dir)), + ("CODEX_HOME".into(), s(&dir.join("home"))), + ("TMPDIR".into(), s(dir)), + ]); + crate::client::child_env(std::env::vars(), &extra) + } + + /// Run `codex exec` once on `input` (see the module docs). + pub async fn run(&self, input: &Input) -> Result { + let real = self.codex_home.join("auth.json"); + let read_login = || async { tokio::fs::read(&real).await.map_err(ExecError::Login) }; + // A run that may refresh the login goes alone; the next one reads + // the login it copied back. + let mut orig = read_login().await?; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX)); + let _alone = if may_refresh(&orig, now) { + let g = refresh_lock().lock().await; + orig = read_login().await?; + Some(g) + } else { + None + }; + let setup = |what| move |e| ExecError::Setup(what, e); + mkdir_private(&self.tmp_root) + .await + .map_err(setup("directory"))?; + let dir = RunDir(self.tmp_root.join(run_name())); + for d in ["", "home", "work", "refs"] { + tokio::fs::DirBuilder::new() + .mode(0o700) + .create(dir.0.join(d)) + .await + .map_err(setup("directory"))?; + } + let ours = dir.0.join("home/auth.json"); + write_private_atomic(&ours, &orig) + .await + .map_err(setup("login"))?; + let mut images = Vec::with_capacity(input.images.len()); + for (i, (name, data)) in input.images.iter().enumerate() { + let p = dir.0.join("refs").join(safe_name(name, i)); + write_private_atomic(&p, data) + .await + .map_err(setup("image"))?; + images.push(p); + } + let out = run_capture( + &self.command, + &self.argv(&dir.0, &images), + &dir.0.join("work"), + &self.child_env(&dir.0), + input.prompt.as_bytes(), + self.timeout, + ) + .await + .map_err(ExecError::Spawn)?; + let copied_back = { + let _g = copy_back_lock().lock().await; + crate::auth::copy_back(&Copied::new(real.clone(), ours, orig)).await + }; + let last_message = tokio::fs::read(dir.0.join("last.txt")) + .await + .unwrap_or_default(); + Ok(Run { + dir, + code: out.0, + stdout: out.1, + stderr: out.2, + last_message, + copied_back, + }) + } +} + +/// Run `command argv` in `cwd` with only `env`, `input` on stdin; its exit +/// code (`None`: killed after `timeout`), stdout and stderr. +async fn run_capture( + command: &str, + argv: &[String], + cwd: &Path, + env: &[(String, String)], + input: &[u8], + timeout: Duration, +) -> std::io::Result<(Option, Vec, Vec)> { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + let mut child = tokio::process::Command::new(command) + .args(argv) + .current_dir(cwd) + .env_clear() + .envs(env.iter().map(|(k, v)| (k, v))) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + if let Some(mut stdin) = child.stdin.take() { + let data = input.to_vec(); + tokio::spawn(async move { + let _ = stdin.write_all(&data).await; + }); + } + let read = |r: Option>| { + tokio::spawn(async move { + let mut b = Vec::new(); + if let Some(mut r) = r { + let _ = r.read_to_end(&mut b).await; + } + b + }) + }; + let t_out = read(child.stdout.take().map(|o| Box::new(o) as _)); + let t_err = read(child.stderr.take().map(|e| Box::new(e) as _)); + let code = if let Ok(s) = tokio::time::timeout(timeout, child.wait()).await { + Some(s?.code().unwrap_or(-1)) + } else { + let _ = child.kill().await; + None + }; + // A grandchild may still hold the pipes open after a kill: don't wait + // on it for long. + let grace = if code.is_some() { + timeout + } else { + Duration::from_secs(1) + }; + let collect = |t: tokio::task::JoinHandle>| async move { + let abort = t.abort_handle(); + tokio::time::timeout(grace, t).await.map_or_else( + |_| { + abort.abort(); + Vec::new() + }, + Result::unwrap_or_default, + ) + }; + let out = collect(t_out).await; + let err = collect(t_err).await; + Ok((code, out, err)) +} + +/// The regular files under `dir` (recursively; no symlinks) of at most +/// `max_bytes`, oldest first. +async fn files_oldest_first(dir: &Path, max_bytes: u64) -> Vec { + let mut files = Vec::new(); + let mut dirs = vec![dir.to_path_buf()]; + while let Some(d) = dirs.pop() { + let Ok(mut rd) = tokio::fs::read_dir(&d).await else { + continue; + }; + while let Ok(Some(e)) = rd.next_entry().await { + let Ok(ft) = e.file_type().await else { + continue; + }; + if ft.is_dir() { + dirs.push(e.path()); + } else if ft.is_file() { + match e.metadata().await { + Ok(m) if m.len() <= max_bytes => { + files.push((m.modified().unwrap_or(std::time::UNIX_EPOCH), e.path())); + } + _ => {} + } + } + } + } + files.sort(); + files.into_iter().map(|(_, p)| p).collect() +} + +/// Why a run failed, from its `--json` events and its log. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Failure { + /// The `ChatGPT` login is spent or revoked (its refresh token was used + /// up, expired or invalidated): someone must run `codex login` again. + LoginDead, + /// The service refused the login (a 401): what it said. Maybe a login. + LoginRefused(String), + /// The turn's failure, or the last error event. + Said(String), +} + +impl fmt::Display for Failure { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::LoginDead => f.write_str( + "Codex's ChatGPT login is no longer valid (its refresh token was used up or \ + revoked): run `codex login` again", + ), + Self::LoginRefused(m) => write!(f, "Codex's ChatGPT login was refused ({m})"), + Self::Said(m) => f.write_str(m), + } + } +} + +/// Why a `codex exec --json` run failed, from its events (`stdout`) and +/// its log (`stderr`) (pure): a dead login; else a 401; else the turn's +/// failure or the last error event (retry notices skipped), if any. +pub fn failure(stdout: &[u8], stderr: &[u8]) -> Option { + const DEAD: [&str; 6] = [ + "refresh_token_reused", + "refresh token was already used", + "refresh_token_expired", + "refresh_token_invalidated", + "Please log out and sign in again", + "Please try signing in again", + ]; + let err = String::from_utf8_lossy(stderr); + let out = String::from_utf8_lossy(stdout); + if DEAD.iter().any(|d| err.contains(d) || out.contains(d)) { + return Some(Failure::LoginDead); + } + let mut said = None; + for line in out.lines() { + let Ok(v) = serde_json::from_str::(line) else { + continue; + }; + let m = match v["type"].as_str() { + Some("turn.failed") => v["error"]["message"].as_str(), + Some("error") => v["message"].as_str(), + _ => None, + }; + if let Some(m) = m.filter(|m| !m.starts_with("Reconnecting...")) { + said = Some(m.to_owned()); + } + } + said.map(|m| { + if m.contains("401") { + Failure::LoginRefused(m) + } else { + Failure::Said(m) + } + }) +} + +/// A prompt that has Codex make exactly `count` images from `brief` and +/// do nothing else; `aspect_ratio` like `16:9`; `with_images` when the +/// run attaches reference images. +pub fn image_prompt( + brief: &str, + count: u64, + aspect_ratio: Option<&str>, + with_images: bool, +) -> String { + use fmt::Write as _; + let mut t = format!( + "Use your image generation tool to create exactly {count} image{} from the brief below", + if count == 1 { "" } else { "s" } + ); + if let Some(a) = aspect_ratio { + let _ = write!(t, ", aspect ratio {a} (width:height)"); + } + if with_images { + t.push_str(", using the attached image(s) as the reference the brief refers to"); + } + t.push_str( + ". Generate each image separately. Do nothing else: no shell commands, no reading or \ + writing files, no web searches, no questions. When the images are made, reply with \ + only the word done.\n\nBrief:\n", + ); + t.push_str(brief); + t +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn image_prompt_says_what() { + let t = image_prompt("a red apple", 2, Some("16:9"), true); + assert!(t.contains("exactly 2 images") && t.contains("16:9") && t.contains("attached")); + assert!(t.ends_with("Brief:\na red apple"), "{t}"); + let t = image_prompt("x", 1, None, false); + assert!( + t.contains("exactly 1 image ") && !t.contains("attached"), + "{t}" + ); + } + + #[test] + fn failure_says_why() { + // A spent login: stderr has the refresh error, the events only + // "unauthorized (401)". + let out = br#"{"type":"thread.started","thread_id":"t"} +{"type":"error","message":"Reconnecting... 2/5 (workspace routing discovery unauthorized (401))"} +{"type":"turn.failed","error":{"message":"workspace routing discovery unauthorized (401)"}}"#; + let err = b"ERROR codex_login::auth::manager: Failed to refresh token: Your access token \ +could not be refreshed because your refresh token was already used. Please log out and sign in again."; + assert_eq!(failure(out, err), Some(Failure::LoginDead)); + assert!(Failure::LoginDead.to_string().contains("codex login")); + assert_eq!( + failure(out, b""), + Some(Failure::LoginRefused( + "workspace routing discovery unauthorized (401)".into() + )) + ); + let out = br#"{"type":"turn.failed","error":{"message":"You've hit your usage limit."}}"#; + assert_eq!( + failure(out, b""), + Some(Failure::Said("You've hit your usage limit.".into())) + ); + assert_eq!(failure(b"not json", b"noise"), None); + } + + #[test] + fn names_and_args() { + assert_eq!(safe_name("../../etc/x.png", 0), "x.png"); + assert_eq!(safe_name("..", 1), "image-2"); + assert_eq!(safe_name(".hidden", 2), "image-3"); + let mut e = Exec::new("codex", "/h", "/t"); + e.model = Some("m".into()); + let a = e.argv(Path::new("/t/r"), &[PathBuf::from("/t/r/refs/a.png")]); + assert_eq!(a[0], "exec"); + assert!(a.windows(2).any(|w| w == ["--sandbox", "read-only"])); + assert!(a.windows(2).any(|w| w == ["--cd", "/t/r/work"])); + assert!(a.windows(2).any(|w| w == ["--model", "m"])); + assert!(a.contains(&"--image=/t/r/refs/a.png".to_owned())); + assert!(a.ends_with(&["--".into(), "-".into()])); + e.env = vec![ + ("PATH".into(), "/bin".into()), + ("HOME".into(), "/nope".into()), + ]; + let env = e.child_env(Path::new("/t/r")); + let get = |k: &str| { + env.iter() + .filter(|(n, _)| n == k) + .map(|(_, v)| v.as_str()) + .collect::>() + }; + assert_eq!(get("HOME"), ["/t/r"]); + assert_eq!(get("CODEX_HOME"), ["/t/r/home"]); + assert_eq!(get("PATH"), ["/bin"]); + assert_eq!(get("LANG"), ["C.UTF-8"]); + } + + #[tokio::test] + async fn oldest_first_regular_files_only() { + let d = tempfile::tempdir().unwrap(); + let g = d.path().join("t1"); + tokio::fs::create_dir_all(&g).await.unwrap(); + tokio::fs::write(g.join("a.png"), b"1").await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + tokio::fs::write(g.join("b.jpg"), b"2").await.unwrap(); + tokio::fs::write(g.join("big.png"), b"0123456789") + .await + .unwrap(); + std::os::unix::fs::symlink("/etc/passwd", g.join("link.png")).unwrap(); + let names: Vec<_> = files_oldest_first(d.path(), 4) + .await + .iter() + .map(|p| p.file_name().unwrap().to_string_lossy().into_owned()) + .collect(); + assert_eq!(names, ["a.png", "b.jpg"]); + assert_eq!( + files_oldest_first(&d.path().join("none"), 4).await, + Vec::::new() + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 9fd12f6..a5201fb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ //! its `VERSION` are in `schema/`, [`CODEX_VERSION`] is the version), [`rpc`] //! frames its messages, [`client`] runs the child and matches responses to //! requests, and [`auth`] gives a session its own copy of `auth.json`. +//! [`exec`] runs one-shot `codex exec` jobs in a private `CODEX_HOME`. //! //! The README follows (its example is compiled as a doctest). //! @@ -26,6 +27,7 @@ extern crate alloc; pub mod auth; pub mod client; +pub mod exec; #[allow( missing_docs, clippy::pedantic, diff --git a/tests/exec.rs b/tests/exec.rs new file mode 100644 index 0000000..97c6441 --- /dev/null +++ b/tests/exec.rs @@ -0,0 +1,114 @@ +//! `exec::Exec::run` against a fake `codex` (`tests/fake_codex_exec`, sh): +//! the private home, the image in and out, the login copied back, the +//! failure, the timeout, and the run's directory removed. + +// `#[tokio::test]` expands to `block_on`. +#![allow(clippy::disallowed_methods)] + +use core::time::Duration; + +use codexwire::exec::{Exec, Failure, Input}; + +const FAKE: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fake_codex_exec"); + +fn setup() -> (tempfile::TempDir, Exec) { + let d = tempfile::tempdir().expect("tempdir"); + let real = d.path().join("codex-home"); + std::fs::create_dir_all(&real).expect("mkdir"); + std::fs::write(real.join("auth.json"), br#"{"tokens":"old"}"#).expect("write"); + std::fs::write(real.join("history.jsonl"), b"secret history").expect("write"); + let mut e = Exec::new(FAKE, &real, d.path().join("state/codex")); + e.timeout = Duration::from_secs(20); + e.env = vec![ + ("PATH".into(), std::env::var("PATH").unwrap_or_default()), + ( + "REAL_CODEX_HOME".into(), + real.to_string_lossy().into_owned(), + ), + ]; + (d, e) +} + +fn left(e: &Exec) -> usize { + std::fs::read_dir(&e.tmp_root).expect("read_dir").count() +} + +#[tokio::test] +async fn runs_in_a_private_home() { + let (_d, e) = setup(); + let input = Input { + prompt: "make an image".into(), + images: vec![("../ref-1.png".into(), b"\x89PNG\r\n\x1a\n0000".to_vec())], + }; + let run = e.run(&input).await.unwrap(); + assert_eq!( + run.code, + Some(0), + "{}", + String::from_utf8_lossy(&run.stderr) + ); + assert_eq!(run.last_message, b"done"); + let imgs = run.generated_images(1 << 20).await; + assert_eq!(imgs.len(), 1); + assert_eq!(std::fs::read(&imgs[0]).unwrap(), b"\x89PNG\r\n\x1a\n0000"); + assert_eq!( + run.generated_images(4).await, + Vec::::new(), + "too big" + ); + // The refreshed login is back. + assert!(*run.copied_back.as_ref().unwrap()); + assert_eq!( + std::fs::read(e.codex_home.join("auth.json")).unwrap(), + br#"{"tokens":"new"}"# + ); + assert_eq!(left(&e), 1); + run.remove().await.unwrap(); + assert_eq!(left(&e), 0); +} + +#[tokio::test] +async fn says_why_it_failed() { + let (_d, e) = setup(); + let run = e + .run(&Input { + prompt: "fail".into(), + images: vec![], + }) + .await + .unwrap(); + assert_eq!(run.code, Some(3)); + assert_eq!( + run.failure(), + Some(Failure::Said("You've hit your usage limit.".into())) + ); + assert!(!*run.copied_back.as_ref().unwrap(), "not refreshed"); + assert_eq!( + std::fs::read(e.codex_home.join("auth.json")).unwrap(), + br#"{"tokens":"old"}"# + ); + drop(run); + assert_eq!(left(&e), 0, "dropped: removed"); +} + +#[tokio::test] +async fn times_out() { + let (_d, mut e) = setup(); + e.timeout = Duration::from_millis(300); + let run = e + .run(&Input { + prompt: "hang".into(), + images: vec![], + }) + .await + .unwrap(); + assert_eq!(run.code, None); +} + +#[tokio::test] +async fn needs_a_login() { + let (_d, mut e) = setup(); + e.codex_home = e.codex_home.join("missing"); + let err = e.run(&Input::default()).await.unwrap_err(); + assert!(matches!(err, codexwire::exec::ExecError::Login(_)), "{err}"); +} diff --git a/tests/fake_codex_exec b/tests/fake_codex_exec new file mode 100755 index 0000000..b9139f4 --- /dev/null +++ b/tests/fake_codex_exec @@ -0,0 +1,31 @@ +#!/bin/sh +# A fake `codex exec` for tests/exec.rs: checks it runs in a private +# CODEX_HOME holding only a copy of the login, then acts per the prompt. +[ "$1" = exec ] || exit 10 +case "$CODEX_HOME" in "$REAL_CODEX_HOME"*) exit 9 ;; esac +[ -f "$CODEX_HOME/history.jsonl" ] && exit 8 +[ "$HOME" = "$TMPDIR" ] || exit 6 +last="" +img="" +while [ $# -gt 0 ]; do + case "$1" in + --output-last-message) last=$2; shift ;; + --image=*) img=${1#--image=} ;; + esac + shift +done +prompt=$(cat) +case "$prompt" in +*image*) + [ -n "$img" ] && [ -f "$img" ] || exit 5 + mkdir -p "$CODEX_HOME/generated_images/t1" + cp "$img" "$CODEX_HOME/generated_images/t1/exec-1.png" + printf '{"tokens":"new"}' >"$CODEX_HOME/auth.json" + printf 'done' >"$last" + ;; +*fail*) + echo '{"type":"turn.failed","error":{"message":"You'"'"'ve hit your usage limit."}}' + exit 3 + ;; +*hang*) sleep 30 ;; +esac From fdac621df4026ebd316d4d1fc0b7b90ad4c5aac9 Mon Sep 17 00:00:00 2001 From: Paolo Santos Date: Fri, 6 Nov 2026 23:32:18 +0800 Subject: [PATCH 2/2] CI: note tests/exec.rs's fake codex --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 39df722..7ba7434 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -66,7 +66,7 @@ jobs: - uses: Swatinem/rust-cache@v2 # Includes tests/corpus.rs (recorded app-server output) and # tests/spawn.rs (the fake codex, tests/fake_codex: python3, which the - # runner has). + # runner has) and tests/exec.rs (tests/fake_codex_exec: sh). - run: cargo test --all-targets --locked # The README's example, compiled (lib.rs includes the README). - run: cargo test --doc --locked