Skip to main content

claude_wrapper/
streaming.rs

1//! NDJSON streaming of `claude` events.
2//!
3//! [`stream_query`] (and its blocking peer `stream_query_sync`) run a
4//! query in `stream-json` mode and hand each decoded event to a
5//! caller-supplied callback as it arrives, rather than buffering the
6//! whole run. Requires the `json` feature.
7
8#[cfg(feature = "json")]
9use std::time::Duration;
10
11#[cfg(all(feature = "json", feature = "async"))]
12use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
13#[cfg(all(feature = "json", feature = "async"))]
14use tokio::process::{ChildStderr, Command};
15#[cfg(feature = "json")]
16use tracing::{debug, warn};
17
18#[cfg(feature = "json")]
19use crate::Claude;
20#[cfg(feature = "json")]
21use crate::error::{Error, Result};
22#[cfg(feature = "json")]
23use crate::exec::CommandOutput;
24
25/// A single line from `--output-format stream-json` output.
26///
27/// Each line is an NDJSON object. The structure varies by message type,
28/// so we provide the raw JSON value and convenience accessors.
29#[cfg(feature = "json")]
30#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
31pub struct StreamEvent {
32    /// The raw JSON object for this event.
33    #[serde(flatten)]
34    pub data: serde_json::Value,
35}
36
37#[cfg(feature = "json")]
38impl StreamEvent {
39    /// Get the event type, if present.
40    pub fn event_type(&self) -> Option<&str> {
41        self.data.get("type").and_then(|v| v.as_str())
42    }
43
44    /// Get the message role, if present.
45    pub fn role(&self) -> Option<&str> {
46        self.data.get("role").and_then(|v| v.as_str())
47    }
48
49    /// Check if this is the final result message.
50    pub fn is_result(&self) -> bool {
51        self.event_type() == Some("result")
52    }
53
54    /// Extract the result text from a result event.
55    pub fn result_text(&self) -> Option<&str> {
56        self.data.get("result").and_then(|v| v.as_str())
57    }
58
59    /// Get the session ID if present.
60    pub fn session_id(&self) -> Option<&str> {
61        self.data.get("session_id").and_then(|v| v.as_str())
62    }
63
64    /// Get the cost in USD if present (usually on result events).
65    ///
66    /// Prefers `total_cost_usd` (the CLI's primary key) and falls back
67    /// to the legacy `cost_usd` alias.
68    pub fn cost_usd(&self) -> Option<f64> {
69        self.data
70            .get("total_cost_usd")
71            .or_else(|| self.data.get("cost_usd"))
72            .and_then(|v| v.as_f64())
73    }
74
75    /// Decode a partial-message event into a typed view.
76    ///
77    /// Returns `Some` when the event is one of the content-block lifecycle
78    /// events surfaced by [`QueryCommand::include_partial_messages`] -- start,
79    /// delta, or stop. Returns `None` for any other event (system, assistant,
80    /// result, message-level stream events, etc).
81    ///
82    /// The CLI wraps each raw streaming event as
83    /// `{"type":"stream_event","event":{...}}`; this accessor unwraps that
84    /// envelope. Unknown block types and unknown delta types fall through to
85    /// [`BlockType::Other`] / [`BlockDelta::Other`] rather than erroring, so
86    /// future content-block kinds remain accessible (just untyped).
87    ///
88    /// # Example
89    ///
90    /// Pull incremental thinking text out of a partial-message event:
91    ///
92    /// ```
93    /// use claude_wrapper::streaming::{BlockDelta, PartialMessageEvent, StreamEvent};
94    /// use serde_json::json;
95    ///
96    /// let event: StreamEvent = serde_json::from_value(json!({
97    ///     "type": "stream_event",
98    ///     "event": {
99    ///         "type": "content_block_delta",
100    ///         "index": 0,
101    ///         "delta": { "type": "thinking_delta", "thinking": "Let me think..." }
102    ///     },
103    ///     "session_id": "abc"
104    /// })).unwrap();
105    ///
106    /// match event.partial_message() {
107    ///     Some(PartialMessageEvent::BlockDelta { delta: BlockDelta::Thinking(t), .. }) => {
108    ///         assert_eq!(t, "Let me think...");
109    ///     }
110    ///     _ => unreachable!(),
111    /// }
112    /// ```
113    ///
114    /// [`QueryCommand::include_partial_messages`]: crate::QueryCommand::include_partial_messages
115    pub fn partial_message(&self) -> Option<PartialMessageEvent> {
116        let event = if self.event_type() == Some("stream_event") {
117            self.data.get("event")?
118        } else {
119            &self.data
120        };
121
122        let inner_type = event.get("type")?.as_str()?;
123        let index = event.get("index").and_then(serde_json::Value::as_u64)?;
124        let index = u32::try_from(index).ok()?;
125
126        match inner_type {
127            "content_block_start" => {
128                let block_type = parse_block_type(event.get("content_block")?);
129                Some(PartialMessageEvent::BlockStart { index, block_type })
130            }
131            "content_block_delta" => {
132                let delta = parse_block_delta(event.get("delta")?);
133                Some(PartialMessageEvent::BlockDelta { index, delta })
134            }
135            "content_block_stop" => Some(PartialMessageEvent::BlockStop { index }),
136            _ => None,
137        }
138    }
139}
140
141/// A decoded partial-message event from a streaming `claude` call.
142///
143/// Surfaced by [`StreamEvent::partial_message`] when `--include-partial-messages`
144/// is set. The three variants correspond to the Anthropic streaming content-block
145/// lifecycle: a block starts, gets one or more deltas, then stops.
146#[cfg(feature = "json")]
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub enum PartialMessageEvent {
149    /// A new content block is starting. `block_type` says what kind.
150    BlockStart {
151        /// Position of this block within the assistant message.
152        index: u32,
153        /// What kind of block is starting (text, thinking, tool use, ...).
154        block_type: BlockType,
155    },
156    /// Incremental content for an in-progress block.
157    BlockDelta {
158        /// Index of the block this delta applies to (matches a prior [`BlockStart`]).
159        ///
160        /// [`BlockStart`]: PartialMessageEvent::BlockStart
161        index: u32,
162        /// The incremental payload.
163        delta: BlockDelta,
164    },
165    /// The block at `index` is complete.
166    BlockStop {
167        /// Index of the block that just finished.
168        index: u32,
169    },
170}
171
172/// The kind of content block reported by a [`PartialMessageEvent::BlockStart`].
173///
174/// Mirrors the `content_block.type` field from the Anthropic streaming API.
175/// New block kinds added upstream surface as [`BlockType::Other`] -- callers
176/// can still recover the type name from the carried string.
177#[cfg(feature = "json")]
178#[derive(Debug, Clone, PartialEq, Eq)]
179pub enum BlockType {
180    /// Regular assistant text -- followed by `text_delta` deltas.
181    Text,
182    /// Extended-thinking block -- followed by `thinking_delta` deltas.
183    Thinking,
184    /// A tool invocation -- followed by `input_json_delta` deltas streaming the JSON input.
185    ToolUse {
186        /// Tool-call id, used to correlate the eventual tool result.
187        id: String,
188        /// Name of the tool being called.
189        name: String,
190    },
191    /// Any block type not yet modelled. Carries the raw `type` string.
192    Other(String),
193}
194
195/// The incremental payload carried by a [`PartialMessageEvent::BlockDelta`].
196///
197/// Mirrors the `delta.type` field from the Anthropic streaming API.
198/// Less-common delta kinds (signature, citations, compaction, ...) collapse to
199/// [`BlockDelta::Other`]; callers that need them can fall back to
200/// [`StreamEvent::data`].
201#[cfg(feature = "json")]
202#[derive(Debug, Clone, PartialEq, Eq)]
203pub enum BlockDelta {
204    /// Chunk of assistant text.
205    Text(String),
206    /// Chunk of extended-thinking text.
207    Thinking(String),
208    /// Chunk of streaming tool-input JSON. Concatenate across deltas to
209    /// reconstruct the full input -- individual chunks are not standalone JSON.
210    InputJson(String),
211    /// Any delta type not modelled above (e.g. `signature_delta`,
212    /// `citations_delta`). Read from [`StreamEvent::data`] for the raw payload.
213    Other,
214}
215
216#[cfg(feature = "json")]
217fn parse_block_type(content_block: &serde_json::Value) -> BlockType {
218    let Some(ty) = content_block
219        .get("type")
220        .and_then(serde_json::Value::as_str)
221    else {
222        return BlockType::Other(String::new());
223    };
224    match ty {
225        "text" => BlockType::Text,
226        "thinking" => BlockType::Thinking,
227        "tool_use" => {
228            let id = content_block
229                .get("id")
230                .and_then(serde_json::Value::as_str)
231                .unwrap_or("")
232                .to_string();
233            let name = content_block
234                .get("name")
235                .and_then(serde_json::Value::as_str)
236                .unwrap_or("")
237                .to_string();
238            BlockType::ToolUse { id, name }
239        }
240        other => BlockType::Other(other.to_string()),
241    }
242}
243
244#[cfg(feature = "json")]
245fn parse_block_delta(delta: &serde_json::Value) -> BlockDelta {
246    let Some(ty) = delta.get("type").and_then(serde_json::Value::as_str) else {
247        return BlockDelta::Other;
248    };
249    match ty {
250        "text_delta" => delta
251            .get("text")
252            .and_then(serde_json::Value::as_str)
253            .map(|s| BlockDelta::Text(s.to_string()))
254            .unwrap_or(BlockDelta::Other),
255        "thinking_delta" => delta
256            .get("thinking")
257            .and_then(serde_json::Value::as_str)
258            .map(|s| BlockDelta::Thinking(s.to_string()))
259            .unwrap_or(BlockDelta::Other),
260        "input_json_delta" => delta
261            .get("partial_json")
262            .and_then(serde_json::Value::as_str)
263            .map(|s| BlockDelta::InputJson(s.to_string()))
264            .unwrap_or(BlockDelta::Other),
265        _ => BlockDelta::Other,
266    }
267}
268
269/// Execute a command with streaming output, calling a handler for each NDJSON line.
270///
271/// This spawns the claude process and reads stdout line-by-line, parsing each
272/// as a JSON event and passing it to the handler. Useful for progress tracking
273/// and real-time output processing.
274///
275/// Dropping the returned future mid-flight kills the spawned `claude`
276/// process and, on Unix, its whole process group (SIGKILL): an
277/// abandoned run does not keep executing in the background, and the
278/// subprocesses it spawned for tool use die with it. Events already
279/// dispatched to the handler are not rolled back.
280///
281/// # Example
282///
283/// ```no_run
284/// use claude_wrapper::{Claude, QueryCommand, OutputFormat};
285/// use claude_wrapper::streaming::{StreamEvent, stream_query};
286///
287/// # async fn example() -> claude_wrapper::Result<()> {
288/// let claude = Claude::builder().build()?;
289///
290/// let cmd = QueryCommand::new("explain quicksort")
291///     .output_format(OutputFormat::StreamJson);
292///
293/// let output = stream_query(&claude, &cmd, |event: StreamEvent| {
294///     if let Some(t) = event.event_type() {
295///         println!("[{t}] {:?}", event.data);
296///     }
297/// }).await?;
298/// # Ok(())
299/// # }
300/// ```
301#[cfg(all(feature = "json", feature = "async"))]
302pub async fn stream_query<F>(
303    claude: &Claude,
304    cmd: &crate::command::query::QueryCommand,
305    handler: F,
306) -> Result<CommandOutput>
307where
308    F: FnMut(StreamEvent),
309{
310    stream_query_impl(claude, cmd, handler, claude.timeout).await
311}
312
313/// Unified streaming implementation with optional timeout.
314///
315/// Reads stderr concurrently in a background task so a chatty child
316/// cannot deadlock by filling the stderr pipe buffer, and so any
317/// captured stderr is available even on timeout or IO error.
318///
319/// On timeout, the child is killed and reaped (`kill().await` sends
320/// SIGKILL and waits), and whatever stderr was produced is logged at
321/// warn level. The returned `Error::Timeout` does not carry partial
322/// output -- streamed stdout events were already dispatched to the
323/// handler as they arrived.
324#[cfg(all(feature = "json", feature = "async"))]
325async fn stream_query_impl<F>(
326    claude: &Claude,
327    cmd: &crate::command::query::QueryCommand,
328    mut handler: F,
329    timeout: Option<Duration>,
330) -> Result<CommandOutput>
331where
332    F: FnMut(StreamEvent),
333{
334    use crate::command::ClaudeCommand;
335
336    let args = cmd.args();
337
338    let mut command_args = Vec::new();
339    command_args.extend(claude.global_args.clone());
340    command_args.extend(args);
341
342    // A stream has three outcomes and the failure one is the easy one
343    // to miss, so `outcome` is recorded on every exit path: completed,
344    // failed, or timeout. `events` counts what the handler actually
345    // saw, which is the difference between "produced nothing" and
346    // "never started".
347    let span = tracing::debug_span!(
348        "claude.stream",
349        command = crate::exec::span_command(&command_args),
350        binary = %claude.binary.display(),
351        cwd = claude.working_dir.as_deref().map(|d| d.display().to_string()),
352        timeout_secs = timeout.map(|t| t.as_secs()),
353        outcome = tracing::field::Empty,
354        events = tracing::field::Empty,
355        exit_code = tracing::field::Empty,
356        duration_ms = tracing::field::Empty,
357    );
358    let _enter = span.enter();
359    let started = std::time::Instant::now();
360    let mut event_count: u64 = 0;
361
362    debug!(
363        binary = %claude.binary.display(),
364        args = ?command_args,
365        timeout = ?timeout,
366        "streaming claude command"
367    );
368
369    let mut cmd = Command::new(&claude.binary);
370    cmd.args(&command_args)
371        .env_remove("CLAUDECODE")
372        .envs(&claude.env)
373        .stdout(std::process::Stdio::piped())
374        .stderr(std::process::Stdio::piped())
375        .stdin(std::process::Stdio::null())
376        // Dropping the in-flight future must kill the child, not leave
377        // the CLI running unattended (see the `exec` module docs).
378        .kill_on_drop(true);
379    // Own process group (Unix) so cancellation can signal the whole
380    // tree, not just the direct child (see exec::GroupKillGuard). Opt
381    // out via ClaudeBuilder::process_group.
382    crate::exec::apply_process_group(&mut cmd, claude.process_group);
383
384    if let Some(ref dir) = claude.working_dir {
385        cmd.current_dir(dir);
386    }
387
388    let mut child = cmd.spawn().map_err(|e| Error::Io {
389        message: format!("failed to spawn claude: {e}"),
390        source: e,
391        working_dir: claude.working_dir.clone(),
392    })?;
393    let mut group =
394        crate::exec::arm_and_notify(claude.process_group, child.id(), claude.on_spawn.as_ref());
395
396    let stdout = child.stdout.take().expect("stdout was piped");
397    let mut stderr = child.stderr.take().expect("stderr was piped");
398
399    let mut reader = BufReader::new(stdout).lines();
400
401    // Run stdout line reading and stderr draining concurrently so a
402    // chatty child can't deadlock by filling the stderr pipe buffer.
403    // tokio::join! polls both futures on the same task (no tokio::spawn
404    // needed, so we avoid pulling in the `rt` feature).
405    let drain = drain_stderr(&mut stderr);
406    // Wrap the caller's handler so the span can report how many events
407    // were dispatched without the handler needing to care.
408    let mut counting_handler = |event: StreamEvent| {
409        event_count += 1;
410        handler(event);
411    };
412    let read_future = read_lines(
413        &mut reader,
414        &mut counting_handler,
415        claude.working_dir.clone(),
416    );
417    let combined = async {
418        let (line_result, stderr_str) = tokio::join!(read_future, drain);
419        (line_result, stderr_str)
420    };
421
422    let (line_result, stderr_str) = match timeout {
423        Some(d) => match tokio::time::timeout(d, combined).await {
424            Ok(pair) => pair,
425            Err(_) => {
426                // Timeout: take down the whole group, honoring the
427                // optional SIGTERM grace, then kill+reap the direct
428                // child, and try to drain whatever stderr remains.
429                // The group kill takes down subprocesses that could
430                // otherwise hold our pipe fds open; the capped drain
431                // below stays as a backstop.
432                crate::exec::kill_group_with_grace(&mut group, claude.kill_grace).await;
433                let _ = child.kill().await;
434                let drain_budget = Duration::from_millis(200);
435                let stderr_str = tokio::time::timeout(drain_budget, drain_stderr(&mut stderr))
436                    .await
437                    .unwrap_or_default();
438                if !stderr_str.is_empty() {
439                    warn!(stderr = %stderr_str, "stderr from timed-out streaming process");
440                }
441                span.record("outcome", "timeout");
442                span.record("events", event_count);
443                span.record("duration_ms", started.elapsed().as_millis() as u64);
444                return Err(Error::Timeout {
445                    timeout_seconds: d.as_secs(),
446                });
447            }
448        },
449        None => combined.await,
450    };
451
452    // If reading lines failed partway through (IO error, not timeout),
453    // clean up the child (and its group) before returning.
454    if let Err(e) = line_result {
455        group.kill_now();
456        let _ = child.kill().await;
457        return Err(e);
458    }
459
460    let status = child.wait().await.map_err(|e| Error::Io {
461        message: "failed to wait for claude process".to_string(),
462        source: e,
463        working_dir: claude.working_dir.clone(),
464    })?;
465    group.disarm();
466
467    let exit_code = status.code().unwrap_or(-1);
468
469    span.record("events", event_count);
470    span.record("exit_code", exit_code);
471    span.record("duration_ms", started.elapsed().as_millis() as u64);
472
473    if !status.success() {
474        span.record("outcome", "failed");
475        return Err(Error::CommandFailed {
476            command: format!("{} {}", claude.binary.display(), command_args.join(" ")),
477            exit_code,
478            stdout: String::new(),
479            stderr: stderr_str,
480            working_dir: claude.working_dir.clone(),
481        });
482    }
483
484    span.record("outcome", "completed");
485    Ok(CommandOutput {
486        stdout: String::new(), // already consumed via streaming
487        stderr: stderr_str,
488        exit_code,
489        success: true,
490    })
491}
492
493#[cfg(all(feature = "json", feature = "async"))]
494async fn drain_stderr(stderr: &mut ChildStderr) -> String {
495    let mut buf = Vec::new();
496    let _ = stderr.read_to_end(&mut buf).await;
497    String::from_utf8_lossy(&buf).into_owned()
498}
499
500#[cfg(all(feature = "json", feature = "async"))]
501async fn read_lines<F>(
502    reader: &mut tokio::io::Lines<BufReader<tokio::process::ChildStdout>>,
503    handler: &mut F,
504    working_dir: Option<std::path::PathBuf>,
505) -> Result<()>
506where
507    F: FnMut(StreamEvent),
508{
509    while let Some(line) = reader.next_line().await.map_err(|e| Error::Io {
510        message: "failed to read stdout line".to_string(),
511        source: e,
512        working_dir: working_dir.clone(),
513    })? {
514        if line.trim().is_empty() {
515            continue;
516        }
517        match serde_json::from_str::<StreamEvent>(&line) {
518            Ok(event) => handler(event),
519            Err(e) => {
520                debug!(line = %line, error = %e, "failed to parse stream event, skipping");
521            }
522        }
523    }
524
525    Ok(())
526}
527
528// ---------- sync streaming ----------
529
530/// Blocking mirror of [`stream_query`]. Reads NDJSON lines from the
531/// child's stdout on a worker thread, dispatches each parsed event
532/// to `handler` on the caller's thread, and drains stderr on a
533/// separate worker thread so the child can't deadlock on a full pipe.
534///
535/// Requires both `sync` and `json` features.
536///
537/// The handler is invoked on the caller's thread — no `Send` bound —
538/// so it can capture non-`Send` state. If a timeout is configured on
539/// the [`Claude`] client, the child's whole process group (Unix) is
540/// SIGKILLed and the child reaped once the deadline passes; partial
541/// events already dispatched to the handler are not rolled back.
542///
543/// # Example
544///
545/// ```no_run
546/// # #[cfg(all(feature = "sync", feature = "json"))]
547/// # {
548/// use claude_wrapper::{Claude, OutputFormat, QueryCommand};
549/// use claude_wrapper::streaming::{StreamEvent, stream_query_sync};
550///
551/// # fn example() -> claude_wrapper::Result<()> {
552/// let claude = Claude::builder().build()?;
553/// let cmd = QueryCommand::new("explain quicksort")
554///     .output_format(OutputFormat::StreamJson);
555///
556/// stream_query_sync(&claude, &cmd, |event: StreamEvent| {
557///     if let Some(t) = event.event_type() {
558///         println!("[{t}] {:?}", event.data);
559///     }
560/// })?;
561/// # Ok(())
562/// # }
563/// # }
564/// ```
565#[cfg(all(feature = "sync", feature = "json"))]
566pub fn stream_query_sync<F>(
567    claude: &Claude,
568    cmd: &crate::command::query::QueryCommand,
569    mut handler: F,
570) -> Result<CommandOutput>
571where
572    F: FnMut(StreamEvent),
573{
574    use std::io::{BufRead as _, Read as _};
575    use std::process::{Command as StdCommand, Stdio};
576    use std::sync::mpsc;
577    use std::thread;
578    use std::time::Instant;
579
580    use crate::command::ClaudeCommand;
581
582    let args = cmd.args();
583    let mut command_args = Vec::new();
584    command_args.extend(claude.global_args.clone());
585    command_args.extend(args);
586
587    debug!(
588        binary = %claude.binary.display(),
589        args = ?command_args,
590        timeout = ?claude.timeout,
591        "streaming claude command (sync)"
592    );
593
594    let mut cmd_builder = StdCommand::new(&claude.binary);
595    cmd_builder
596        .args(&command_args)
597        .env_remove("CLAUDECODE")
598        .env_remove("CLAUDE_CODE_ENTRYPOINT")
599        .envs(&claude.env)
600        .stdin(Stdio::null())
601        .stdout(Stdio::piped())
602        .stderr(Stdio::piped());
603    // Own process group (Unix) so a kill can signal the whole tree,
604    // not just the direct child (see exec::GroupKillGuard). Opt out
605    // via ClaudeBuilder::process_group.
606    crate::exec::apply_process_group_sync(&mut cmd_builder, claude.process_group);
607
608    if let Some(ref dir) = claude.working_dir {
609        cmd_builder.current_dir(dir);
610    }
611
612    let mut child = cmd_builder.spawn().map_err(|e| Error::Io {
613        message: format!("failed to spawn claude: {e}"),
614        source: e,
615        working_dir: claude.working_dir.clone(),
616    })?;
617    let mut group = crate::exec::arm_and_notify(
618        claude.process_group,
619        Some(child.id()),
620        claude.on_spawn.as_ref(),
621    );
622
623    let stdout = child.stdout.take().expect("stdout was piped");
624    let stderr = child.stderr.take().expect("stderr was piped");
625
626    // Reader thread: parse NDJSON lines and push StreamEvents through
627    // the channel. Handler runs on the caller's thread so it doesn't
628    // need Send. Bubbles IO errors out via the thread's return value.
629    let (tx, rx) = mpsc::channel::<StreamEvent>();
630    let reader_wd = claude.working_dir.clone();
631    let reader_thread = thread::spawn(move || -> Result<()> {
632        let reader = std::io::BufReader::new(stdout);
633        for line_res in reader.lines() {
634            let line = line_res.map_err(|e| Error::Io {
635                message: "failed to read stdout line".to_string(),
636                source: e,
637                working_dir: reader_wd.clone(),
638            })?;
639            if line.trim().is_empty() {
640                continue;
641            }
642            match serde_json::from_str::<StreamEvent>(&line) {
643                Ok(event) => {
644                    if tx.send(event).is_err() {
645                        // Receiver gone — main thread has bailed out.
646                        return Ok(());
647                    }
648                }
649                Err(e) => {
650                    debug!(line = %line, error = %e, "failed to parse stream event, skipping");
651                }
652            }
653        }
654        Ok(())
655    });
656
657    let stderr_thread = thread::spawn(move || -> String {
658        let mut buf = Vec::new();
659        let mut stderr = stderr;
660        let _ = stderr.read_to_end(&mut buf);
661        String::from_utf8_lossy(&buf).into_owned()
662    });
663
664    // Main loop: dispatch events on the caller's thread, honouring the
665    // configured timeout. Break on disconnect (reader done) or timeout.
666    let deadline = claude.timeout.map(|d| Instant::now() + d);
667    let mut timed_out = false;
668
669    loop {
670        let recv_result = match deadline {
671            Some(d) => {
672                let now = Instant::now();
673                if now >= d {
674                    timed_out = true;
675                    break;
676                }
677                rx.recv_timeout(d - now)
678            }
679            None => rx.recv().map_err(|_| mpsc::RecvTimeoutError::Disconnected),
680        };
681
682        match recv_result {
683            Ok(event) => handler(event),
684            Err(mpsc::RecvTimeoutError::Timeout) => {
685                timed_out = true;
686                break;
687            }
688            Err(mpsc::RecvTimeoutError::Disconnected) => break,
689        }
690    }
691
692    if timed_out {
693        // Take down the whole group first, honoring the optional
694        // SIGTERM grace, so grandchildren holding our pipe fds die
695        // too, then kill+reap the direct child.
696        crate::exec::kill_group_with_grace_sync(&mut group, claude.kill_grace);
697        let _ = child.kill();
698        let _ = child.wait();
699        // Both worker threads can block indefinitely if an orphaned
700        // grandchild inherited our pipe fds and keeps the write end
701        // open (e.g. a `bash` script whose `sleep` subprocess outlives
702        // the SIGKILLed shell). Cap the joins so the timeout error
703        // still returns promptly; any thread that misses the deadline
704        // leaks its JoinHandle, which is acceptable for this edge.
705        let budget = Duration::from_millis(200);
706        let stderr_str = join_with_budget(stderr_thread, budget).unwrap_or_default();
707        let _ = join_with_budget(reader_thread, budget);
708        if !stderr_str.is_empty() {
709            warn!(stderr = %stderr_str, "stderr from timed-out streaming process");
710        }
711        return Err(Error::Timeout {
712            timeout_seconds: claude.timeout.map(|d| d.as_secs()).unwrap_or_default(),
713        });
714    }
715
716    // Normal completion: collect reader result (may carry IO error).
717    let reader_result = reader_thread.join().unwrap_or(Ok(()));
718    if let Err(e) = reader_result {
719        group.kill_now();
720        let _ = child.kill();
721        let _ = child.wait();
722        let _ = stderr_thread.join();
723        return Err(e);
724    }
725
726    let status = child.wait().map_err(|e| Error::Io {
727        message: "failed to wait for claude process".to_string(),
728        source: e,
729        working_dir: claude.working_dir.clone(),
730    })?;
731    group.disarm();
732    let stderr_str = stderr_thread.join().unwrap_or_default();
733    let exit_code = status.code().unwrap_or(-1);
734
735    if !status.success() {
736        return Err(Error::CommandFailed {
737            command: format!("{} {}", claude.binary.display(), command_args.join(" ")),
738            exit_code,
739            stdout: String::new(),
740            stderr: stderr_str,
741            working_dir: claude.working_dir.clone(),
742        });
743    }
744
745    Ok(CommandOutput {
746        stdout: String::new(),
747        stderr: stderr_str,
748        exit_code,
749        success: true,
750    })
751}
752
753/// Join a worker thread with a time budget. Returns `Some(value)` if
754/// the thread finished in time, `None` if the deadline passed first.
755/// A missed deadline leaks the `JoinHandle`; the thread completes
756/// eventually and its value is dropped.
757#[cfg(all(feature = "sync", feature = "json"))]
758fn join_with_budget<T: Send + 'static>(
759    handle: std::thread::JoinHandle<T>,
760    budget: Duration,
761) -> Option<T> {
762    use std::sync::mpsc;
763    use std::thread;
764
765    let (tx, rx) = mpsc::channel::<T>();
766    thread::spawn(move || {
767        if let Ok(v) = handle.join() {
768            let _ = tx.send(v);
769        }
770    });
771    rx.recv_timeout(budget).ok()
772}
773
774#[cfg(all(test, feature = "json"))]
775mod tests {
776    use super::*;
777    use serde_json::json;
778
779    fn parse(v: serde_json::Value) -> StreamEvent {
780        serde_json::from_value(v).expect("valid StreamEvent")
781    }
782
783    fn wrap(inner: serde_json::Value) -> StreamEvent {
784        parse(json!({
785            "type": "stream_event",
786            "event": inner,
787            "session_id": "sess-1",
788            "parent_tool_use_id": null,
789            "uuid": "11111111-1111-1111-1111-111111111111"
790        }))
791    }
792
793    #[test]
794    fn partial_message_text_block_lifecycle() {
795        let start = wrap(json!({
796            "type": "content_block_start",
797            "index": 0,
798            "content_block": { "type": "text", "text": "" }
799        }));
800        assert_eq!(
801            start.partial_message(),
802            Some(PartialMessageEvent::BlockStart {
803                index: 0,
804                block_type: BlockType::Text,
805            })
806        );
807
808        let delta = wrap(json!({
809            "type": "content_block_delta",
810            "index": 0,
811            "delta": { "type": "text_delta", "text": "Hello" }
812        }));
813        assert_eq!(
814            delta.partial_message(),
815            Some(PartialMessageEvent::BlockDelta {
816                index: 0,
817                delta: BlockDelta::Text("Hello".into()),
818            })
819        );
820
821        let stop = wrap(json!({ "type": "content_block_stop", "index": 0 }));
822        assert_eq!(
823            stop.partial_message(),
824            Some(PartialMessageEvent::BlockStop { index: 0 })
825        );
826    }
827
828    #[test]
829    fn partial_message_thinking_block_lifecycle() {
830        let start = wrap(json!({
831            "type": "content_block_start",
832            "index": 1,
833            "content_block": { "type": "thinking", "thinking": "", "signature": "" }
834        }));
835        assert_eq!(
836            start.partial_message(),
837            Some(PartialMessageEvent::BlockStart {
838                index: 1,
839                block_type: BlockType::Thinking,
840            })
841        );
842
843        let delta = wrap(json!({
844            "type": "content_block_delta",
845            "index": 1,
846            "delta": { "type": "thinking_delta", "thinking": "weighing options" }
847        }));
848        assert_eq!(
849            delta.partial_message(),
850            Some(PartialMessageEvent::BlockDelta {
851                index: 1,
852                delta: BlockDelta::Thinking("weighing options".into()),
853            })
854        );
855
856        let stop = wrap(json!({ "type": "content_block_stop", "index": 1 }));
857        assert_eq!(
858            stop.partial_message(),
859            Some(PartialMessageEvent::BlockStop { index: 1 })
860        );
861    }
862
863    #[test]
864    fn partial_message_tool_use_block_carries_id_and_name() {
865        let start = wrap(json!({
866            "type": "content_block_start",
867            "index": 2,
868            "content_block": {
869                "type": "tool_use",
870                "id": "toolu_abc",
871                "name": "Bash",
872                "input": {}
873            }
874        }));
875        assert_eq!(
876            start.partial_message(),
877            Some(PartialMessageEvent::BlockStart {
878                index: 2,
879                block_type: BlockType::ToolUse {
880                    id: "toolu_abc".into(),
881                    name: "Bash".into(),
882                },
883            })
884        );
885
886        let delta = wrap(json!({
887            "type": "content_block_delta",
888            "index": 2,
889            "delta": { "type": "input_json_delta", "partial_json": "{\"cmd\":" }
890        }));
891        assert_eq!(
892            delta.partial_message(),
893            Some(PartialMessageEvent::BlockDelta {
894                index: 2,
895                delta: BlockDelta::InputJson("{\"cmd\":".into()),
896            })
897        );
898    }
899
900    #[test]
901    fn partial_message_unknown_kinds_fall_through_to_other() {
902        let unknown_block = wrap(json!({
903            "type": "content_block_start",
904            "index": 3,
905            "content_block": { "type": "redacted_thinking", "data": "..." }
906        }));
907        assert_eq!(
908            unknown_block.partial_message(),
909            Some(PartialMessageEvent::BlockStart {
910                index: 3,
911                block_type: BlockType::Other("redacted_thinking".into()),
912            })
913        );
914
915        let unknown_delta = wrap(json!({
916            "type": "content_block_delta",
917            "index": 3,
918            "delta": { "type": "signature_delta", "signature": "sig" }
919        }));
920        assert_eq!(
921            unknown_delta.partial_message(),
922            Some(PartialMessageEvent::BlockDelta {
923                index: 3,
924                delta: BlockDelta::Other,
925            })
926        );
927    }
928
929    #[test]
930    fn partial_message_returns_none_for_non_partial_events() {
931        let result = parse(json!({
932            "type": "result",
933            "result": "done",
934            "session_id": "sess-1",
935            "total_cost_usd": 0.01
936        }));
937        assert!(result.partial_message().is_none());
938
939        let assistant = parse(json!({
940            "type": "assistant",
941            "message": { "role": "assistant", "content": [] },
942            "session_id": "sess-1"
943        }));
944        assert!(assistant.partial_message().is_none());
945
946        let message_start = wrap(json!({
947            "type": "message_start",
948            "message": { "id": "msg_1", "role": "assistant", "content": [] }
949        }));
950        assert!(message_start.partial_message().is_none());
951    }
952
953    #[test]
954    fn partial_message_accepts_unwrapped_event() {
955        let raw = parse(json!({
956            "type": "content_block_delta",
957            "index": 0,
958            "delta": { "type": "text_delta", "text": "hi" }
959        }));
960        assert_eq!(
961            raw.partial_message(),
962            Some(PartialMessageEvent::BlockDelta {
963                index: 0,
964                delta: BlockDelta::Text("hi".into()),
965            })
966        );
967    }
968}