omni_dev/cli/
codex_wrap.rs

1//! `omni-dev codex-wrap` — the Codex TUI on a private app-server, observed.
2//!
3//! Runs the Codex TUI against a **private** Codex app-server that this process
4//! starts and owns, and reports the exact thread status it polls from that
5//! server to the daemon's `sessions` service (#1910, ADR-0088). The Codex
6//! analogue of [`omni-dev claude-wrap`](super::claude_wrap).
7//!
8//! The Codex VS Code extension and Desktop each run their own private stdio
9//! app-server, so no outside process can read their threads' status. A session
10//! omni-dev launches can be different: this wrapper starts
11//! `codex app-server --listen unix://<runtime-dir>/codex-wrap-<pid>.sock`, runs
12//! `codex --remote unix://… <args>` in the foreground on the caller's terminal,
13//! and polls that server with [`crate::sessions::codex_app_server`] — never
14//! subscribing and never answering a server request, so it cannot see, answer
15//! or duplicate an approval. It attaches only to the server it started.
16//!
17//! **Fail-open is the hard rule**, as for `claude-wrap`: if the app-server does
18//! not come up, this process is replaced by plain `codex <args>`; if the
19//! observer cannot connect, the TUI still runs, unobserved. An invocation the
20//! remote TUI does not serve (a non-interactive subcommand, no terminal, or an
21//! explicit `--remote`) is passed straight to `codex` the same way.
22//!
23//! The server lives exactly as long as the wrapper: when the TUI exits, the
24//! wrapper ends the sessions it reported, stops the server and removes its
25//! socket. Nothing is logged or persisted.
26
27use std::io::IsTerminal;
28use std::os::unix::process::{CommandExt, ExitStatusExt};
29use std::path::{Path, PathBuf};
30use std::process::{ExitStatus, Stdio};
31use std::time::Duration;
32
33use anyhow::{bail, Result};
34use clap::Parser;
35use serde_json::{json, Value};
36use tokio::process::{Child, Command};
37use tokio_util::sync::CancellationToken;
38
39use crate::daemon::client::DaemonClient;
40use crate::daemon::paths;
41use crate::daemon::protocol::DaemonEnvelope;
42use crate::daemon::server;
43use crate::sessions::codex_app_server::{AppServerClient, Report, StatusTracker};
44
45/// The `sessions` service routing key on the daemon control socket.
46const SERVICE: &str = "sessions";
47
48/// How long a fire-and-forget report waits for the daemon.
49const REPORT_TIMEOUT: Duration = Duration::from_secs(2);
50
51/// How often the server is polled for thread status.
52const POLL_INTERVAL: Duration = Duration::from_secs(1);
53
54/// How long connecting to the server (socket, WebSocket upgrade, `initialize`)
55/// may take before the wrapper gives up observing.
56const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
57
58/// Consecutive failed polls after which the server is taken to be gone.
59const MAX_POLL_FAILURES: u32 = 3;
60
61/// How long the app-server may take to start listening before the wrapper
62/// gives up on it and runs plain `codex` instead.
63const SERVER_READY_TIMEOUT: Duration = Duration::from_secs(10);
64
65/// How long the app-server gets to exit on `SIGTERM` before it is killed.
66const SERVER_STOP_GRACE: Duration = Duration::from_secs(2);
67
68/// Exit status for a child killed by a signal (`128 + signo`).
69const SIGNAL_EXIT_BASE: i32 = 128;
70
71/// `codex` subcommands that are not the interactive TUI, so `--remote` does not
72/// apply and the wrapper passes them straight through. `resume` and `fork` open
73/// the TUI and are wrapped.
74const PASSTHROUGH_SUBCOMMANDS: &[&str] = &[
75    "agents",
76    "app",
77    "app-server",
78    "apply",
79    "a",
80    "archive",
81    "cloud",
82    "completion",
83    "debug",
84    "delete",
85    "doctor",
86    "e",
87    "exec",
88    "exec-server",
89    "features",
90    "help",
91    "login",
92    "logout",
93    "mcp",
94    "migrate-rollouts",
95    "plugin",
96    "remote-control",
97    "review",
98    "sandbox",
99    "unarchive",
100    "update",
101];
102
103/// Runs the Codex TUI against a private app-server, reporting its exact state.
104#[derive(Parser)]
105pub struct CodexWrapCommand {
106    /// Path to the daemon control socket. Defaults to the standard per-user path.
107    #[arg(long, value_name = "PATH")]
108    pub socket: Option<PathBuf>,
109
110    /// The Codex command and its arguments, after a `--` separator
111    /// (e.g. `-- codex`, `-- codex resume <id>`).
112    #[arg(
113        trailing_var_arg = true,
114        allow_hyphen_values = true,
115        value_name = "CMD"
116    )]
117    pub argv: Vec<String>,
118}
119
120impl CodexWrapCommand {
121    /// Executes the wrapper and exits with the TUI's own status — like
122    /// `claude-wrap`, via [`std::process::exit`], so a wrapped run writes no
123    /// request-log record.
124    pub async fn execute(self) -> Result<()> {
125        let code = self.run().await?;
126        std::process::exit(code)
127    }
128
129    async fn run(self) -> Result<i32> {
130        let Some((program, args)) = self.argv.split_first() else {
131            bail!("codex-wrap needs a command to run, e.g. `omni-dev codex-wrap -- codex`");
132        };
133        if !std::io::stdout().is_terminal() || !is_tui_invocation(args) {
134            return Err(exec_replace(program, args));
135        }
136        let Some(listen) = listen_path() else {
137            return Err(exec_replace(program, args));
138        };
139        let Some(server) = start_server(program, &listen).await else {
140            return Err(exec_replace(program, args));
141        };
142        Ok(wrap(program, args, &listen, server, self.socket).await)
143    }
144}
145
146/// Whether `args` (after the `codex` program) open the interactive TUI the
147/// `--remote` flag serves: no pass-through subcommand anywhere among the
148/// positionals, and no `--remote` of the caller's own.
149fn is_tui_invocation(args: &[String]) -> bool {
150    !args.iter().any(|arg| {
151        arg == "--remote"
152            || arg.starts_with("--remote=")
153            || PASSTHROUGH_SUBCOMMANDS.contains(&arg.as_str())
154    })
155}
156
157/// The private server's socket: `<runtime-dir>/codex-wrap-<pid>.sock`, in the
158/// daemon's `0700` directory. `None` (so the wrapper passes through) when the
159/// directory cannot be made or the path is too long for a `sockaddr_un`.
160fn listen_path() -> Option<PathBuf> {
161    let dir = paths::runtime_dir().ok()?;
162    paths::ensure_dir_0700(&dir).ok()?;
163    let path = dir.join(format!("codex-wrap-{}.sock", std::process::id()));
164    (path.as_os_str().len() < paths::MAX_SOCKET_PATH_LEN).then_some(path)
165}
166
167/// Starts `<program> app-server --listen unix://<listen>` and waits for it to
168/// accept connections. `None` if it exits or does not come up in time — the
169/// caller then runs plain `codex`.
170async fn start_server(program: &str, listen: &Path) -> Option<Child> {
171    let _ = std::fs::remove_file(listen);
172    let mut child = Command::new(program)
173        .arg("app-server")
174        .arg("--listen")
175        .arg(format!("unix://{}", listen.display()))
176        .stdin(Stdio::null())
177        .stdout(Stdio::null())
178        .stderr(Stdio::null())
179        .kill_on_drop(true)
180        .spawn()
181        .ok()?;
182    let deadline = tokio::time::Instant::now() + SERVER_READY_TIMEOUT;
183    while tokio::time::Instant::now() < deadline {
184        if tokio::net::UnixStream::connect(listen).await.is_ok() {
185            return Some(child);
186        }
187        if matches!(child.try_wait(), Ok(Some(_))) {
188            break;
189        }
190        tokio::time::sleep(Duration::from_millis(50)).await;
191    }
192    stop_server(&mut child, listen).await;
193    None
194}
195
196/// Runs the TUI against the server, observes until it exits, then cleans up.
197/// Returns the TUI's exit code.
198async fn wrap(
199    program: &str,
200    args: &[String],
201    listen: &Path,
202    mut server: Child,
203    socket: Option<PathBuf>,
204) -> i32 {
205    let tui = std::process::Command::new(program)
206        .arg("--remote")
207        .arg(format!("unix://{}", listen.display()))
208        .args(args)
209        .spawn();
210    let Ok(mut tui) = tui else {
211        stop_server(&mut server, listen).await;
212        // The same program just failed to spawn; report it as the shell would.
213        return 127;
214    };
215    let stop = CancellationToken::new();
216    let observer = tokio::spawn(observe(listen.to_path_buf(), socket, stop.clone()));
217    let signals = tokio::spawn(forward_signals(tui.id()));
218    // The TUI owns the terminal; wait for it off the async runtime.
219    let status = tokio::task::spawn_blocking(move || tui.wait())
220        .await
221        .ok()
222        .and_then(Result::ok);
223    signals.abort();
224    stop.cancel();
225    let _ = observer.await;
226    stop_server(&mut server, listen).await;
227    status.map_or(1, exit_code)
228}
229
230/// Polls the server, re-asserting every live session's state each poll, and
231/// ends every reported session once `stop` fires or the server stops
232/// answering. Connecting and each poll race `stop`, so a stalled server can
233/// never keep the wrapper from exiting.
234async fn observe(listen: PathBuf, socket: Option<PathBuf>, stop: CancellationToken) {
235    let Ok(socket) = server::resolve_socket(socket) else {
236        return;
237    };
238    let connect = tokio::time::timeout(CONNECT_TIMEOUT, AppServerClient::connect(&listen));
239    let mut client = tokio::select! {
240        () = stop.cancelled() => return,
241        connected = connect => match connected {
242            Ok(Ok(client)) => client,
243            _ => return,
244        },
245    };
246    let mut tracker = StatusTracker::default();
247    let mut poll = tokio::time::interval(POLL_INTERVAL);
248    let mut failures = 0;
249    loop {
250        tokio::select! {
251            () = stop.cancelled() => break,
252            _ = poll.tick() => {}
253        }
254        let polled = tokio::select! {
255            () = stop.cancelled() => break,
256            polled = client.snapshot() => polled,
257        };
258        if let Ok(snapshot) = polled {
259            failures = 0;
260            send(&socket, tracker.update(&snapshot)).await;
261        } else {
262            // A closed connection fails every call at once, so a few in a row
263            // means the server is gone (the TUI will follow it); one slow poll
264            // does not.
265            failures += 1;
266            if failures >= MAX_POLL_FAILURES {
267                break;
268            }
269        }
270    }
271    send(&socket, tracker.finish()).await;
272}
273
274/// Sends reports to the daemon, fire-and-forget: a missing or wedged daemon is
275/// a silent no-op.
276async fn send(socket: &Path, reports: Vec<Report>) {
277    for report in reports {
278        let (op, payload): (&str, Value) = match report {
279            Report::Observe(request) => match serde_json::to_value(request) {
280                Ok(payload) => ("observe", payload),
281                Err(_) => continue,
282            },
283            Report::End(id) => ("end", json!({ "session_id": id })),
284        };
285        let envelope = DaemonEnvelope::service(SERVICE, op, payload);
286        let _ =
287            tokio::time::timeout(REPORT_TIMEOUT, DaemonClient::new(socket).request(envelope)).await;
288    }
289}
290
291/// Stops the server — `SIGTERM`, then a kill after [`SERVER_STOP_GRACE`] — and
292/// removes its socket. The server was started by this process and is its child,
293/// so the pid is its own.
294async fn stop_server(server: &mut Child, listen: &Path) {
295    if let Some(pid) = server.id().and_then(|pid| i32::try_from(pid).ok()) {
296        let _ = nix::sys::signal::kill(
297            nix::unistd::Pid::from_raw(pid),
298            nix::sys::signal::Signal::SIGTERM,
299        );
300    }
301    if tokio::time::timeout(SERVER_STOP_GRACE, server.wait())
302        .await
303        .is_err()
304    {
305        let _ = server.kill().await;
306    }
307    let _ = std::fs::remove_file(listen);
308}
309
310/// Relays `SIGTERM`/`SIGHUP` to the TUI, so a caller that signals the wrapper
311/// stops Codex rather than orphaning it, and swallows `SIGINT` — the TUI shares
312/// the terminal's process group and gets it directly; the wrapper must survive
313/// to clean up.
314async fn forward_signals(child_pid: u32) {
315    use nix::sys::signal::Signal;
316    use tokio::signal::unix::{signal, SignalKind};
317
318    let Ok(pid) = i32::try_from(child_pid) else {
319        return;
320    };
321    let pid = nix::unistd::Pid::from_raw(pid);
322    let (Ok(mut terminate), Ok(mut hangup), Ok(mut interrupt)) = (
323        signal(SignalKind::terminate()),
324        signal(SignalKind::hangup()),
325        signal(SignalKind::interrupt()),
326    ) else {
327        return;
328    };
329    loop {
330        let relay = tokio::select! {
331            _ = terminate.recv() => Some(Signal::SIGTERM),
332            _ = hangup.recv() => Some(Signal::SIGHUP),
333            _ = interrupt.recv() => None,
334        };
335        if let Some(signal) = relay {
336            let _ = nix::sys::signal::kill(pid, signal);
337        }
338    }
339}
340
341/// Replaces this process with `program`, returning the error only if the `exec`
342/// itself failed.
343fn exec_replace(program: &str, args: &[String]) -> anyhow::Error {
344    let error = std::process::Command::new(program).args(args).exec();
345    anyhow::Error::new(error).context(format!("failed to exec {program}"))
346}
347
348/// The exit code to leave with: the TUI's own, or `128 + signo`.
349fn exit_code(status: ExitStatus) -> i32 {
350    status
351        .code()
352        .or_else(|| status.signal().map(|signal| SIGNAL_EXIT_BASE + signal))
353        .unwrap_or(1)
354}
355
356#[cfg(test)]
357#[allow(clippy::unwrap_used, clippy::expect_used)]
358mod tests {
359    use super::*;
360
361    fn args(list: &[&str]) -> Vec<String> {
362        list.iter().map(|s| (*s).to_string()).collect()
363    }
364
365    #[test]
366    fn only_tui_invocations_are_wrapped() {
367        assert!(is_tui_invocation(&args(&[])));
368        assert!(is_tui_invocation(&args(&["fix the build"])));
369        assert!(is_tui_invocation(&args(&[
370            "-m",
371            "gpt-5.5",
372            "--no-alt-screen"
373        ])));
374        assert!(is_tui_invocation(&args(&["resume", "019a-id"])));
375        assert!(is_tui_invocation(&args(&["fork", "--last"])));
376        assert!(!is_tui_invocation(&args(&["exec", "say hi"])));
377        assert!(!is_tui_invocation(&args(&["-c", "x=1", "login"])));
378        assert!(!is_tui_invocation(&args(&["--remote", "ws://h:1"])));
379        assert!(!is_tui_invocation(&args(&["--remote=ws://h:1"])));
380    }
381
382    #[test]
383    fn codex_wrap_parses_its_trailing_command() {
384        #[derive(Parser)]
385        struct Wrapper {
386            #[command(flatten)]
387            cmd: CodexWrapCommand,
388        }
389        let parsed = Wrapper::parse_from(["x", "--", "codex", "resume", "--last"]);
390        assert_eq!(parsed.cmd.argv, args(&["codex", "resume", "--last"]));
391        assert!(parsed.cmd.socket.is_none());
392    }
393
394    #[test]
395    fn exit_codes_follow_the_shell_convention() {
396        assert_eq!(exit_code(ExitStatus::from_raw(3 << 8)), 3);
397        assert_eq!(exit_code(ExitStatus::from_raw(9)), SIGNAL_EXIT_BASE + 9);
398    }
399
400    #[tokio::test]
401    async fn a_server_that_never_listens_is_abandoned_for_plain_codex() {
402        // `false app-server …` exits at once without listening.
403        let tmp = tempfile::tempdir_in("/tmp").unwrap();
404        let listen = tmp.path().join("as.sock");
405        assert!(start_server("false", &listen).await.is_none());
406        assert!(!listen.exists());
407        // A program that cannot even be spawned is the same pass-through.
408        assert!(start_server("/nonexistent/codex", &listen).await.is_none());
409    }
410
411    #[tokio::test]
412    async fn a_run_without_a_command_is_an_error() {
413        let cmd = CodexWrapCommand {
414            socket: None,
415            argv: Vec::new(),
416        };
417        let err = cmd.run().await.unwrap_err();
418        assert!(err.to_string().contains("needs a command"), "{err}");
419    }
420}
omni_dev/cli/codex_wrap.rs

omni_dev/cli/
codex_wrap.rs
