1#[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#[cfg(feature = "json")]
30#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
31pub struct StreamEvent {
32 #[serde(flatten)]
34 pub data: serde_json::Value,
35}
36
37#[cfg(feature = "json")]
38impl StreamEvent {
39 pub fn event_type(&self) -> Option<&str> {
41 self.data.get("type").and_then(|v| v.as_str())
42 }
43
44 pub fn role(&self) -> Option<&str> {
46 self.data.get("role").and_then(|v| v.as_str())
47 }
48
49 pub fn is_result(&self) -> bool {
51 self.event_type() == Some("result")
52 }
53
54 pub fn result_text(&self) -> Option<&str> {
56 self.data.get("result").and_then(|v| v.as_str())
57 }
58
59 pub fn session_id(&self) -> Option<&str> {
61 self.data.get("session_id").and_then(|v| v.as_str())
62 }
63
64 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 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#[cfg(feature = "json")]
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub enum PartialMessageEvent {
149 BlockStart {
151 index: u32,
153 block_type: BlockType,
155 },
156 BlockDelta {
158 index: u32,
162 delta: BlockDelta,
164 },
165 BlockStop {
167 index: u32,
169 },
170}
171
172#[cfg(feature = "json")]
178#[derive(Debug, Clone, PartialEq, Eq)]
179pub enum BlockType {
180 Text,
182 Thinking,
184 ToolUse {
186 id: String,
188 name: String,
190 },
191 Other(String),
193}
194
195#[cfg(feature = "json")]
202#[derive(Debug, Clone, PartialEq, Eq)]
203pub enum BlockDelta {
204 Text(String),
206 Thinking(String),
208 InputJson(String),
211 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#[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#[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 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 .kill_on_drop(true);
379 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 let drain = drain_stderr(&mut stderr);
406 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 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 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(), 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#[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 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 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 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 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 crate::exec::kill_group_with_grace_sync(&mut group, claude.kill_grace);
697 let _ = child.kill();
698 let _ = child.wait();
699 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 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#[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}