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