1use std::io::IsTerminal;
26use std::os::unix::process::{CommandExt, ExitStatusExt};
27use std::path::{Path, PathBuf};
28use std::process::{ExitStatus, Stdio};
29use std::time::Duration;
30
31use anyhow::{anyhow, bail, Context, Result};
32use clap::Parser;
33use serde_json::{json, Value};
34use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
35use tokio::sync::{mpsc, watch};
36
37use crate::claude::model_config::get_model_registry;
38use crate::daemon::client::DaemonClient;
39use crate::daemon::protocol::DaemonEnvelope;
40use crate::daemon::server;
41use crate::sessions::stream::{Direction, StreamTracker};
42
43const SERVICE: &str = "sessions";
45
46mod diagnostics;
47use diagnostics::{increment, Diagnostics};
48
49const REPORT_TIMEOUT: Duration = Duration::from_secs(2);
52
53const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
61
62const TEE_CAPACITY: usize = 256;
66
67const MAX_LINE_BYTES: usize = 1024 * 1024;
71
72const PUMP_BUF_BYTES: usize = 64 * 1024;
74
75const SIGNAL_EXIT_BASE: i32 = 128;
78
79const MAX_OSC_SEQUENCE_BYTES: usize = 8 * 1024;
85
86const NO_TITLE_REWRITE_ENV: &str = "OMNI_DEV_CLAUDE_WRAP_NO_TITLE_REWRITE";
91
92#[derive(Parser)]
94pub struct ClaudeWrapCommand {
95 #[arg(long, value_name = "PATH")]
97 pub socket: Option<PathBuf>,
98
99 #[arg(
101 trailing_var_arg = true,
102 allow_hyphen_values = true,
103 value_name = "CMD"
104 )]
105 pub argv: Vec<String>,
106}
107
108impl ClaudeWrapCommand {
109 pub async fn execute(self) -> Result<()> {
118 let code = self.run().await?;
119 std::process::exit(code)
120 }
121
122 async fn run(self) -> Result<i32> {
126 let Some((program, args)) = self.argv.split_first() else {
127 bail!(
128 "claude-wrap needs a command to run, e.g. \
129 `omni-dev claude-wrap -- claude --output-format stream-json`"
130 );
131 };
132
133 if std::io::stdout().is_terminal() {
138 return Err(exec_replace(program, args));
139 }
140
141 wrap(program, args, self.socket).await
142 }
143}
144
145fn exec_replace(program: &str, args: &[String]) -> anyhow::Error {
148 let error = std::process::Command::new(program).args(args).exec();
149 anyhow::Error::new(error).context(format!("failed to exec {program}"))
150}
151
152async fn wrap(program: &str, args: &[String], socket: Option<PathBuf>) -> Result<i32> {
154 wrap_io_diagnostics(
156 program,
157 args,
158 socket,
159 tokio::io::stdin(),
160 tokio::io::stdout(),
161 std::env::var_os("OMNI_DEV_CLAUDE_WRAP_LOG").map(PathBuf::from),
162 )
163 .await
164 }
166
167#[cfg(test)]
175async fn wrap_io<R, W>(
176 program: &str,
177 args: &[String],
178 socket: Option<PathBuf>,
179 input: R,
180 output: W,
181) -> Result<i32>
182where
183 R: AsyncRead + Unpin + Send + 'static,
184 W: AsyncWrite + Unpin + Send + 'static,
185{
186 wrap_io_diagnostics(program, args, socket, input, output, None).await
187}
188
189async fn wrap_io_diagnostics<R, W>(
190 program: &str,
191 args: &[String],
192 socket: Option<PathBuf>,
193 input: R,
194 output: W,
195 log_path: Option<PathBuf>,
196) -> Result<i32>
197where
198 R: AsyncRead + Unpin + Send + 'static,
199 W: AsyncWrite + Unpin + Send + 'static,
200{
201 let (diagnostics, log_done) = Diagnostics::open(log_path.as_deref());
202 diagnostics.record(|| json!({"event":"process_start"}));
203 let mut child = tokio::process::Command::new(program)
209 .args(args)
210 .stdin(Stdio::piped())
211 .stdout(Stdio::piped())
212 .stderr(Stdio::inherit())
213 .spawn()
214 .with_context(|| format!("failed to spawn {program}"))?;
215
216 let child_stdin = child
217 .stdin
218 .take()
219 .ok_or_else(|| anyhow!("failed to capture the wrapped process's stdin"))?;
220 let child_stdout = child
221 .stdout
222 .take()
223 .ok_or_else(|| anyhow!("failed to capture the wrapped process's stdout"))?;
224 let child_pid = child.id();
225
226 let (tee, lines) = mpsc::channel::<(Direction, String)>(TEE_CAPACITY);
227 let (title_tx, title_rx) = watch::channel::<Option<String>>(None);
228 let observer = tokio::spawn(observe_diagnostics(
229 lines,
230 socket,
231 KEEPALIVE_INTERVAL,
232 title_tx,
233 child_pid,
234 diagnostics.clone(),
235 ));
236 let signals = tokio::spawn(forward_signals(child_pid));
237
238 let from_child_title_rx = title_rewrite_enabled().then_some(title_rx);
241 let from_child = tokio::spawn(pump_diagnostics(
242 child_stdout,
243 output,
244 Direction::FromClaude,
245 tee.clone(),
246 from_child_title_rx,
247 diagnostics.clone(),
248 ));
249 let to_child = tokio::spawn(pump_diagnostics(
250 input,
251 child_stdin,
252 Direction::ToClaude,
253 tee.clone(),
254 None,
255 diagnostics.clone(),
256 ));
257 drop(tee);
258
259 let _ = from_child.await;
263 to_child.abort();
267 let _ = to_child.await;
268 signals.abort();
269
270 let status = child
271 .wait()
272 .await
273 .context("failed to wait for the wrapped process")?;
274 let _ = observer.await;
275 let code = exit_code(status);
276 diagnostics.record(|| json!({"event":"process_exit", "pid":child_pid, "code":code}));
277 drop(diagnostics);
278 if let Some(done) = log_done {
279 let _ = tokio::time::timeout(Duration::from_millis(200), done).await;
280 }
281 Ok(code)
282}
283
284#[cfg(test)]
313async fn pump<R, W>(
314 reader: R,
315 writer: W,
316 direction: Direction,
317 tee: mpsc::Sender<(Direction, String)>,
318 title_rx: Option<watch::Receiver<Option<String>>>,
319) where
320 R: AsyncRead + Unpin,
321 W: AsyncWrite + Unpin,
322{
323 pump_diagnostics(
324 reader,
325 writer,
326 direction,
327 tee,
328 title_rx,
329 Diagnostics::default(),
330 )
331 .await;
332}
333
334async fn pump_diagnostics<R, W>(
335 mut reader: R,
336 mut writer: W,
337 direction: Direction,
338 tee: mpsc::Sender<(Direction, String)>,
339 mut title_rx: Option<watch::Receiver<Option<String>>>,
340 diagnostics: Diagnostics,
341) where
342 R: AsyncRead + Unpin,
343 W: AsyncWrite + Unpin,
344{
345 let mut buffer = vec![0u8; PUMP_BUF_BYTES];
346 let mut line = Vec::new();
347 let mut rewriter = TitleRewriter::default();
348 loop {
349 let title_changed = async {
350 match title_rx.as_mut() {
351 Some(rx) => rx.changed().await,
352 None => std::future::pending().await,
355 }
356 };
357 tokio::select! {
358 read_result = reader.read(&mut buffer) => {
359 let read = match read_result {
360 Ok(0) | Err(_) => break,
361 Ok(read) => read,
362 };
363 let chunk = &buffer[..read];
364 let write_result = match &title_rx {
365 Some(rx) => {
366 let prefix = rx.borrow().clone();
367 writer
368 .write_all(&rewriter.rewrite(chunk, prefix.as_deref()))
369 .await
370 }
371 None => writer.write_all(chunk).await,
372 };
373 if write_result.is_err() || writer.flush().await.is_err() {
374 break;
375 }
376 tee_chunk_diagnostics(chunk, direction, &tee, &mut line, &diagnostics);
377 }
378 changed = title_changed => {
379 match changed {
380 Ok(()) => {
381 let Some(rx) = &title_rx else { continue };
382 let prefix = rx.borrow().clone();
383 let Some(bytes) = rewriter.resync(prefix.as_deref()) else { continue };
384 if writer.write_all(&bytes).await.is_err()
385 || writer.flush().await.is_err()
386 {
387 break;
388 }
389 }
390 Err(_) => title_rx = None,
394 }
395 }
396 }
397 }
398 let _ = writer.shutdown().await;
401}
402
403#[derive(Debug, Default)]
433struct TitleRewriter {
434 pending: Vec<u8>,
438 title: Vec<u8>,
441 code: u8,
443 phase: Phase,
444 last_code: Option<u8>,
448 last_title: Option<Vec<u8>>,
449 last_terminator: Option<Vec<u8>>,
450}
451
452#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
455enum Phase {
456 #[default]
458 Idle,
459 Esc,
461 Open,
463 Code,
465 Title,
467 TitleEsc,
470}
471
472const ESC: u8 = 0x1B;
473const BEL: u8 = 0x07;
474
475impl TitleRewriter {
476 fn rewrite(&mut self, chunk: &[u8], prefix: Option<&str>) -> Vec<u8> {
481 let mut out = Vec::with_capacity(chunk.len());
482 for &byte in chunk {
483 match self.phase {
484 Phase::Idle => {
485 if byte == ESC {
486 self.pending.push(byte);
487 self.phase = Phase::Esc;
488 } else {
489 out.push(byte);
490 }
491 }
492 Phase::Esc => {
493 self.pending.push(byte);
494 if byte == b']' {
495 self.phase = Phase::Open;
496 } else {
497 self.flush_pending(&mut out);
498 }
499 }
500 Phase::Open => {
501 self.pending.push(byte);
502 if byte == b'0' || byte == b'2' {
503 self.code = byte;
504 self.phase = Phase::Code;
505 } else {
506 self.flush_pending(&mut out);
507 }
508 }
509 Phase::Code => {
510 self.pending.push(byte);
511 if byte == b';' {
512 self.phase = Phase::Title;
513 } else {
514 self.flush_pending(&mut out);
515 }
516 }
517 Phase::Title => {
518 self.pending.push(byte);
519 if byte == BEL {
520 self.finish(&mut out, prefix, &[BEL]);
521 } else if byte == ESC {
522 self.phase = Phase::TitleEsc;
523 } else {
524 self.title.push(byte);
525 }
526 }
527 Phase::TitleEsc => {
528 self.pending.push(byte);
529 if byte == b'\\' {
530 self.finish(&mut out, prefix, &[ESC, b'\\']);
531 } else {
532 self.flush_pending(&mut out);
535 }
536 }
537 }
538 if self.pending.len() > MAX_OSC_SEQUENCE_BYTES {
539 self.flush_pending(&mut out);
540 }
541 }
542 out
543 }
544
545 fn finish(&mut self, out: &mut Vec<u8>, prefix: Option<&str>, terminator: &[u8]) {
549 let rewritten = prefix.zip(std::str::from_utf8(&self.title).ok());
550 match rewritten {
551 Some((prefix, title)) => {
552 out.push(ESC);
553 out.push(b']');
554 out.push(self.code);
555 out.push(b';');
556 out.extend_from_slice(prefix.as_bytes());
557 out.extend_from_slice(title.as_bytes());
558 out.extend_from_slice(terminator);
559 }
560 None => out.extend_from_slice(&self.pending),
561 }
562 self.last_code = Some(self.code);
563 self.last_title = Some(self.title.clone());
564 self.last_terminator = Some(terminator.to_vec());
565 self.reset();
566 }
567
568 fn resync(&self, prefix: Option<&str>) -> Option<Vec<u8>> {
576 let prefix = prefix?;
577 let code = self.last_code?;
578 let title = std::str::from_utf8(self.last_title.as_ref()?).ok()?;
579 let terminator = self.last_terminator.as_ref()?;
580 let mut out = vec![ESC, b']', code, b';'];
581 out.extend_from_slice(prefix.as_bytes());
582 out.extend_from_slice(title.as_bytes());
583 out.extend_from_slice(terminator);
584 Some(out)
585 }
586
587 fn flush_pending(&mut self, out: &mut Vec<u8>) {
589 out.extend_from_slice(&self.pending);
590 self.reset();
591 }
592
593 fn reset(&mut self) {
594 self.pending.clear();
595 self.title.clear();
596 self.code = 0;
597 self.phase = Phase::Idle;
598 }
599}
600
601fn title_rewrite_enabled() -> bool {
604 !std::env::var(NO_TITLE_REWRITE_ENV).is_ok_and(|v| flag_is_truthy(&v))
605}
606
607fn flag_is_truthy(value: &str) -> bool {
611 matches!(
612 value.trim().to_ascii_lowercase().as_str(),
613 "1" | "true" | "yes" | "on"
614 )
615}
616
617fn title_prefix_for_model(model_id: &str) -> String {
623 let family = get_model_registry().get_model_family(model_id);
624 format!("{} {} · ", family.glyph(), family.label())
625}
626
627#[cfg(test)]
631fn tee_chunk(
632 chunk: &[u8],
633 direction: Direction,
634 tee: &mpsc::Sender<(Direction, String)>,
635 line: &mut Vec<u8>,
636) {
637 tee_chunk_diagnostics(chunk, direction, tee, line, &Diagnostics::default());
638}
639
640fn tee_chunk_diagnostics(
641 chunk: &[u8],
642 direction: Direction,
643 tee: &mpsc::Sender<(Direction, String)>,
644 line: &mut Vec<u8>,
645 diagnostics: &Diagnostics,
646) {
647 for &byte in chunk {
648 if byte == b'\n' {
649 if line.len() < MAX_LINE_BYTES {
650 if let Ok(text) = std::str::from_utf8(line) {
651 if let Err(error) = tee.try_send((direction, text.to_string())) {
652 if let Some(c) = diagnostics.counts() {
653 match error {
654 mpsc::error::TrySendError::Full(_) => increment(&c.full),
655 mpsc::error::TrySendError::Closed(_) => increment(&c.closed),
656 }
657 }
658 }
659 } else if let Some(c) = diagnostics.counts() {
660 increment(&c.non_utf8);
661 }
662 } else if let Some(c) = diagnostics.counts() {
663 increment(&c.oversize);
664 }
665 line.clear();
666 } else if line.len() < MAX_LINE_BYTES {
667 line.push(byte);
668 }
669 }
670}
671
672#[cfg(test)]
683async fn observe(
684 lines: mpsc::Receiver<(Direction, String)>,
685 socket: Option<PathBuf>,
686 every: Duration,
687 title_tx: watch::Sender<Option<String>>,
688 child_pid: Option<u32>,
689) {
690 observe_diagnostics(
691 lines,
692 socket,
693 every,
694 title_tx,
695 child_pid,
696 Diagnostics::default(),
697 )
698 .await;
699}
700
701async fn observe_diagnostics(
702 mut lines: mpsc::Receiver<(Direction, String)>,
703 socket: Option<PathBuf>,
704 every: Duration,
705 title_tx: watch::Sender<Option<String>>,
706 child_pid: Option<u32>,
707 diagnostics: Diagnostics,
708) {
709 let Ok(socket) = server::resolve_socket(socket) else {
710 diagnostics
712 .record(|| json!({"event":"observer_stopped", "outcome":"socket_resolution_failed"}));
713 return;
714 };
716 let mut tracker = StreamTracker::new();
717 let mut last_model: Option<String> = None;
718 let mut last_identity = None;
719 let mut keepalive = tokio::time::interval(every);
720 keepalive.tick().await;
723
724 loop {
725 let observed = tokio::select! {
726 line = lines.recv() => match line {
727 Some((direction, text)) => tracker.observe_line(direction, &text),
728 None => break,
729 },
730 _ = keepalive.tick() => {
731 diagnostics.summary(tracker.diagnostics());
732 tracker.keepalive()
733 },
734 };
735 if tracker.model() != last_model.as_deref() {
740 last_model = tracker.model().map(str::to_string);
741 let prefix = last_model.as_deref().map(title_prefix_for_model);
742 let _ = title_tx.send(prefix);
743 }
744 if let Some(mut request) = observed {
745 request.pid = child_pid;
746 if diagnostics.enabled() {
747 let identity = (
748 request.session_id.clone(),
749 request.cwd.clone(),
750 request.model.clone(),
751 );
752 if last_identity.as_ref() != Some(&identity) {
753 diagnostics.record(|| {
754 json!({"event":"session_identity", "session_id":identity.0,
755 "cwd":identity.1, "model":identity.2, "pid":child_pid})
756 });
757 last_identity = Some(identity);
758 }
759 }
760 diagnostics.record(|| {
761 json!({"event":"state_report", "session_id":request.session_id,
762 "state":request.event, "pid":child_pid})
763 });
764 if let Ok(payload) = serde_json::to_value(request) {
765 report_diagnostics(&socket, "observe", payload, &diagnostics).await;
766 } else {
767 diagnostics.record(
769 || json!({"event":"report", "op":"observe", "outcome":"serialization_failed"}),
770 );
771 }
773 }
774 }
775
776 if let Some(session_id) = tracker.session_id() {
777 let mut payload = json!({ "session_id": session_id });
778 if let Some(pid) = child_pid {
779 payload["pid"] = Value::from(pid);
780 }
781 report_diagnostics(&socket, "end", payload, &diagnostics).await;
782 }
783 diagnostics.summary(tracker.diagnostics());
784}
785
786async fn report_diagnostics(socket: &Path, op: &str, payload: Value, diagnostics: &Diagnostics) {
789 let envelope = DaemonEnvelope::service(SERVICE, op, payload);
790 let outcome =
791 match tokio::time::timeout(REPORT_TIMEOUT, DaemonClient::new(socket).request(envelope))
792 .await
793 {
794 Err(_) => "timeout",
795 Ok(Err(_)) => "transport_failed",
796 Ok(Ok(reply)) if !reply.ok => "daemon_rejected",
797 Ok(Ok(_)) => "delivered",
798 };
799 diagnostics.record(|| json!({"event":"report", "op":op, "outcome":outcome}));
800}
801
802async fn forward_signals(child_pid: Option<u32>) {
805 let Some(pid) = child_pid.and_then(|pid| i32::try_from(pid).ok()) else {
806 return;
807 };
808 let pid = nix::unistd::Pid::from_raw(pid);
809 let (Ok(mut terminate), Ok(mut interrupt)) = (
810 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()),
811 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()),
812 ) else {
813 return;
814 };
815 loop {
816 let signal = tokio::select! {
817 _ = terminate.recv() => nix::sys::signal::Signal::SIGTERM,
818 _ = interrupt.recv() => nix::sys::signal::Signal::SIGINT,
819 };
820 let _ = nix::sys::signal::kill(pid, signal);
821 }
822}
823
824fn exit_code(status: ExitStatus) -> i32 {
827 status
828 .code()
829 .or_else(|| status.signal().map(|signal| SIGNAL_EXIT_BASE + signal))
830 .unwrap_or(1)
831}
832
833#[cfg(test)]
834#[allow(clippy::unwrap_used, clippy::expect_used)]
835mod tests {
836 use std::pin::Pin;
837 use std::sync::{Arc, Mutex};
838 use std::task::{Context as TaskContext, Poll};
839
840 use tokio::io::AsyncBufReadExt;
841 use tokio::net::UnixListener;
842
843 use super::*;
844
845 #[tokio::test]
846 async fn diagnostic_tee_counts_each_drop_reason() {
847 use std::sync::atomic::Ordering;
848 let dir = tempfile::tempdir().unwrap();
849 let (diagnostics, done) = Diagnostics::open(Some(&dir.path().join("log")));
850 let (tee, lines) = mpsc::channel(1);
851 let mut buffer = Vec::new();
852 tee_chunk_diagnostics(
853 b"a\nb\n",
854 Direction::FromClaude,
855 &tee,
856 &mut buffer,
857 &diagnostics,
858 );
859 drop(lines);
860 tee_chunk_diagnostics(
861 b"c\n",
862 Direction::FromClaude,
863 &tee,
864 &mut buffer,
865 &diagnostics,
866 );
867 tee_chunk_diagnostics(
868 b"\xff\n",
869 Direction::FromClaude,
870 &tee,
871 &mut buffer,
872 &diagnostics,
873 );
874 tee_chunk_diagnostics(
875 &vec![b'x'; MAX_LINE_BYTES],
876 Direction::FromClaude,
877 &tee,
878 &mut buffer,
879 &diagnostics,
880 );
881 tee_chunk_diagnostics(
882 b"\n",
883 Direction::FromClaude,
884 &tee,
885 &mut buffer,
886 &diagnostics,
887 );
888 let counts = diagnostics.counts().unwrap();
889 assert_eq!(counts.full.load(Ordering::Relaxed), 1);
890 assert_eq!(counts.closed.load(Ordering::Relaxed), 1);
891 assert_eq!(counts.oversize.load(Ordering::Relaxed), 1);
892 assert_eq!(counts.non_utf8.load(Ordering::Relaxed), 1);
893 drop(diagnostics);
894 done.unwrap().await.unwrap();
895 }
896
897 #[tokio::test]
898 async fn enabled_diagnostics_preserve_bytes_and_exclude_conversation_content() {
899 let (dir, socket, _seen) = fake_daemon();
900 let log = dir.path().join("wrapper.jsonl");
901 let output = Sink::default();
902 let stream = concat!(
903 r#"{"type":"system","subtype":"init","session_id":"private-test","cwd":"/project","model":"test"}"#,
904 "\n",
905 r#"{"type":"assistant","message":{"content":"CONVERSATION_SECRET"}}"#,
906 "\n",
907 r#"{"type":"control_request","request_id":"perm","request":{"subtype":"can_use_tool","input":"TOOL_SECRET"}}"#,
908 "\n",
909 r#"{"type":"result"}"#,
910 "\n"
911 );
912 let code = wrap_io_diagnostics(
913 "/bin/cat",
914 &[],
915 Some(socket),
916 stream.as_bytes(),
917 output.clone(),
918 Some(log.clone()),
919 )
920 .await
921 .unwrap();
922 assert_eq!(code, 0);
923 assert_eq!(output.contents(), stream);
924 let text = std::fs::read_to_string(log).unwrap();
925 assert!(text.contains("session_identity"));
926 assert!(text.contains("process_exit"));
927 assert!(text.contains("state_report"));
928 assert!(!text.contains("CONVERSATION_SECRET"));
929 assert!(!text.contains("TOOL_SECRET"));
930 for line in text.lines() {
931 serde_json::from_str::<Value>(line).unwrap();
932 }
933 }
934
935 #[tokio::test]
936 async fn diagnostic_report_rejection_does_not_persist_daemon_error_content() {
937 let dir = tempfile::tempdir().unwrap();
938 let socket = dir.path().join("d.sock");
939 let listener = UnixListener::bind(&socket).unwrap();
940 let server = tokio::spawn(async move {
941 let (stream, _) = listener.accept().await.unwrap();
942 let mut framed =
943 tokio_util::codec::Framed::new(stream, tokio_util::codec::LinesCodec::new());
944 use futures::{SinkExt, StreamExt};
945 framed.next().await.unwrap().unwrap();
946 framed
947 .send(r#"{"ok":false,"error":"DAEMON_CONTENT_SECRET"}"#.to_string())
948 .await
949 .unwrap();
950 });
951 let path = dir.path().join("log");
952 let (diagnostics, done) = Diagnostics::open(Some(&path));
953 report_diagnostics(&socket, "observe", json!({}), &diagnostics).await;
954 server.await.unwrap();
955 report_diagnostics(
956 &dir.path().join("missing.sock"),
957 "observe",
958 json!({}),
959 &diagnostics,
960 )
961 .await;
962 drop(diagnostics);
963 done.unwrap().await.unwrap();
964 let text = std::fs::read_to_string(path).unwrap();
965 assert!(text.contains("daemon_rejected"));
966 assert!(text.contains("transport_failed"));
967 assert!(!text.contains("DAEMON_CONTENT_SECRET"));
968 }
969
970 #[tokio::test(start_paused = true)]
971 async fn diagnostic_report_gives_up_on_a_wedged_daemon() {
972 let dir = tempfile::tempdir_in("/tmp").unwrap();
973 let socket = dir.path().join("d.sock");
974 let _wedged = UnixListener::bind(&socket).unwrap();
978 let path = dir.path().join("log");
979 let (diagnostics, done) = Diagnostics::open(Some(&path));
980 report_diagnostics(&socket, "observe", json!({}), &diagnostics).await;
981 drop(diagnostics);
982 done.unwrap().await.unwrap();
983 assert!(std::fs::read_to_string(path).unwrap().contains("timeout"));
984 }
985
986 #[derive(Clone, Default)]
989 struct Sink(Arc<Mutex<Vec<u8>>>);
990
991 impl Sink {
992 fn contents(&self) -> String {
993 String::from_utf8(self.0.lock().unwrap().clone()).unwrap()
994 }
995 }
996
997 impl AsyncWrite for Sink {
998 fn poll_write(
999 self: Pin<&mut Self>,
1000 _cx: &mut TaskContext<'_>,
1001 buf: &[u8],
1002 ) -> Poll<std::io::Result<usize>> {
1003 self.0.lock().unwrap().extend_from_slice(buf);
1004 Poll::Ready(Ok(buf.len()))
1005 }
1006
1007 fn poll_flush(
1008 self: Pin<&mut Self>,
1009 _cx: &mut TaskContext<'_>,
1010 ) -> Poll<std::io::Result<()>> {
1011 Poll::Ready(Ok(()))
1012 }
1013
1014 fn poll_shutdown(
1015 self: Pin<&mut Self>,
1016 _cx: &mut TaskContext<'_>,
1017 ) -> Poll<std::io::Result<()>> {
1018 Poll::Ready(Ok(()))
1019 }
1020 }
1021
1022 async fn wrap_script(script: &str, socket: Option<PathBuf>) -> (i32, String) {
1025 let sink = Sink::default();
1026 let code = wrap_io(
1027 "/bin/sh",
1028 &["-c".to_string(), script.to_string()],
1029 socket,
1030 tokio::io::empty(),
1031 sink.clone(),
1032 )
1033 .await
1034 .unwrap();
1035 (code, sink.contents())
1036 }
1037
1038 fn parse(args: &[&str]) -> ClaudeWrapCommand {
1040 let mut argv = vec!["claude-wrap"];
1041 argv.extend_from_slice(args);
1042 ClaudeWrapCommand::try_parse_from(argv).unwrap()
1043 }
1044
1045 fn fake_daemon() -> (
1047 tempfile::TempDir,
1048 PathBuf,
1049 Arc<tokio::sync::Mutex<Vec<Value>>>,
1050 ) {
1051 let dir = tempfile::tempdir_in("/tmp").unwrap();
1052 let socket = dir.path().join("d.sock");
1053 let listener = UnixListener::bind(&socket).unwrap();
1054 let seen = Arc::new(tokio::sync::Mutex::new(Vec::new()));
1055 let recorder = Arc::clone(&seen);
1056 tokio::spawn(async move {
1057 loop {
1058 let Ok((stream, _)) = listener.accept().await else {
1059 return;
1060 };
1061 let recorder = Arc::clone(&recorder);
1062 tokio::spawn(async move {
1063 let (read, mut write) = stream.into_split();
1064 let mut line = String::new();
1065 if tokio::io::BufReader::new(read)
1066 .read_line(&mut line)
1067 .await
1068 .is_ok()
1069 {
1070 if let Ok(value) = serde_json::from_str::<Value>(&line) {
1071 recorder.lock().await.push(value);
1072 }
1073 }
1074 let _ = write.write_all(b"{\"ok\":true,\"payload\":{}}\n").await;
1075 });
1076 }
1077 });
1078 (dir, socket, seen)
1079 }
1080
1081 #[test]
1082 fn trailing_args_are_captured_verbatim_after_a_separator() {
1083 let cmd = parse(&["--", "node", "cli.js", "--output-format", "stream-json"]);
1084 assert_eq!(
1085 cmd.argv,
1086 vec!["node", "cli.js", "--output-format", "stream-json"]
1087 );
1088 assert!(cmd.socket.is_none());
1089 }
1090
1091 #[test]
1092 fn the_socket_override_is_parsed_before_the_separator() {
1093 let cmd = parse(&["--socket", "/tmp/d.sock", "--", "claude", "-p"]);
1094 assert_eq!(cmd.socket.unwrap(), PathBuf::from("/tmp/d.sock"));
1095 assert_eq!(cmd.argv, vec!["claude", "-p"]);
1096 }
1097
1098 #[tokio::test]
1099 async fn an_empty_command_is_a_usage_error() {
1100 let error = parse(&[]).run().await.unwrap_err();
1101 assert!(error.to_string().contains("needs a command to run"));
1102 }
1103
1104 #[tokio::test]
1118 async fn signal_forwarding_gives_up_when_there_is_no_child() {
1119 forward_signals(None).await;
1122 }
1123
1124 #[tokio::test]
1125 async fn the_child_is_wrapped_losslessly_and_its_exit_code_survives() {
1126 let (_dir, socket, _seen) = fake_daemon();
1127 let (code, forwarded) =
1130 wrap_script("printf 'hello\\nnot json\\ntrailing'; exit 7", Some(socket)).await;
1131 assert_eq!(code, 7);
1132 assert_eq!(forwarded, "hello\nnot json\ntrailing");
1133 }
1134
1135 #[tokio::test]
1136 async fn stream_state_transitions_reach_the_daemon() {
1137 let (_dir, socket, seen) = fake_daemon();
1138 let script = concat!(
1140 r#"printf '{"type":"system","subtype":"init","session_id":"s-1","cwd":"/w"}\n';"#,
1141 r#"printf '{"type":"assistant","session_id":"s-1"}\n';"#,
1142 r#"printf '{"type":"control_request","request_id":"r1","request":{"subtype":"can_use_tool"}}\n';"#,
1143 r#"printf '{"type":"result"}\n'"#,
1144 );
1145 let (code, _forwarded) = wrap_script(script, Some(socket)).await;
1146 assert_eq!(code, 0);
1147
1148 let envelopes = seen.lock().await.clone();
1149 let states: Vec<String> = envelopes
1150 .iter()
1151 .filter(|e| e["op"] == "observe")
1152 .filter_map(|e| e["payload"]["event"]["stream_state"].as_str())
1153 .map(ToString::to_string)
1154 .collect();
1155 assert_eq!(
1156 states,
1157 vec!["idle", "working", "waiting_for_permission", "idle"]
1158 );
1159 assert_eq!(envelopes[0]["service"], "sessions");
1161 assert_eq!(envelopes[0]["payload"]["session_id"], "s-1");
1162 assert_eq!(envelopes[0]["payload"]["cwd"], "/w");
1163 let ended = envelopes.last().unwrap();
1164 assert_eq!(ended["op"], "end");
1165 assert_eq!(ended["payload"]["session_id"], "s-1");
1166 }
1167
1168 #[tokio::test]
1169 async fn a_missing_daemon_never_affects_the_child() {
1170 let dir = tempfile::tempdir_in("/tmp").unwrap();
1172 let script = concat!(
1173 r#"printf '{"type":"system","subtype":"init","session_id":"s-1"}\n';"#,
1174 r#"printf 'not json at all\n';"#,
1175 r#"exit 3"#,
1176 );
1177 let (code, forwarded) = wrap_script(script, Some(dir.path().join("absent.sock"))).await;
1178 assert_eq!(code, 3);
1179 assert!(forwarded.ends_with("not json at all\n"));
1180 }
1181
1182 #[tokio::test]
1183 async fn a_child_killed_by_a_signal_reports_the_shell_convention() {
1184 let (_dir, socket, _seen) = fake_daemon();
1185 let (code, _forwarded) = wrap_script("kill -TERM $$", Some(socket)).await;
1186 assert_eq!(code, SIGNAL_EXIT_BASE + 15);
1187 }
1188
1189 #[tokio::test]
1190 async fn spawning_a_missing_program_is_an_error_not_a_panic() {
1191 let error = wrap_io(
1192 "/nonexistent/omni-dev-test-binary",
1193 &[],
1194 None,
1195 tokio::io::empty(),
1196 Sink::default(),
1197 )
1198 .await
1199 .unwrap_err();
1200 assert!(error.to_string().contains("failed to spawn"));
1201 }
1202
1203 #[test]
1204 fn over_long_and_non_utf8_lines_are_dropped_but_still_forwarded() {
1205 let (tee, mut lines) = mpsc::channel(4);
1206 let mut buffer = Vec::new();
1207 let mut chunk = vec![b'x'; MAX_LINE_BYTES + 1];
1208 chunk.push(b'\n');
1209 chunk.extend_from_slice(&[0xff, 0xfe, b'\n']);
1210 chunk.extend_from_slice(b"{\"type\":\"result\"}\n");
1211 tee_chunk(&chunk, Direction::FromClaude, &tee, &mut buffer);
1212 let (direction, text) = lines.try_recv().unwrap();
1214 assert_eq!(direction, Direction::FromClaude);
1215 assert_eq!(text, r#"{"type":"result"}"#);
1216 assert!(lines.try_recv().is_err());
1217 }
1218
1219 #[tokio::test]
1220 async fn the_keepalive_re_reports_a_silent_working_session() {
1221 let (_dir, socket, seen) = fake_daemon();
1224 let (tee, lines) = mpsc::channel(4);
1225 let (title_tx, _title_rx) = watch::channel(None);
1226 let observer = tokio::spawn(observe(
1227 lines,
1228 Some(socket),
1229 Duration::from_millis(20),
1230 title_tx,
1231 Some(4242),
1232 ));
1233 for line in [
1234 r#"{"type":"system","subtype":"init","session_id":"ka-1"}"#,
1235 r#"{"type":"assistant","session_id":"ka-1"}"#,
1236 ] {
1237 tee.send((Direction::FromClaude, line.to_string()))
1238 .await
1239 .unwrap();
1240 }
1241 tokio::time::sleep(Duration::from_millis(120)).await;
1242 drop(tee);
1243 observer.await.unwrap();
1244
1245 let envelopes = seen.lock().await.clone();
1246 let observes: Vec<_> = envelopes.iter().filter(|e| e["op"] == "observe").collect();
1247 assert!(observes.len() > 2, "expected keep-alives, got {observes:?}");
1249 for envelope in &observes[1..] {
1250 assert_eq!(envelope["payload"]["session_id"], "ka-1");
1251 assert_eq!(envelope["payload"]["event"]["stream_state"], "working");
1252 assert_eq!(envelope["payload"]["pid"], 4242);
1254 }
1255 let end = envelopes.last().unwrap();
1256 assert_eq!(end["op"], "end");
1257 assert_eq!(end["payload"]["pid"], 4242);
1258 }
1259
1260 #[tokio::test]
1261 async fn the_keepalive_leaves_an_idle_session_to_age_out() {
1262 let (_dir, socket, seen) = fake_daemon();
1265 let (tee, lines) = mpsc::channel(4);
1266 let (title_tx, _title_rx) = watch::channel(None);
1267 let observer = tokio::spawn(observe(
1268 lines,
1269 Some(socket),
1270 Duration::from_millis(20),
1271 title_tx,
1272 Some(4242),
1273 ));
1274 tee.send((
1275 Direction::FromClaude,
1276 r#"{"type":"system","subtype":"init","session_id":"ka-2"}"#.to_string(),
1277 ))
1278 .await
1279 .unwrap();
1280 tokio::time::sleep(Duration::from_millis(120)).await;
1281 drop(tee);
1282 observer.await.unwrap();
1283
1284 let envelopes = seen.lock().await.clone();
1285 let ops: Vec<_> = envelopes
1286 .iter()
1287 .map(|e| e["op"].as_str().unwrap())
1288 .collect();
1289 assert_eq!(ops, ["observe", "end"], "got {envelopes:?}");
1290 assert_eq!(envelopes[0]["payload"]["event"]["stream_state"], "idle");
1291 }
1292
1293 #[tokio::test]
1294 async fn the_pump_stops_when_the_far_side_goes_away() {
1295 struct Closed;
1298 impl AsyncWrite for Closed {
1299 fn poll_write(
1300 self: Pin<&mut Self>,
1301 _cx: &mut TaskContext<'_>,
1302 _buf: &[u8],
1303 ) -> Poll<std::io::Result<usize>> {
1304 Poll::Ready(Err(std::io::Error::from(std::io::ErrorKind::BrokenPipe)))
1305 }
1306 fn poll_flush(
1307 self: Pin<&mut Self>,
1308 _cx: &mut TaskContext<'_>,
1309 ) -> Poll<std::io::Result<()>> {
1310 Poll::Ready(Ok(()))
1311 }
1312 fn poll_shutdown(
1313 self: Pin<&mut Self>,
1314 _cx: &mut TaskContext<'_>,
1315 ) -> Poll<std::io::Result<()>> {
1316 Poll::Ready(Ok(()))
1317 }
1318 }
1319 let (tee, mut lines) = mpsc::channel(4);
1320 pump(
1322 &b"{\"type\":\"result\"}\n"[..],
1323 Closed,
1324 Direction::FromClaude,
1325 tee,
1326 None,
1327 )
1328 .await;
1329 assert!(lines.try_recv().is_err());
1330 }
1331
1332 #[tokio::test]
1333 async fn pump_resyncs_the_title_from_the_watch_channel_alone_with_no_further_reads() {
1334 let (tee, _lines) = mpsc::channel(4);
1339 let (title_tx, title_rx) = watch::channel(None);
1340 let sink = Sink::default();
1341 let (mut writer_end, reader_end) = tokio::io::duplex(1024);
1342 writer_end.write_all(b"\x1b]0;2.1.132\x07").await.unwrap();
1343
1344 let pump_task = tokio::spawn(pump(
1345 reader_end,
1346 sink.clone(),
1347 Direction::FromClaude,
1348 tee,
1349 Some(title_rx),
1350 ));
1351
1352 tokio::time::sleep(Duration::from_millis(20)).await;
1353 assert_eq!(sink.contents(), "\x1b]0;2.1.132\x07");
1354
1355 title_tx.send(Some(PREFIX.to_string())).unwrap();
1356 tokio::time::sleep(Duration::from_millis(20)).await;
1357 assert!(
1358 sink.contents().contains("\x1b]0;🟡 Opus · 2.1.132\x07"),
1359 "expected a resynced title, got: {:?}",
1360 sink.contents()
1361 );
1362
1363 drop(writer_end);
1364 let _ = pump_task.await;
1365 }
1366
1367 #[test]
1368 fn a_full_tee_drops_lines_rather_than_blocking() {
1369 let (tee, _lines) = mpsc::channel(1);
1370 let mut buffer = Vec::new();
1371 tee_chunk(b"a\nb\nc\n", Direction::ToClaude, &tee, &mut buffer);
1372 assert_eq!(tee.capacity(), 0);
1374 }
1375
1376 const PREFIX: &str = "🟡 Opus · ";
1377
1378 #[test]
1379 fn bel_terminated_title_is_rewritten_with_the_model_prefix() {
1380 let mut rewriter = TitleRewriter::default();
1381 let out = rewriter.rewrite(b"\x1b]0;2.1.132\x07", Some(PREFIX));
1382 assert_eq!(out, b"\x1b]0;\xf0\x9f\x9f\xa1 Opus \xc2\xb7 2.1.132\x07");
1383 assert_eq!(
1384 String::from_utf8(out).unwrap(),
1385 "\x1b]0;🟡 Opus · 2.1.132\x07"
1386 );
1387 }
1388
1389 #[test]
1390 fn st_terminated_title_is_rewritten_with_the_model_prefix() {
1391 let mut rewriter = TitleRewriter::default();
1392 let out = rewriter.rewrite(b"\x1b]2;2.1.132\x1b\\", Some(PREFIX));
1393 assert_eq!(
1394 String::from_utf8(out).unwrap(),
1395 "\x1b]2;🟡 Opus · 2.1.132\x1b\\"
1396 );
1397 }
1398
1399 #[test]
1400 fn no_model_known_yet_leaves_the_title_unchanged() {
1401 let mut rewriter = TitleRewriter::default();
1402 let chunk = b"\x1b]0;2.1.132\x07";
1403 let out = rewriter.rewrite(chunk, None);
1404 assert_eq!(out, chunk);
1405 }
1406
1407 #[test]
1408 fn a_non_osc_chunk_is_untouched() {
1409 let mut rewriter = TitleRewriter::default();
1410 let chunk = b"just some ordinary assistant output, no escapes here\n";
1411 let out = rewriter.rewrite(chunk, Some(PREFIX));
1412 assert_eq!(out, chunk);
1413 }
1414
1415 #[test]
1416 fn a_sequence_split_across_two_reads_is_still_rewritten() {
1417 let mut rewriter = TitleRewriter::default();
1418 let mut out = rewriter.rewrite(b"before \x1b]0;2.1", Some(PREFIX));
1419 out.extend(rewriter.rewrite(b".132\x07 after", Some(PREFIX)));
1420 assert_eq!(
1421 String::from_utf8(out).unwrap(),
1422 "before \x1b]0;🟡 Opus · 2.1.132\x07 after"
1423 );
1424 }
1425
1426 #[test]
1427 fn an_osc_code_other_than_0_or_2_is_left_untouched() {
1428 let mut rewriter = TitleRewriter::default();
1430 let chunk = b"\x1b]8;;https://example.invalid\x07link text\x1b]8;;\x07";
1431 let out = rewriter.rewrite(chunk, Some(PREFIX));
1432 assert_eq!(out, chunk);
1433 }
1434
1435 #[test]
1436 fn an_overlong_sequence_without_a_terminator_is_flushed_unchanged() {
1437 let mut rewriter = TitleRewriter::default();
1438 let mut chunk = b"\x1b]0;".to_vec();
1439 chunk.extend(std::iter::repeat_n(b'x', MAX_OSC_SEQUENCE_BYTES + 16));
1440 let out = rewriter.rewrite(&chunk, Some(PREFIX));
1441 assert_eq!(out, chunk);
1442 }
1443
1444 #[test]
1445 fn a_malformed_st_terminator_is_flushed_unchanged() {
1446 let mut rewriter = TitleRewriter::default();
1449 let chunk = b"\x1b]0;oops\x1bnope";
1450 let out = rewriter.rewrite(chunk, Some(PREFIX));
1451 assert_eq!(out, chunk);
1452 }
1453
1454 #[test]
1455 fn resync_returns_none_before_anything_has_completed() {
1456 let rewriter = TitleRewriter::default();
1457 assert_eq!(rewriter.resync(Some(PREFIX)), None);
1458 }
1459
1460 #[test]
1461 fn resync_returns_none_without_a_prefix() {
1462 let mut rewriter = TitleRewriter::default();
1463 rewriter.rewrite(b"\x1b]0;2.1.132\x07", None);
1464 assert_eq!(rewriter.resync(None), None);
1465 }
1466
1467 #[test]
1468 fn resync_reformats_a_title_that_was_forwarded_before_the_model_was_known() {
1469 let mut rewriter = TitleRewriter::default();
1474 let forwarded = rewriter.rewrite(b"\x1b]0;2.1.132\x07", None);
1475 assert_eq!(forwarded, b"\x1b]0;2.1.132\x07");
1476 let corrected = rewriter.resync(Some(PREFIX)).unwrap();
1477 assert_eq!(
1478 String::from_utf8(corrected).unwrap(),
1479 "\x1b]0;🟡 Opus · 2.1.132\x07"
1480 );
1481 }
1482
1483 #[test]
1484 fn resync_reformats_under_a_changed_prefix_on_a_mid_session_model_switch() {
1485 let mut rewriter = TitleRewriter::default();
1486 rewriter.rewrite(b"\x1b]0;2.1.132\x07", Some(PREFIX));
1487 let resynced = rewriter.resync(Some("🟢 Sonnet · ")).unwrap();
1488 assert_eq!(
1489 String::from_utf8(resynced).unwrap(),
1490 "\x1b]0;🟢 Sonnet · 2.1.132\x07"
1491 );
1492 }
1493
1494 #[test]
1495 fn flag_is_truthy_recognises_on_and_off_values() {
1496 for value in ["1", "true", "TRUE", "yes", "on", " on "] {
1497 assert!(flag_is_truthy(value), "{value:?} should be truthy");
1498 }
1499 for value in ["0", "false", "no", "off", ""] {
1500 assert!(!flag_is_truthy(value), "{value:?} should not be truthy");
1501 }
1502 }
1503
1504 #[tokio::test]
1505 async fn the_wrapper_rewrites_claudes_title_end_to_end() {
1506 let (_dir, socket, _seen) = fake_daemon();
1507 let script = concat!(
1517 r#"printf '\033]0;2.1.132\007';"#,
1518 r#"printf '{"type":"system","subtype":"init","session_id":"s-1","model":"claude-opus-4-8"}\n';"#,
1519 r#"sleep 0.2;"#,
1520 r#"printf 'plain text after the title\n'"#,
1521 );
1522 let (_code, forwarded) = wrap_script(script, Some(socket)).await;
1523 assert!(
1527 forwarded.contains("\x1b]0;2.1.132\x07"),
1528 "the original title should still have been forwarded verbatim: {forwarded:?}"
1529 );
1530 assert!(
1533 forwarded.contains("\x1b]0;🟡 Opus · 2.1.132\x07"),
1534 "the title should have been corrected once the model became known: {forwarded:?}"
1535 );
1536 let uncorrected_at = forwarded.find("\x1b]0;2.1.132\x07").unwrap();
1537 let corrected_at = forwarded.find("\x1b]0;🟡 Opus · 2.1.132\x07").unwrap();
1538 assert!(uncorrected_at < corrected_at);
1539 assert!(forwarded.ends_with("plain text after the title\n"));
1540 }
1541
1542 }