Skip to main content

omni_dev/cli/
claude_wrap.rs

1//! `omni-dev claude-wrap` — a transparent wrapper around the `claude` process,
2//! teeing its `--output-format stream-json` stdio to the daemon's `sessions`
3//! service as the authoritative Feed 4.
4//!
5//! The Claude Code VS Code extension launches Claude through whatever executable
6//! its `claudeCode.claudeProcessWrapper` setting names, as
7//! `<wrapper> <real-cmd> <real-args…>`. Pointing that at this command puts us in
8//! the stream the extension itself reads, which is the only place the exact
9//! session state — in particular the `can_use_tool` permission prompt — is
10//! visible. See ADR-0057, and [`crate::sessions::stream`] for the state machine.
11//!
12//! **Fail-open is the hard rule.** This sits in Claude's launch path, so the
13//! worst case must be "lose state visibility", never "Claude won't launch". The
14//! byte-forwarding path therefore never awaits the parser, the daemon, or
15//! anything else: lines are handed to the observer through a *bounded* channel
16//! with [`try_send`](tokio::sync::mpsc::Sender::try_send) and dropped when it is
17//! full, and every reporting error is swallowed exactly as the `sessions hook`
18//! sink swallows its own.
19//!
20//! Conversation content is never logged or persisted. An optional metadata-only
21//! file sink is enabled by `OMNI_DEV_CLAUDE_WRAP_LOG`; it never runs on the byte pump. The observer reads only the state, the
22//! `session_id`, the `cwd` and the model out of the stream, and reports them to
23//! the daemon's existing `0600` Unix socket.
24
25use 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
43/// The `sessions` service routing key on the daemon control socket.
44const SERVICE: &str = "sessions";
45
46mod diagnostics;
47use diagnostics::{increment, Diagnostics};
48
49/// How long a fire-and-forget report waits for the daemon before giving up.
50/// Short, and on a task that no byte ever waits behind.
51const REPORT_TIMEOUT: Duration = Duration::from_secs(2);
52
53/// How often a *busy* state is re-reported while the child lives.
54///
55/// Comfortably inside the registry's 300s session TTL, so a long silent turn or
56/// an unanswered permission prompt never ages out. An idle session is not
57/// re-reported (#1454, see [`StreamTracker::keepalive`]): the Claude VS Code
58/// extension keeps a process per chat rather than per visible tab, so
59/// process lifetime is not a sign anyone can still see the chat.
60const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
61
62/// How many teed lines may be in flight before further ones are dropped.
63/// Generous for the two-second worst case of a slow daemon, and bounded so a
64/// wedged observer can never grow memory or slow the child's I/O.
65const TEE_CAPACITY: usize = 256;
66
67/// Longest stream-json line the tee will reassemble; longer lines are forwarded
68/// verbatim like every other byte but not parsed. Matches the daemon control
69/// socket's own line ceiling.
70const MAX_LINE_BYTES: usize = 1024 * 1024;
71
72/// Read-buffer size for each direction of the byte pump.
73const PUMP_BUF_BYTES: usize = 64 * 1024;
74
75/// Exit status used when the child is killed by a signal, mirroring the shell's
76/// `128 + signo` convention.
77const SIGNAL_EXIT_BASE: i32 = 128;
78
79/// Longest we will hold bytes back while looking for the parameter or
80/// terminator of a candidate OSC `0`/`2` title sequence, across any number of
81/// `read()` calls. Bounds memory and is the fail-open backstop for the title
82/// rewrite: a malformed or pathologically split sequence is flushed
83/// unchanged, rather than held back indefinitely, once this is exceeded.
84const MAX_OSC_SEQUENCE_BYTES: usize = 8 * 1024;
85
86/// Environment variable that, set to a truthy value, disables the OSC
87/// title rewrite entirely: Claude's own title sequence is then forwarded
88/// unchanged, exactly as it was before issue #1445. An escape hatch for the
89/// rare terminal/font where the rewritten title misrenders.
90const NO_TITLE_REWRITE_ENV: &str = "OMNI_DEV_CLAUDE_WRAP_NO_TITLE_REWRITE";
91
92/// Runs a command transparently, reporting the Claude session state it streams.
93#[derive(Parser)]
94pub struct ClaudeWrapCommand {
95    /// Path to the daemon control socket. Defaults to the standard per-user path.
96    #[arg(long, value_name = "PATH")]
97    pub socket: Option<PathBuf>,
98
99    /// The command to run and its arguments, after a `--` separator.
100    #[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    /// Executes the wrapper: runs the given command, forwards its stdio
110    /// verbatim, and exits with the child's own status.
111    ///
112    /// This never returns normally on the wrapped path — it calls
113    /// [`std::process::exit`] so the child's exit code survives, which the
114    /// framework's `Result`-to-0/1 collapse in `main` could not express. The
115    /// consequence is that a `claude-wrap` invocation writes no request-log
116    /// record, which is deliberate: it is a process shim, not a user command.
117    pub async fn execute(self) -> Result<()> {
118        let code = self.run().await?;
119        std::process::exit(code)
120    }
121
122    /// Runs the wrapped command and returns the exit code it should be reported
123    /// with. Split out of [`execute`](Self::execute) so everything but the
124    /// process-ending `exit` itself is testable.
125    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        // An interactive launch is not the stream-json protocol, so there is
134        // nothing to observe and no reason to sit in the middle of it: replace
135        // this process with the child outright, which is as transparent as it
136        // gets (same pid, same terminal, same job control, same exit status).
137        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
145/// Replaces this process with `program`, returning the error only if the `exec`
146/// itself failed (on success it never returns).
147fn 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
152/// Wraps `program`, joining it to this process's own stdin and stdout.
153async fn wrap(program: &str, args: &[String], socket: Option<PathBuf>) -> Result<i32> {
154    // omni-dev: coverage ignore reason="process-bound wiring shell: it joins this process's own stdin/stdout and reads the real OMNI_DEV_CLAUDE_WRAP_LOG, which a test must not take over (see the note on run/wrap in the tests module); wrap_io_diagnostics beneath it is covered directly"
155    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    // omni-dev: coverage end
165}
166
167/// Spawns the child with piped stdio, pumps `input` and `output` through it
168/// verbatim while teeing both directions to the observer, and returns the
169/// child's exit code.
170///
171/// Generic over the two endpoints so the wrapping can be tested against a real
172/// child process without touching the test runner's own stdio — which is also
173/// the only way to assert that the forwarding is byte-for-byte lossless.
174#[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    // stderr is inherited and env is left untouched, so the child sees exactly
204    // the environment it would have without the wrapper. The child is
205    // deliberately *not* put in its own process group (unlike the managed
206    // `claude -p` subprocess in `crate::claude::ai::claude_cli`): a transparent
207    // wrapper must leave the child in the caller's group or job control breaks.
208    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    // The title rewrite only ever applies to the FromClaude direction — it
239    // rewrites what Claude asserts about itself, never what we send it.
240    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    // Draining the child's stdout to EOF *before* reaping is what makes the
260    // wrapper lossless: every byte the child wrote reaches our stdout, exactly
261    // as it would have without us in the middle.
262    let _ = from_child.await;
263    // Our own stdin never reaches EOF on its own, so the input pump has to be
264    // cancelled; awaiting the aborted handle drops its tee sender, which is what
265    // eventually lets the observer finish.
266    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/// Copies `reader` to `writer` byte-for-byte, teeing complete lines to the
285/// observer as it goes.
286///
287/// The copy is the priority: bytes are written and flushed before the tee is
288/// even attempted, and the tee is a non-blocking [`try_send`] whose failure is
289/// ignored. No I/O here can be delayed by the observer or the daemon.
290///
291/// `title_rx` is the one deliberate, narrow exception to "never depends on
292/// the observer" (see the [`TitleRewriter`] doc comment and the ADR-0057
293/// amendment for #1445): when present, each chunk is passed through
294/// [`TitleRewriter`] before being written, which does a single non-blocking
295/// [`watch::Receiver::borrow`] of a value the observer publishes — never an
296/// `.await` on the observer or the daemon. `None` (always the `ToClaude`
297/// direction, and `FromClaude` too when the rewrite is disabled) skips the
298/// rewrite entirely and writes the chunk verbatim, exactly as before #1445.
299///
300/// A second, independent branch races [`watch::Receiver::changed`] against
301/// the read: real Claude asserts its title once at startup, not on a timer,
302/// so waiting only on chunks read *from Claude* would very often lose the
303/// race against the observer classifying the model — with nothing left to
304/// reassert and catch it later. When the model becomes known (or changes),
305/// this branch calls [`TitleRewriter::resync`] to correct the cached title
306/// immediately, with no new bytes from Claude required. It is still
307/// non-blocking on the observer in the sense that matters: a value already
308/// sitting in the channel is read instantly, and if the observer never
309/// updates it, this branch simply never fires — it cannot hang the pump.
310///
311/// [`try_send`]: tokio::sync::mpsc::Sender::try_send
312#[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                // No receiver (ToClaude, or the rewrite is disabled): never
353                // resolves, so `select!` only ever waits on the read.
354                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                    // The observer's sender was dropped; stop selecting on
391                    // this branch rather than spinning on a channel that
392                    // will report closed forever.
393                    Err(_) => title_rx = None,
394                }
395            }
396        }
397    }
398    // Closing the write end propagates EOF (this is how the child learns its
399    // input has ended); a failure here means the peer is already gone.
400    let _ = writer.shutdown().await;
401}
402
403/// Rewrites Claude's own OSC `0`/`2` terminal-title sequence in flight,
404/// prepending a colour-glyph + model-family prefix, while leaving every other
405/// byte untouched (issue #1445).
406///
407/// **Fail-open by construction**, matching every other guarantee this module
408/// makes: a candidate sequence is forwarded byte-for-byte unchanged, exactly
409/// as received, whenever it
410/// - is not an OSC `0` or `2` sequence (any other OSC code, or a bare `ESC`
411///   not followed by `]`),
412/// - is not terminated with `BEL` or the two-byte `ESC \` (ST) form,
413/// - carries a title that is not valid UTF-8,
414/// - exceeds [`MAX_OSC_SEQUENCE_BYTES`] before terminating, or
415/// - arrives while no model is known yet (`prefix` is `None`).
416///
417/// State (`pending`/`title`/`code`/`phase`) persists across calls so a
418/// sequence split unfavourably across a `read()` buffer boundary is still
419/// recognized and rewritten correctly.
420///
421/// Real Claude Code asserts its title **once**, as part of its terminal-mode
422/// setup at startup, not on a timer — measured empirically, not the
423/// "continuously re-asserted" behavior issue #1445 assumed. Left as
424/// rewrite-on-assert alone, that one assertion almost always races ahead of
425/// the model becoming known (the observer needs to receive, tee, and parse
426/// the `system`/`init` line first), and since there is no later reassertion
427/// to catch, the title is never corrected for the rest of the session. Every
428/// completed sequence is therefore also cached — regardless of whether it
429/// was rewritten or forwarded unchanged — so [`Self::resync`] can re-emit a
430/// corrected title the moment the model becomes known (or changes again),
431/// without waiting for Claude to assert anything else.
432#[derive(Debug, Default)]
433struct TitleRewriter {
434    /// Every raw byte seen since leaving [`Phase::Idle`] — the exact bytes
435    /// flushed verbatim whenever a candidate sequence does not end up being
436    /// rewritten.
437    pending: Vec<u8>,
438    /// The title text accumulated once past `ESC ] <code> ;`, excluding the
439    /// terminator.
440    title: Vec<u8>,
441    /// The OSC parameter byte (`b'0'` or `b'2'`), once known.
442    code: u8,
443    phase: Phase,
444    /// The most recently completed sequence's parameter byte, title text,
445    /// and terminator — cached regardless of whether it was rewritten, so
446    /// [`Self::resync`] can reformat it under a new prefix.
447    last_code: Option<u8>,
448    last_title: Option<Vec<u8>>,
449    last_terminator: Option<Vec<u8>>,
450}
451
452/// `TitleRewriter`'s position within (or outside of) a candidate OSC title
453/// sequence.
454#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
455enum Phase {
456    /// Not inside any candidate escape sequence; bytes pass straight through.
457    #[default]
458    Idle,
459    /// Saw `ESC`; deciding whether `]` follows.
460    Esc,
461    /// Saw `ESC ]`; deciding whether a `0`/`2` title code follows.
462    Open,
463    /// Saw `ESC ] <code>`; deciding whether `;` follows.
464    Code,
465    /// Accumulating title text after `ESC ] <code> ;`, until `BEL` or ST.
466    Title,
467    /// Saw `ESC` while accumulating title text; deciding whether `\` (the
468    /// second byte of the ST terminator) follows.
469    TitleEsc,
470}
471
472const ESC: u8 = 0x1B;
473const BEL: u8 = 0x07;
474
475impl TitleRewriter {
476    /// Passes `chunk` through the rewriter, returning the bytes that should
477    /// actually be written downstream. `prefix` is the current
478    /// `"{glyph} {label} · "` string to prepend to a rewritten title, or
479    /// `None` while no model is known yet.
480    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                        // Not a valid ST terminator: the whole candidate is
533                        // malformed, so flush verbatim rather than guess.
534                        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    /// A completed sequence: rewrites it when `prefix` and a valid UTF-8
546    /// title are both available, otherwise flushes the original bytes
547    /// unchanged. Cached either way, for [`Self::resync`].
548    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    /// Re-emits the most recently completed title sequence under a new
569    /// `prefix`, so a title Claude already asserted — before any model was
570    /// known, or under a now-stale one — is corrected without waiting for
571    /// Claude to reassert it on its own (which, empirically, it may never
572    /// do again for the rest of the session). Returns `None` when nothing
573    /// has completed yet, `prefix` is still unavailable, or the cached title
574    /// is not valid UTF-8.
575    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    /// Flushes everything buffered so far, unchanged, and returns to `Idle`.
588    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
601/// Whether the OSC title rewrite is enabled — the default, unless
602/// [`NO_TITLE_REWRITE_ENV`] is set to a truthy value.
603fn title_rewrite_enabled() -> bool {
604    !std::env::var(NO_TITLE_REWRITE_ENV).is_ok_and(|v| flag_is_truthy(&v))
605}
606
607/// Whether a `NO_TITLE_REWRITE_ENV`-style flag value should be read as "on".
608/// Split out from [`title_rewrite_enabled`] so the parsing itself is testable
609/// without mutating process-wide environment state.
610fn flag_is_truthy(value: &str) -> bool {
611    matches!(
612        value.trim().to_ascii_lowercase().as_str(),
613        "1" | "true" | "yes" | "on"
614    )
615}
616
617/// Builds the `"{glyph} {label} · "` prefix to prepend to Claude's own OSC
618/// title text for `model_id`, classified via the shared [`ModelRegistry`]
619/// (issue #1445).
620///
621/// [`ModelRegistry`]: crate::claude::model_config::ModelRegistry
622fn 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/// Reassembles newline-delimited lines out of a forwarded chunk and offers each
628/// to the observer, dropping any that is over-long, not UTF-8, or arrives while
629/// the channel is full.
630#[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/// Consumes teed lines, tracks the session state, and reports every change to
673/// the daemon — plus a keep-alive while a busy state is unchanged, and an `end`
674/// once the stream is over.
675/// `every` is the keep-alive cadence, injected so tests can drive it without
676/// waiting out the real [`KEEPALIVE_INTERVAL`]. `child_pid` is the wrapped
677/// `claude`, sent on every report — the same pid its hooks send — so the `end`
678/// this wrapper sends when an in-place resume has already replaced its child
679/// cannot end the resumed session (#1948). It is also the seed for pid-based
680/// liveness (#1916), whose identity-token reading is entirely the daemon's own
681/// [`pid_watcher`](crate::sessions::pid_watcher) work, never this wrapper's.
682#[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        // omni-dev: coverage ignore reason="resolve_socket fails only when the platform has no data directory to put the default socket in (no resolvable home), which a test cannot reproduce on macOS or Linux; the fail-open return is what keeps the wrapper forwarding"
711        diagnostics
712            .record(|| json!({"event":"observer_stopped", "outcome":"socket_resolution_failed"}));
713        return;
714        // omni-dev: coverage end
715    };
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    // The first tick of a tokio interval completes immediately; consume it so
721    // the keep-alive does not fire before anything has been observed.
722    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        // Publish the classified title prefix whenever the model changes —
736        // independent of `observed`, since a mid-session model switch does
737        // not necessarily move `SessionState` (issue #1445). The pump's
738        // `watch::Receiver::borrow` on the other end never blocks on this.
739        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                // omni-dev: coverage ignore reason="to_value on an ObserveRequest fails only for a non-UTF-8 cwd, and the tracker takes cwd from a JSON string, so it is always UTF-8; the arm exists so a future non-string field cannot silently drop the report"
768                diagnostics.record(
769                    || json!({"event":"report", "op":"observe", "outcome":"serialization_failed"}),
770                );
771                // omni-dev: coverage end
772            }
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
786/// Sends one bounded, fire-and-forget op to the daemon's `sessions` service,
787/// swallowing every failure — a missing or wedged daemon must be a silent no-op.
788async 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
802/// Relays `SIGINT`/`SIGTERM` to the child, so a caller that signals the wrapper
803/// stops Claude rather than orphaning it.
804async 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
824/// The exit code to leave with: the child's own, or the shell's `128 + signo`
825/// when it died on a signal.
826fn 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        // Bound but never accepted: the connect lands in the backlog and the
975        // request in the socket buffer, so only the report timeout can end the
976        // exchange. With the clock paused, that two-second wait costs no wall time.
977        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    /// An [`AsyncWrite`] that appends into a shared buffer, so a test can assert
987    /// on exactly the bytes the wrapper forwarded.
988    #[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    /// Runs `/bin/sh -c script` through the wrapper, returning its exit code and
1023    /// everything it wrote downstream.
1024    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    /// Parses a `claude-wrap` argument list the way the real CLI would.
1039    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    /// A one-shot daemon stand-in that records every envelope it is sent.
1046    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    // Note: there is deliberately no test that drives `run`/`wrap` to completion.
1105    // Both join the *process's own* stdin and stdout, which a test must not take
1106    // over — and `run`'s terminal branch would `exec`-replace the test binary
1107    // outright when the suite is run from an interactive shell. They are thin
1108    // delegations; `wrap_io` underneath them is covered directly.
1109
1110    // Note: `exec_replace` is deliberately untested. On success it never returns,
1111    // and its failure path cannot be exercised in-process either: `exec` resets
1112    // `SIGPIPE` to `SIG_DFL` before the `execve`, and a failed `exec` leaves that
1113    // reset in place (the "broken state" its docs warn about). Calling it from a
1114    // test poisons the whole binary — every later test that writes to a closed
1115    // pipe dies on signal 13.
1116
1117    #[tokio::test]
1118    async fn signal_forwarding_gives_up_when_there_is_no_child() {
1119        // A child already reaped has no pid to signal; this must return rather
1120        // than spin.
1121        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        // A partial trailing line and a non-JSON line: neither is parseable, and
1128        // both must still arrive downstream byte-for-byte.
1129        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        // A minimal but real turn: announce, work, ask permission, finish.
1139        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        // Identity rides the reports, and the stream's end is announced.
1160        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        // A socket path that does not exist: every report fails silently.
1171        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        // Only the last (short, valid UTF-8) line is offered to the observer.
1213        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        // A turn that goes quiet (a long tool call, an unanswered prompt) must
1222        // keep being reported, or the registry ages it out at 5 minutes.
1223        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        // One for init, one for the state change, then repeats of that state.
1248        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            // The wrapped child's pid, so a resume can tell it apart (#1948).
1253            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        // #1454: the extension keeps a process per chat, not per visible tab, so
1263        // an idle one is reported once and then not pinned.
1264        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        // A closed peer (the editor exiting mid-turn) must end the copy, not
1296        // spin on a writer that can never succeed again.
1297        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        // Returns rather than hanging, and tees nothing it could not forward.
1321        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        // Isolates pump()'s select! loop from the rest of the observe()
1335        // pipeline: drives the watch channel directly, with a reader that
1336        // never offers a second chunk, to prove the resync path does not
1337        // depend on Claude sending anything else.
1338        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        // Two of the three were dropped; nothing blocked, nothing panicked.
1373        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        // OSC 8 (hyperlink) must never be mistaken for a title sequence.
1429        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        // ESC not immediately followed by `\` is not a valid ST: the whole
1447        // candidate must be forwarded verbatim rather than guessed at.
1448        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        // The exact real-world case this exists for: Claude's one-and-only
1470        // title assertion raced ahead of the model becoming known, so it
1471        // went out unrewritten — resync must still be able to correct it
1472        // later, from the cache, with no further help from Claude.
1473        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        // Models the real race, not the happy path: real Claude asserts its
1508        // title once, at startup, *before* the model can possibly be known
1509        // (the observer has not even received the init line yet) — and,
1510        // empirically, does not reassert it later in the session. So the
1511        // very first assertion here is forwarded unrewritten (correct,
1512        // fail-open behavior), and the `sleep` that follows leaves the child
1513        // otherwise idle long enough for the async observer pipeline to
1514        // classify the model and for the pump's watch-driven resync to
1515        // correct the title on its own, with no further help from Claude.
1516        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        // The original, un-decorated assertion still reaches stdout exactly
1524        // as Claude sent it — fail-open means "forward unchanged", not "hold
1525        // back and hope".
1526        assert!(
1527            forwarded.contains("\x1b]0;2.1.132\x07"),
1528            "the original title should still have been forwarded verbatim: {forwarded:?}"
1529        );
1530        // And it is followed, with no further title assertion from Claude,
1531        // by a corrected one carrying the classified model prefix.
1532        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    // Note: there is deliberately no end-to-end test that sets
1543    // `NO_TITLE_REWRITE_ENV` and drives `wrap_script` — mutating that real
1544    // variable would race any other test in this module concurrently
1545    // spawning `wrap_io` (which reads it too), exactly the hazard
1546    // `request_log.rs`'s env-var tests avoid by mutating a private test-only
1547    // name instead. `flag_is_truthy_recognises_on_and_off_values` above
1548    // covers the parsing directly; the one-line `.then_some(title_rx)` wiring
1549    // in `wrap_io` needs no separate test.
1550}