claude_codes/
client_async.rs

1//! Asynchronous client for Claude communication
2
3use crate::cli::ClaudeCliBuilder;
4use crate::error::{Error, Result};
5use crate::io::{
6    ClaudeInput, ClaudeOutput, ContentBlock, ControlRequestMessage, ControlResponse,
7    ControlResponseMessage,
8};
9use crate::protocol::Protocol;
10use log::{debug, error, info, warn};
11use serde::{Deserialize, Serialize};
12use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufReader as AsyncBufReader};
13use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout};
14use uuid::Uuid;
15
16/// Asynchronous client for communicating with Claude
17pub struct AsyncClient {
18    child: Child,
19    stdin: ChildStdin,
20    stdout: BufReader<ChildStdout>,
21    stderr: Option<BufReader<ChildStderr>>,
22    session_uuid: Option<Uuid>,
23    /// Whether tool approval protocol has been initialized
24    tool_approval_enabled: bool,
25}
26
27/// Buffer size for reading Claude's stdout (10MB).
28const STDOUT_BUFFER_SIZE: usize = 10 * 1024 * 1024;
29
30impl AsyncClient {
31    /// Create a new async client from a tokio Child process
32    pub fn new(mut child: Child) -> Result<Self> {
33        let stdin = child
34            .stdin
35            .take()
36            .ok_or_else(|| Error::Io(std::io::Error::other("Failed to get stdin handle")))?;
37
38        let stdout = BufReader::with_capacity(
39            STDOUT_BUFFER_SIZE,
40            child
41                .stdout
42                .take()
43                .ok_or_else(|| Error::Io(std::io::Error::other("Failed to get stdout handle")))?,
44        );
45
46        let stderr = child.stderr.take().map(BufReader::new);
47
48        Ok(Self {
49            child,
50            stdin,
51            stdout,
52            stderr,
53            session_uuid: None,
54            tool_approval_enabled: false,
55        })
56    }
57
58    /// Create a client with default settings (using logic from start_claude)
59    pub async fn with_defaults() -> Result<Self> {
60        // Check Claude version (only warns once per session)
61        // NOTE: The claude-codes API is in high flux. If you wish to work around
62        // this version check, you can use AsyncClient::new() directly with:
63        //   let child = ClaudeCliBuilder::new().model("sonnet").spawn().await?;
64        //   AsyncClient::new(child)
65        crate::version::check_claude_version_async().await?;
66        Self::with_model("sonnet").await
67    }
68
69    /// Create a client with a specific model
70    pub async fn with_model(model: &str) -> Result<Self> {
71        let child = ClaudeCliBuilder::new().model(model).spawn().await?;
72
73        info!("Started Claude process with model: {}", model);
74        Self::new(child)
75    }
76
77    /// Create a client from a custom builder
78    pub async fn from_builder(builder: ClaudeCliBuilder) -> Result<Self> {
79        let child = builder.spawn().await?;
80        info!("Started Claude process from custom builder");
81        Self::new(child)
82    }
83
84    /// Resume a previous session by UUID
85    /// This creates a new client that resumes an existing session
86    pub async fn resume_session(session_uuid: Uuid) -> Result<Self> {
87        let child = ClaudeCliBuilder::new()
88            .resume(Some(session_uuid.to_string()))
89            .spawn()
90            .await?;
91
92        info!("Resuming Claude session with UUID: {}", session_uuid);
93        let mut client = Self::new(child)?;
94        // Pre-populate the session UUID since we're resuming
95        client.session_uuid = Some(session_uuid);
96        Ok(client)
97    }
98
99    /// Resume a previous session with a specific model
100    pub async fn resume_session_with_model(session_uuid: Uuid, model: &str) -> Result<Self> {
101        let child = ClaudeCliBuilder::new()
102            .model(model)
103            .resume(Some(session_uuid.to_string()))
104            .spawn()
105            .await?;
106
107        info!(
108            "Resuming Claude session with UUID: {} and model: {}",
109            session_uuid, model
110        );
111        let mut client = Self::new(child)?;
112        // Pre-populate the session UUID since we're resuming
113        client.session_uuid = Some(session_uuid);
114        Ok(client)
115    }
116
117    /// Send a query and collect all responses until Result message
118    /// This is the simplified version that collects all responses
119    pub async fn query(&mut self, text: &str) -> Result<Vec<ClaudeOutput>> {
120        let session_id = Uuid::new_v4();
121        self.query_with_session(text, session_id).await
122    }
123
124    /// Send a query with a custom session ID and collect all responses
125    pub async fn query_with_session(
126        &mut self,
127        text: &str,
128        session_id: Uuid,
129    ) -> Result<Vec<ClaudeOutput>> {
130        // Send the query
131        let input = ClaudeInput::user_message(text, session_id);
132        self.send(&input).await?;
133
134        // Collect responses until we get a Result message
135        let mut responses = Vec::new();
136
137        loop {
138            let output = self.receive().await?;
139            let is_result = matches!(&output, ClaudeOutput::Result(_));
140            responses.push(output);
141
142            if is_result {
143                break;
144            }
145        }
146
147        Ok(responses)
148    }
149
150    /// Send a query and return an async iterator over responses
151    /// Returns a stream that yields ClaudeOutput until Result message is received
152    pub async fn query_stream(&mut self, text: &str) -> Result<ResponseStream<'_>> {
153        let session_id = Uuid::new_v4();
154        self.query_stream_with_session(text, session_id).await
155    }
156
157    /// Send a query with session ID and return an async iterator over responses
158    pub async fn query_stream_with_session(
159        &mut self,
160        text: &str,
161        session_id: Uuid,
162    ) -> Result<ResponseStream<'_>> {
163        // Send the query first
164        let input = ClaudeInput::user_message(text, session_id);
165        self.send(&input).await?;
166
167        // Return a stream that will read responses
168        Ok(ResponseStream {
169            client: self,
170            finished: false,
171        })
172    }
173
174    /// Send a ClaudeInput directly
175    pub async fn send(&mut self, input: &ClaudeInput) -> Result<()> {
176        let json_line = Protocol::serialize(input)?;
177        debug!("[OUTGOING] Sending JSON to Claude: {}", json_line.trim());
178
179        self.stdin
180            .write_all(json_line.as_bytes())
181            .await
182            .map_err(Error::Io)?;
183
184        self.stdin.flush().await.map_err(Error::Io)?;
185        Ok(())
186    }
187
188    /// Send an interrupt to gracefully stop the current response.
189    ///
190    /// This writes a `control_request` with subtype `interrupt` to stdin,
191    /// telling Claude to stop without killing the session. Returns the
192    /// generated `request_id`; the CLI acknowledges with a
193    /// `control_response` carrying the same id and ends the turn with a
194    /// `result` message.
195    pub async fn interrupt(&mut self) -> Result<String> {
196        let request_id = format!("interrupt-{}", Uuid::new_v4());
197        self.send(&ClaudeInput::interrupt(&request_id)).await?;
198        Ok(request_id)
199    }
200
201    /// Merge settings into the live session (`apply_flag_settings` control
202    /// request), e.g. change the effort level or toggle ultracode without
203    /// restarting.
204    ///
205    /// Returns the generated `request_id`; the CLI answers with a
206    /// `control_response` carrying the same id (an error response if it
207    /// rejects a value). The new effort takes effect on the next request.
208    pub async fn apply_flag_settings(
209        &mut self,
210        settings: crate::io::ApplyFlagSettingsRequest,
211    ) -> Result<String> {
212        let request_id = format!("apply-flag-settings-{}", Uuid::new_v4());
213        self.send(&ClaudeInput::apply_flag_settings(&request_id, settings))
214            .await?;
215        Ok(request_id)
216    }
217
218    /// Change the session's effort level mid-session. Shorthand for
219    /// [`apply_flag_settings`](Self::apply_flag_settings) with only
220    /// `effortLevel` set.
221    pub async fn set_effort(&mut self, level: crate::EffortLevel) -> Result<String> {
222        self.apply_flag_settings(crate::io::ApplyFlagSettingsRequest::new().effort_level(level))
223            .await
224    }
225
226    /// Receive a single response from Claude.
227    ///
228    /// # Important: Polling Frequency
229    ///
230    /// This method should be polled frequently to prevent the OS pipe buffer from
231    /// filling up. Claude can emit very large JSON messages (hundreds of KB), and
232    /// if the pipe buffer overflows, data may be truncated.
233    ///
234    /// In a `tokio::select!` loop with other async operations, ensure `receive()`
235    /// is given priority or called frequently. For high-throughput scenarios,
236    /// consider spawning a dedicated task to drain stdout into an unbounded channel.
237    ///
238    /// # Returns
239    ///
240    /// - `Ok(ClaudeOutput)` - A parsed message from Claude
241    /// - `Err(Error::ConnectionClosed)` - Claude process has exited
242    /// - `Err(Error::Deserialization)` - Failed to parse the message
243    pub async fn receive(&mut self) -> Result<ClaudeOutput> {
244        let trimmed = self.read_frame_line().await?;
245        debug!("[INCOMING] Received JSON from Claude: {}", trimmed);
246
247        // Use the parse_json_tolerant method which handles ANSI escape codes
248        match ClaudeOutput::parse_json_tolerant(&trimmed) {
249            Ok(output) => {
250                debug!("[INCOMING] Parsed output type: {}", output.message_type());
251
252                // Capture UUID from first response if not already set
253                if self.session_uuid.is_none() {
254                    if let ClaudeOutput::Assistant(ref msg) = output {
255                        if let Some(ref uuid_str) = msg.uuid {
256                            if let Ok(uuid) = Uuid::parse_str(uuid_str) {
257                                debug!("[INCOMING] Captured session UUID: {}", uuid);
258                                self.session_uuid = Some(uuid);
259                            }
260                        }
261                    } else if let ClaudeOutput::Result(ref msg) = output {
262                        if let Some(ref uuid_str) = msg.uuid {
263                            if let Ok(uuid) = Uuid::parse_str(uuid_str) {
264                                debug!("[INCOMING] Captured session UUID: {}", uuid);
265                                self.session_uuid = Some(uuid);
266                            }
267                        }
268                    }
269                }
270
271                Ok(output)
272            }
273            Err(parse_error) => {
274                warn!("[INCOMING] Failed to deserialize message from Claude CLI. Please report this at https://github.com/meawoppl/rust-claude-codes/issues with the raw message below.");
275                warn!("[INCOMING] Parse error: {}", parse_error.error_message);
276                warn!("[INCOMING] Raw message: {}", trimmed);
277                Err(parse_error.into())
278            }
279        }
280    }
281
282    /// Read the next non-empty line from Claude's stdout, trimmed.
283    ///
284    /// Returns `Err(Error::ConnectionClosed)` at EOF. Shared by [`receive`] and
285    /// [`receive_raw`].
286    ///
287    /// [`receive`]: Self::receive
288    /// [`receive_raw`]: Self::receive_raw
289    async fn read_frame_line(&mut self) -> Result<String> {
290        let mut line = String::new();
291        loop {
292            line.clear();
293            let bytes_read = self.stdout.read_line(&mut line).await.map_err(Error::Io)?;
294            if bytes_read == 0 {
295                return Err(Error::ConnectionClosed);
296            }
297            let trimmed = line.trim();
298            if trimmed.is_empty() {
299                continue;
300            }
301            return Ok(trimmed.to_string());
302        }
303    }
304
305    /// Receive the next frame as a raw `serde_json::Value`, before it is mapped
306    /// into a typed [`ClaudeOutput`].
307    ///
308    /// Useful for auditing wire fidelity: pair it with
309    /// [`audit_frame`](crate::io::audit_frame) to confirm the typed model
310    /// captures every field the CLI emitted. Applies the same leading-prefix
311    /// tolerance as [`receive`](Self::receive).
312    pub async fn receive_raw(&mut self) -> Result<serde_json::Value> {
313        let trimmed = self.read_frame_line().await?;
314        match serde_json::from_str::<serde_json::Value>(&trimmed) {
315            Ok(value) => Ok(value),
316            Err(e) => match trimmed.find('{') {
317                Some(start) => Ok(serde_json::from_str::<serde_json::Value>(&trimmed[start..])
318                    .map_err(|_| Error::Json(e))?),
319                None => Err(Error::Json(e)),
320            },
321        }
322    }
323
324    /// Check if the Claude process is still running
325    pub fn is_alive(&mut self) -> bool {
326        self.child.try_wait().ok().flatten().is_none()
327    }
328
329    /// Gracefully shutdown the client
330    pub async fn shutdown(mut self) -> Result<()> {
331        info!("Shutting down Claude process...");
332        self.child.kill().await.map_err(Error::Io)?;
333        Ok(())
334    }
335
336    /// Get the process ID
337    pub fn pid(&self) -> Option<u32> {
338        self.child.id()
339    }
340
341    /// Take the stderr reader (can only be called once)
342    pub fn take_stderr(&mut self) -> Option<BufReader<ChildStderr>> {
343        self.stderr.take()
344    }
345
346    /// Get the session UUID if available
347    /// Returns an error if no response has been received yet
348    pub fn session_uuid(&self) -> Result<Uuid> {
349        self.session_uuid.ok_or(Error::SessionNotInitialized)
350    }
351
352    /// Test if the Claude connection is working by sending a ping message
353    /// Returns true if Claude responds with "pong", false otherwise
354    pub async fn ping(&mut self) -> bool {
355        // Send a simple ping request
356        let ping_input = ClaudeInput::user_message(
357            "ping - respond with just the word 'pong' and nothing else",
358            self.session_uuid.unwrap_or_else(Uuid::new_v4),
359        );
360
361        // Try to send the ping
362        if let Err(e) = self.send(&ping_input).await {
363            debug!("Ping failed to send: {}", e);
364            return false;
365        }
366
367        // Try to receive responses until we get a result or error
368        let mut found_pong = false;
369        let mut message_count = 0;
370        const MAX_MESSAGES: usize = 10;
371
372        loop {
373            match self.receive().await {
374                Ok(output) => {
375                    message_count += 1;
376
377                    // Check if it's an assistant message containing "pong"
378                    if let ClaudeOutput::Assistant(msg) = &output {
379                        for content in &msg.message.content {
380                            if let ContentBlock::Text(text) = content {
381                                if text.text.to_lowercase().contains("pong") {
382                                    found_pong = true;
383                                }
384                            }
385                        }
386                    }
387
388                    // Stop on result message
389                    if matches!(output, ClaudeOutput::Result(_)) {
390                        break;
391                    }
392
393                    // Safety limit
394                    if message_count >= MAX_MESSAGES {
395                        debug!("Ping exceeded message limit");
396                        break;
397                    }
398                }
399                Err(e) => {
400                    debug!("Ping failed to receive response: {}", e);
401                    break;
402                }
403            }
404        }
405
406        found_pong
407    }
408
409    // =========================================================================
410    // Tool Approval Protocol
411    // =========================================================================
412
413    /// Enable the tool approval protocol by performing the initialization handshake.
414    ///
415    /// After calling this method, the CLI will send `ControlRequest` messages when
416    /// Claude wants to use a tool. You must handle these by calling
417    /// `send_control_response()` with an appropriate response.
418    ///
419    /// **Important**: The client must have been created with
420    /// `ClaudeCliBuilder::permission_prompt_tool("stdio")` for this to work.
421    ///
422    /// # Example
423    ///
424    /// ```no_run
425    /// use claude_codes::{AsyncClient, ClaudeCliBuilder, ClaudeOutput, ControlRequestPayload};
426    ///
427    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
428    /// let child = ClaudeCliBuilder::new()
429    ///     .model("sonnet")
430    ///     .permission_prompt_tool("stdio")
431    ///     .spawn()
432    ///     .await?;
433    ///
434    /// let mut client = AsyncClient::new(child)?;
435    /// client.enable_tool_approval().await?;
436    ///
437    /// // Now when you receive messages, you may get ControlRequest messages
438    /// // that need responses
439    /// # Ok(())
440    /// # }
441    /// ```
442    pub async fn enable_tool_approval(&mut self) -> Result<()> {
443        if self.tool_approval_enabled {
444            debug!("[TOOL_APPROVAL] Already enabled, skipping initialization");
445            return Ok(());
446        }
447
448        let request_id = format!("init-{}", Uuid::new_v4());
449        let init_request = ControlRequestMessage::initialize(&request_id);
450
451        debug!("[TOOL_APPROVAL] Sending initialization handshake");
452        let json_line = Protocol::serialize(&init_request)?;
453        self.stdin
454            .write_all(json_line.as_bytes())
455            .await
456            .map_err(Error::Io)?;
457        self.stdin.flush().await.map_err(Error::Io)?;
458
459        // Wait for the initialization response
460        loop {
461            let mut line = String::new();
462            let bytes_read = self.stdout.read_line(&mut line).await.map_err(Error::Io)?;
463
464            if bytes_read == 0 {
465                return Err(Error::ConnectionClosed);
466            }
467
468            let trimmed = line.trim();
469            if trimmed.is_empty() {
470                continue;
471            }
472
473            debug!("[TOOL_APPROVAL] Received: {}", trimmed);
474
475            // Try to parse as ClaudeOutput
476            match ClaudeOutput::parse_json_tolerant(trimmed) {
477                Ok(ClaudeOutput::ControlResponse(resp)) => {
478                    use crate::io::ControlResponsePayload;
479                    match &resp.response {
480                        ControlResponsePayload::Success {
481                            request_id: rid, ..
482                        } if rid == &request_id => {
483                            debug!("[TOOL_APPROVAL] Initialization successful");
484                            self.tool_approval_enabled = true;
485                            return Ok(());
486                        }
487                        ControlResponsePayload::Error { error, .. } => {
488                            return Err(Error::Protocol(format!(
489                                "Tool approval initialization failed: {}",
490                                error
491                            )));
492                        }
493                        _ => {
494                            // Different request_id, keep waiting
495                            continue;
496                        }
497                    }
498                }
499                Ok(_) => {
500                    // Got a different message type (system, etc.), keep waiting
501                    continue;
502                }
503                Err(e) => {
504                    return Err(e.into());
505                }
506            }
507        }
508    }
509
510    /// Send a control response back to the CLI.
511    ///
512    /// Use this to respond to `ControlRequest` messages received during tool approval.
513    /// The easiest way to create responses is using the helper methods on
514    /// `ToolPermissionRequest`:
515    ///
516    /// # Example
517    ///
518    /// ```no_run
519    /// use claude_codes::{AsyncClient, ClaudeOutput, ControlRequestPayload};
520    ///
521    /// # async fn example(client: &mut AsyncClient) -> Result<(), Box<dyn std::error::Error>> {
522    /// # let output = client.receive().await?;
523    /// if let ClaudeOutput::ControlRequest(req) = output {
524    ///     if let ControlRequestPayload::CanUseTool(perm_req) = &req.request {
525    ///         // Use the ergonomic helpers on ToolPermissionRequest
526    ///         let response = if perm_req.tool_name == "Bash" {
527    ///             perm_req.deny("Bash commands not allowed", &req.request_id)
528    ///         } else {
529    ///             perm_req.allow(&req.request_id)
530    ///         };
531    ///         client.send_control_response(response).await?;
532    ///     }
533    /// }
534    /// # Ok(())
535    /// # }
536    /// ```
537    pub async fn send_control_response(&mut self, response: ControlResponse) -> Result<()> {
538        let message: ControlResponseMessage = response.into();
539        let json_line = Protocol::serialize(&message)?;
540        debug!(
541            "[TOOL_APPROVAL] Sending control response: {}",
542            json_line.trim()
543        );
544
545        self.stdin
546            .write_all(json_line.as_bytes())
547            .await
548            .map_err(Error::Io)?;
549        self.stdin.flush().await.map_err(Error::Io)?;
550        Ok(())
551    }
552
553    /// Check if tool approval protocol is enabled
554    pub fn is_tool_approval_enabled(&self) -> bool {
555        self.tool_approval_enabled
556    }
557}
558
559/// A response stream that yields ClaudeOutput messages
560/// Holds a reference to the client to read from
561pub struct ResponseStream<'a> {
562    client: &'a mut AsyncClient,
563    finished: bool,
564}
565
566impl ResponseStream<'_> {
567    /// Convert to a vector by collecting all responses
568    pub async fn collect(mut self) -> Result<Vec<ClaudeOutput>> {
569        let mut responses = Vec::new();
570
571        while !self.finished {
572            let output = self.client.receive().await?;
573            let is_result = matches!(&output, ClaudeOutput::Result(_));
574            responses.push(output);
575
576            if is_result {
577                self.finished = true;
578                break;
579            }
580        }
581
582        Ok(responses)
583    }
584
585    /// Get the next response
586    pub async fn next(&mut self) -> Option<Result<ClaudeOutput>> {
587        if self.finished {
588            return None;
589        }
590
591        match self.client.receive().await {
592            Ok(output) => {
593                if matches!(&output, ClaudeOutput::Result(_)) {
594                    self.finished = true;
595                }
596                Some(Ok(output))
597            }
598            Err(e) => {
599                self.finished = true;
600                Some(Err(e))
601            }
602        }
603    }
604}
605
606impl Drop for AsyncClient {
607    fn drop(&mut self) {
608        if self.is_alive() {
609            // Try to kill the process
610            if let Err(e) = self.child.start_kill() {
611                error!("Failed to kill Claude process on drop: {}", e);
612            }
613        }
614    }
615}
616
617// Protocol extension methods for asynchronous I/O
618impl Protocol {
619    /// Write a message to an async writer
620    pub async fn write_async<W: AsyncWriteExt + Unpin, T: Serialize>(
621        writer: &mut W,
622        message: &T,
623    ) -> Result<()> {
624        let line = Self::serialize(message)?;
625        debug!("[PROTOCOL] Sending async: {}", line.trim());
626        writer.write_all(line.as_bytes()).await?;
627        writer.flush().await?;
628        Ok(())
629    }
630
631    /// Read a message from an async reader
632    pub async fn read_async<R: AsyncBufReadExt + Unpin, T: for<'de> Deserialize<'de>>(
633        reader: &mut R,
634    ) -> Result<T> {
635        let mut line = String::new();
636        let bytes_read = reader.read_line(&mut line).await?;
637        if bytes_read == 0 {
638            return Err(Error::ConnectionClosed);
639        }
640        debug!("[PROTOCOL] Received async: {}", line.trim());
641        Self::deserialize(&line)
642    }
643}
644
645/// Async stream processor for handling continuous message streams
646pub struct AsyncStreamProcessor<R> {
647    reader: AsyncBufReader<R>,
648}
649
650impl<R: tokio::io::AsyncRead + Unpin> AsyncStreamProcessor<R> {
651    /// Create a new async stream processor
652    pub fn new(reader: R) -> Self {
653        Self {
654            reader: AsyncBufReader::new(reader),
655        }
656    }
657
658    /// Process the next message from the stream
659    pub async fn next_message<T: for<'de> Deserialize<'de>>(&mut self) -> Result<T> {
660        Protocol::read_async(&mut self.reader).await
661    }
662
663    /// Process all messages in the stream
664    pub async fn process_all<T, F, Fut>(&mut self, mut handler: F) -> Result<()>
665    where
666        T: for<'de> Deserialize<'de>,
667        F: FnMut(T) -> Fut,
668        Fut: std::future::Future<Output = Result<()>>,
669    {
670        loop {
671            match self.next_message().await {
672                Ok(message) => handler(message).await?,
673                Err(Error::ConnectionClosed) => break,
674                Err(e) => return Err(e),
675            }
676        }
677        Ok(())
678    }
679}
claude_codes/client_async.rs

claude_codes/
client_async.rs
