use std::io::IsTerminal;
use std::os::unix::process::{CommandExt, ExitStatusExt};
use std::path::{Path, PathBuf};
use std::process::{ExitStatus, Stdio};
use std::time::Duration;
use anyhow::{anyhow, bail, Context, Result};
use clap::Parser;
use serde_json::{json, Value};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::mpsc;
use crate::daemon::client::DaemonClient;
use crate::daemon::protocol::DaemonEnvelope;
use crate::daemon::server;
use crate::sessions::stream::{Direction, StreamTracker};
const SERVICE: &str = "sessions";
const REPORT_TIMEOUT: Duration = Duration::from_secs(2);
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
const TEE_CAPACITY: usize = 256;
const MAX_LINE_BYTES: usize = 1024 * 1024;
const PUMP_BUF_BYTES: usize = 64 * 1024;
const SIGNAL_EXIT_BASE: i32 = 128;
#[derive(Parser)]
pub struct ClaudeWrapCommand {
#[arg(long, value_name = "PATH")]
pub socket: Option<PathBuf>,
#[arg(
trailing_var_arg = true,
allow_hyphen_values = true,
value_name = "CMD"
)]
pub argv: Vec<String>,
}
impl ClaudeWrapCommand {
pub async fn execute(self) -> Result<()> {
let code = self.run().await?;
std::process::exit(code)
}
async fn run(self) -> Result<i32> {
let Some((program, args)) = self.argv.split_first() else {
bail!(
"claude-wrap needs a command to run, e.g. \
`omni-dev claude-wrap -- claude --output-format stream-json`"
);
};
if std::io::stdout().is_terminal() {
return Err(exec_replace(program, args));
}
wrap(program, args, self.socket).await
}
}
fn exec_replace(program: &str, args: &[String]) -> anyhow::Error {
let error = std::process::Command::new(program).args(args).exec();
anyhow::Error::new(error).context(format!("failed to exec {program}"))
}
async fn wrap(program: &str, args: &[String], socket: Option<PathBuf>) -> Result<i32> {
wrap_io(
program,
args,
socket,
tokio::io::stdin(),
tokio::io::stdout(),
)
.await
}
async fn wrap_io<R, W>(
program: &str,
args: &[String],
socket: Option<PathBuf>,
input: R,
output: W,
) -> Result<i32>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let mut child = tokio::process::Command::new(program)
.args(args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.with_context(|| format!("failed to spawn {program}"))?;
let child_stdin = child
.stdin
.take()
.ok_or_else(|| anyhow!("failed to capture the wrapped process's stdin"))?;
let child_stdout = child
.stdout
.take()
.ok_or_else(|| anyhow!("failed to capture the wrapped process's stdout"))?;
let child_pid = child.id();
let (tee, lines) = mpsc::channel::<(Direction, String)>(TEE_CAPACITY);
let observer = tokio::spawn(observe(lines, socket, KEEPALIVE_INTERVAL));
let signals = tokio::spawn(forward_signals(child_pid));
let from_child = tokio::spawn(pump(
child_stdout,
output,
Direction::FromClaude,
tee.clone(),
));
let to_child = tokio::spawn(pump(input, child_stdin, Direction::ToClaude, tee.clone()));
drop(tee);
let _ = from_child.await;
to_child.abort();
let _ = to_child.await;
signals.abort();
let status = child
.wait()
.await
.context("failed to wait for the wrapped process")?;
let _ = observer.await;
Ok(exit_code(status))
}
async fn pump<R, W>(
mut reader: R,
mut writer: W,
direction: Direction,
tee: mpsc::Sender<(Direction, String)>,
) where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let mut buffer = vec![0u8; PUMP_BUF_BYTES];
let mut line = Vec::new();
loop {
let read = match reader.read(&mut buffer).await {
Ok(0) | Err(_) => break,
Ok(read) => read,
};
let chunk = &buffer[..read];
if writer.write_all(chunk).await.is_err() || writer.flush().await.is_err() {
break;
}
tee_chunk(chunk, direction, &tee, &mut line);
}
let _ = writer.shutdown().await;
}
fn tee_chunk(
chunk: &[u8],
direction: Direction,
tee: &mpsc::Sender<(Direction, String)>,
line: &mut Vec<u8>,
) {
for &byte in chunk {
if byte == b'\n' {
if line.len() < MAX_LINE_BYTES {
if let Ok(text) = std::str::from_utf8(line) {
let _ = tee.try_send((direction, text.to_string()));
}
}
line.clear();
} else if line.len() < MAX_LINE_BYTES {
line.push(byte);
}
}
}
async fn observe(
mut lines: mpsc::Receiver<(Direction, String)>,
socket: Option<PathBuf>,
every: Duration,
) {
let Ok(socket) = server::resolve_socket(socket) else {
return;
};
let mut tracker = StreamTracker::new();
let mut keepalive = tokio::time::interval(every);
keepalive.tick().await;
loop {
let observed = tokio::select! {
line = lines.recv() => match line {
Some((direction, text)) => tracker.observe_line(direction, &text),
None => break,
},
_ = keepalive.tick() => tracker.keepalive(),
};
if let Some(request) = observed {
if let Ok(payload) = serde_json::to_value(request) {
report(&socket, "observe", payload).await;
}
}
}
if let Some(session_id) = tracker.session_id() {
report(&socket, "end", json!({ "session_id": session_id })).await;
}
}
async fn report(socket: &Path, op: &str, payload: Value) {
let envelope = DaemonEnvelope::service(SERVICE, op, payload);
let _ = tokio::time::timeout(REPORT_TIMEOUT, DaemonClient::new(socket).request(envelope)).await;
}
async fn forward_signals(child_pid: Option<u32>) {
let Some(pid) = child_pid.and_then(|pid| i32::try_from(pid).ok()) else {
return;
};
let pid = nix::unistd::Pid::from_raw(pid);
let (Ok(mut terminate), Ok(mut interrupt)) = (
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()),
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()),
) else {
return;
};
loop {
let signal = tokio::select! {
_ = terminate.recv() => nix::sys::signal::Signal::SIGTERM,
_ = interrupt.recv() => nix::sys::signal::Signal::SIGINT,
};
let _ = nix::sys::signal::kill(pid, signal);
}
}
fn exit_code(status: ExitStatus) -> i32 {
status
.code()
.or_else(|| status.signal().map(|signal| SIGNAL_EXIT_BASE + signal))
.unwrap_or(1)
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context as TaskContext, Poll};
use tokio::io::AsyncBufReadExt;
use tokio::net::UnixListener;
use super::*;
#[derive(Clone, Default)]
struct Sink(Arc<Mutex<Vec<u8>>>);
impl Sink {
fn contents(&self) -> String {
String::from_utf8(self.0.lock().unwrap().clone()).unwrap()
}
}
impl AsyncWrite for Sink {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
self.0.lock().unwrap().extend_from_slice(buf);
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
}
async fn wrap_script(script: &str, socket: Option<PathBuf>) -> (i32, String) {
let sink = Sink::default();
let code = wrap_io(
"/bin/sh",
&["-c".to_string(), script.to_string()],
socket,
tokio::io::empty(),
sink.clone(),
)
.await
.unwrap();
(code, sink.contents())
}
fn parse(args: &[&str]) -> ClaudeWrapCommand {
let mut argv = vec!["claude-wrap"];
argv.extend_from_slice(args);
ClaudeWrapCommand::try_parse_from(argv).unwrap()
}
fn fake_daemon() -> (
tempfile::TempDir,
PathBuf,
Arc<tokio::sync::Mutex<Vec<Value>>>,
) {
let dir = tempfile::tempdir_in("/tmp").unwrap();
let socket = dir.path().join("d.sock");
let listener = UnixListener::bind(&socket).unwrap();
let seen = Arc::new(tokio::sync::Mutex::new(Vec::new()));
let recorder = Arc::clone(&seen);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
return;
};
let recorder = Arc::clone(&recorder);
tokio::spawn(async move {
let (read, mut write) = stream.into_split();
let mut line = String::new();
if tokio::io::BufReader::new(read)
.read_line(&mut line)
.await
.is_ok()
{
if let Ok(value) = serde_json::from_str::<Value>(&line) {
recorder.lock().await.push(value);
}
}
let _ = write.write_all(b"{\"ok\":true,\"payload\":{}}\n").await;
});
}
});
(dir, socket, seen)
}
#[test]
fn trailing_args_are_captured_verbatim_after_a_separator() {
let cmd = parse(&["--", "node", "cli.js", "--output-format", "stream-json"]);
assert_eq!(
cmd.argv,
vec!["node", "cli.js", "--output-format", "stream-json"]
);
assert!(cmd.socket.is_none());
}
#[test]
fn the_socket_override_is_parsed_before_the_separator() {
let cmd = parse(&["--socket", "/tmp/d.sock", "--", "claude", "-p"]);
assert_eq!(cmd.socket.unwrap(), PathBuf::from("/tmp/d.sock"));
assert_eq!(cmd.argv, vec!["claude", "-p"]);
}
#[tokio::test]
async fn an_empty_command_is_a_usage_error() {
let error = parse(&[]).run().await.unwrap_err();
assert!(error.to_string().contains("needs a command to run"));
}
#[tokio::test]
async fn signal_forwarding_gives_up_when_there_is_no_child() {
forward_signals(None).await;
}
#[tokio::test]
async fn the_child_is_wrapped_losslessly_and_its_exit_code_survives() {
let (_dir, socket, _seen) = fake_daemon();
let (code, forwarded) =
wrap_script("printf 'hello\\nnot json\\ntrailing'; exit 7", Some(socket)).await;
assert_eq!(code, 7);
assert_eq!(forwarded, "hello\nnot json\ntrailing");
}
#[tokio::test]
async fn stream_state_transitions_reach_the_daemon() {
let (_dir, socket, seen) = fake_daemon();
let script = concat!(
r#"printf '{"type":"system","subtype":"init","session_id":"s-1","cwd":"/w"}\n';"#,
r#"printf '{"type":"assistant","session_id":"s-1"}\n';"#,
r#"printf '{"type":"control_request","request_id":"r1","request":{"subtype":"can_use_tool"}}\n';"#,
r#"printf '{"type":"result"}\n'"#,
);
let (code, _forwarded) = wrap_script(script, Some(socket)).await;
assert_eq!(code, 0);
let envelopes = seen.lock().await.clone();
let states: Vec<String> = envelopes
.iter()
.filter(|e| e["op"] == "observe")
.filter_map(|e| e["payload"]["event"]["stream_state"].as_str())
.map(ToString::to_string)
.collect();
assert_eq!(
states,
vec!["idle", "working", "waiting_for_permission", "idle"]
);
assert_eq!(envelopes[0]["service"], "sessions");
assert_eq!(envelopes[0]["payload"]["session_id"], "s-1");
assert_eq!(envelopes[0]["payload"]["cwd"], "/w");
let ended = envelopes.last().unwrap();
assert_eq!(ended["op"], "end");
assert_eq!(ended["payload"]["session_id"], "s-1");
}
#[tokio::test]
async fn a_missing_daemon_never_affects_the_child() {
let dir = tempfile::tempdir_in("/tmp").unwrap();
let script = concat!(
r#"printf '{"type":"system","subtype":"init","session_id":"s-1"}\n';"#,
r#"printf 'not json at all\n';"#,
r#"exit 3"#,
);
let (code, forwarded) = wrap_script(script, Some(dir.path().join("absent.sock"))).await;
assert_eq!(code, 3);
assert!(forwarded.ends_with("not json at all\n"));
}
#[tokio::test]
async fn a_child_killed_by_a_signal_reports_the_shell_convention() {
let (_dir, socket, _seen) = fake_daemon();
let (code, _forwarded) = wrap_script("kill -TERM $$", Some(socket)).await;
assert_eq!(code, SIGNAL_EXIT_BASE + 15);
}
#[tokio::test]
async fn spawning_a_missing_program_is_an_error_not_a_panic() {
let error = wrap_io(
"/nonexistent/omni-dev-test-binary",
&[],
None,
tokio::io::empty(),
Sink::default(),
)
.await
.unwrap_err();
assert!(error.to_string().contains("failed to spawn"));
}
#[test]
fn over_long_and_non_utf8_lines_are_dropped_but_still_forwarded() {
let (tee, mut lines) = mpsc::channel(4);
let mut buffer = Vec::new();
let mut chunk = vec![b'x'; MAX_LINE_BYTES + 1];
chunk.push(b'\n');
chunk.extend_from_slice(&[0xff, 0xfe, b'\n']);
chunk.extend_from_slice(b"{\"type\":\"result\"}\n");
tee_chunk(&chunk, Direction::FromClaude, &tee, &mut buffer);
let (direction, text) = lines.try_recv().unwrap();
assert_eq!(direction, Direction::FromClaude);
assert_eq!(text, r#"{"type":"result"}"#);
assert!(lines.try_recv().is_err());
}
#[tokio::test]
async fn the_keepalive_re_reports_a_silent_session() {
let (_dir, socket, seen) = fake_daemon();
let (tee, lines) = mpsc::channel(4);
let observer = tokio::spawn(observe(lines, Some(socket), Duration::from_millis(20)));
tee.send((
Direction::FromClaude,
r#"{"type":"system","subtype":"init","session_id":"ka-1"}"#.to_string(),
))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(120)).await;
drop(tee);
observer.await.unwrap();
let envelopes = seen.lock().await.clone();
let observes: Vec<_> = envelopes.iter().filter(|e| e["op"] == "observe").collect();
assert!(observes.len() > 1, "expected keep-alives, got {observes:?}");
for envelope in &observes {
assert_eq!(envelope["payload"]["session_id"], "ka-1");
assert_eq!(envelope["payload"]["event"]["stream_state"], "idle");
}
assert_eq!(envelopes.last().unwrap()["op"], "end");
}
#[tokio::test]
async fn the_pump_stops_when_the_far_side_goes_away() {
struct Closed;
impl AsyncWrite for Closed {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
_buf: &[u8],
) -> Poll<std::io::Result<usize>> {
Poll::Ready(Err(std::io::Error::from(std::io::ErrorKind::BrokenPipe)))
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(
self: Pin<&mut Self>,
_cx: &mut TaskContext<'_>,
) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
}
let (tee, mut lines) = mpsc::channel(4);
pump(
&b"{\"type\":\"result\"}\n"[..],
Closed,
Direction::FromClaude,
tee,
)
.await;
assert!(lines.try_recv().is_err());
}
#[test]
fn a_full_tee_drops_lines_rather_than_blocking() {
let (tee, _lines) = mpsc::channel(1);
let mut buffer = Vec::new();
tee_chunk(b"a\nb\nc\n", Direction::ToClaude, &tee, &mut buffer);
assert_eq!(tee.capacity(), 0);
}
}