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    /// Receive a single response from Claude.
202    ///
203    /// # Important: Polling Frequency
204    ///
205    /// This method should be polled frequently to prevent the OS pipe buffer from
206    /// filling up. Claude can emit very large JSON messages (hundreds of KB), and
207    /// if the pipe buffer overflows, data may be truncated.
208    ///
209    /// In a `tokio::select!` loop with other async operations, ensure `receive()`
210    /// is given priority or called frequently. For high-throughput scenarios,
211    /// consider spawning a dedicated task to drain stdout into an unbounded channel.
212    ///
213    /// # Returns
214    ///
215    /// - `Ok(ClaudeOutput)` - A parsed message from Claude
216    /// - `Err(Error::ConnectionClosed)` - Claude process has exited
217    /// - `Err(Error::Deserialization)` - Failed to parse the message
218    pub async fn receive(&mut self) -> Result<ClaudeOutput> {
219        let trimmed = self.read_frame_line().await?;
220        debug!("[INCOMING] Received JSON from Claude: {}", trimmed);
221
222        // Use the parse_json_tolerant method which handles ANSI escape codes
223        match ClaudeOutput::parse_json_tolerant(&trimmed) {
224            Ok(output) => {
225                debug!("[INCOMING] Parsed output type: {}", output.message_type());
226
227                // Capture UUID from first response if not already set
228                if self.session_uuid.is_none() {
229                    if let ClaudeOutput::Assistant(ref msg) = output {
230                        if let Some(ref uuid_str) = msg.uuid {
231                            if let Ok(uuid) = Uuid::parse_str(uuid_str) {
232                                debug!("[INCOMING] Captured session UUID: {}", uuid);
233                                self.session_uuid = Some(uuid);
234                            }
235                        }
236                    } else if let ClaudeOutput::Result(ref msg) = output {
237                        if let Some(ref uuid_str) = msg.uuid {
238                            if let Ok(uuid) = Uuid::parse_str(uuid_str) {
239                                debug!("[INCOMING] Captured session UUID: {}", uuid);
240                                self.session_uuid = Some(uuid);
241                            }
242                        }
243                    }
244                }
245
246                Ok(output)
247            }
248            Err(parse_error) => {
249                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.");
250                warn!("[INCOMING] Parse error: {}", parse_error.error_message);
251                warn!("[INCOMING] Raw message: {}", trimmed);
252                Err(parse_error.into())
253            }
254        }
255    }
256
257    /// Read the next non-empty line from Claude's stdout, trimmed.
258    ///
259    /// Returns `Err(Error::ConnectionClosed)` at EOF. Shared by [`receive`] and
260    /// [`receive_raw`].
261    ///
262    /// [`receive`]: Self::receive
263    /// [`receive_raw`]: Self::receive_raw
264    async fn read_frame_line(&mut self) -> Result<String> {
265        let mut line = String::new();
266        loop {
267            line.clear();
268            let bytes_read = self.stdout.read_line(&mut line).await.map_err(Error::Io)?;
269            if bytes_read == 0 {
270                return Err(Error::ConnectionClosed);
271            }
272            let trimmed = line.trim();
273            if trimmed.is_empty() {
274                continue;
275            }
276            return Ok(trimmed.to_string());
277        }
278    }
279
280    /// Receive the next frame as a raw `serde_json::Value`, before it is mapped
281    /// into a typed [`ClaudeOutput`].
282    ///
283    /// Useful for auditing wire fidelity: pair it with
284    /// [`audit_frame`](crate::io::audit_frame) to confirm the typed model
285    /// captures every field the CLI emitted. Applies the same leading-prefix
286    /// tolerance as [`receive`](Self::receive).
287    pub async fn receive_raw(&mut self) -> Result<serde_json::Value> {
288        let trimmed = self.read_frame_line().await?;
289        match serde_json::from_str::<serde_json::Value>(&trimmed) {
290            Ok(value) => Ok(value),
291            Err(e) => match trimmed.find('{') {
292                Some(start) => Ok(serde_json::from_str::<serde_json::Value>(&trimmed[start..])
293                    .map_err(|_| Error::Json(e))?),
294                None => Err(Error::Json(e)),
295            },
296        }
297    }
298
299    /// Check if the Claude process is still running
300    pub fn is_alive(&mut self) -> bool {
301        self.child.try_wait().ok().flatten().is_none()
302    }
303
304    /// Gracefully shutdown the client
305    pub async fn shutdown(mut self) -> Result<()> {
306        info!("Shutting down Claude process...");
307        self.child.kill().await.map_err(Error::Io)?;
308        Ok(())
309    }
310
311    /// Get the process ID
312    pub fn pid(&self) -> Option<u32> {
313        self.child.id()
314    }
315
316    /// Take the stderr reader (can only be called once)
317    pub fn take_stderr(&mut self) -> Option<BufReader<ChildStderr>> {
318        self.stderr.take()
319    }
320
321    /// Get the session UUID if available
322    /// Returns an error if no response has been received yet
323    pub fn session_uuid(&self) -> Result<Uuid> {
324        self.session_uuid.ok_or(Error::SessionNotInitialized)
325    }
326
327    /// Test if the Claude connection is working by sending a ping message
328    /// Returns true if Claude responds with "pong", false otherwise
329    pub async fn ping(&mut self) -> bool {
330        // Send a simple ping request
331        let ping_input = ClaudeInput::user_message(
332            "ping - respond with just the word 'pong' and nothing else",
333            self.session_uuid.unwrap_or_else(Uuid::new_v4),
334        );
335
336        // Try to send the ping
337        if let Err(e) = self.send(&ping_input).await {
338            debug!("Ping failed to send: {}", e);
339            return false;
340        }
341
342        // Try to receive responses until we get a result or error
343        let mut found_pong = false;
344        let mut message_count = 0;
345        const MAX_MESSAGES: usize = 10;
346
347        loop {
348            match self.receive().await {
349                Ok(output) => {
350                    message_count += 1;
351
352                    // Check if it's an assistant message containing "pong"
353                    if let ClaudeOutput::Assistant(msg) = &output {
354                        for content in &msg.message.content {
355                            if let ContentBlock::Text(text) = content {
356                                if text.text.to_lowercase().contains("pong") {
357                                    found_pong = true;
358                                }
359                            }
360                        }
361                    }
362
363                    // Stop on result message
364                    if matches!(output, ClaudeOutput::Result(_)) {
365                        break;
366                    }
367
368                    // Safety limit
369                    if message_count >= MAX_MESSAGES {
370                        debug!("Ping exceeded message limit");
371                        break;
372                    }
373                }
374                Err(e) => {
375                    debug!("Ping failed to receive response: {}", e);
376                    break;
377                }
378            }
379        }
380
381        found_pong
382    }
383
384    // =========================================================================
385    // Tool Approval Protocol
386    // =========================================================================
387
388    /// Enable the tool approval protocol by performing the initialization handshake.
389    ///
390    /// After calling this method, the CLI will send `ControlRequest` messages when
391    /// Claude wants to use a tool. You must handle these by calling
392    /// `send_control_response()` with an appropriate response.
393    ///
394    /// **Important**: The client must have been created with
395    /// `ClaudeCliBuilder::permission_prompt_tool("stdio")` for this to work.
396    ///
397    /// # Example
398    ///
399    /// ```no_run
400    /// use claude_codes::{AsyncClient, ClaudeCliBuilder, ClaudeOutput, ControlRequestPayload};
401    ///
402    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
403    /// let child = ClaudeCliBuilder::new()
404    ///     .model("sonnet")
405    ///     .permission_prompt_tool("stdio")
406    ///     .spawn()
407    ///     .await?;
408    ///
409    /// let mut client = AsyncClient::new(child)?;
410    /// client.enable_tool_approval().await?;
411    ///
412    /// // Now when you receive messages, you may get ControlRequest messages
413    /// // that need responses
414    /// # Ok(())
415    /// # }
416    /// ```
417    pub async fn enable_tool_approval(&mut self) -> Result<()> {
418        if self.tool_approval_enabled {
419            debug!("[TOOL_APPROVAL] Already enabled, skipping initialization");
420            return Ok(());
421        }
422
423        let request_id = format!("init-{}", Uuid::new_v4());
424        let init_request = ControlRequestMessage::initialize(&request_id);
425
426        debug!("[TOOL_APPROVAL] Sending initialization handshake");
427        let json_line = Protocol::serialize(&init_request)?;
428        self.stdin
429            .write_all(json_line.as_bytes())
430            .await
431            .map_err(Error::Io)?;
432        self.stdin.flush().await.map_err(Error::Io)?;
433
434        // Wait for the initialization response
435        loop {
436            let mut line = String::new();
437            let bytes_read = self.stdout.read_line(&mut line).await.map_err(Error::Io)?;
438
439            if bytes_read == 0 {
440                return Err(Error::ConnectionClosed);
441            }
442
443            let trimmed = line.trim();
444            if trimmed.is_empty() {
445                continue;
446            }
447
448            debug!("[TOOL_APPROVAL] Received: {}", trimmed);
449
450            // Try to parse as ClaudeOutput
451            match ClaudeOutput::parse_json_tolerant(trimmed) {
452                Ok(ClaudeOutput::ControlResponse(resp)) => {
453                    use crate::io::ControlResponsePayload;
454                    match &resp.response {
455                        ControlResponsePayload::Success {
456                            request_id: rid, ..
457                        } if rid == &request_id => {
458                            debug!("[TOOL_APPROVAL] Initialization successful");
459                            self.tool_approval_enabled = true;
460                            return Ok(());
461                        }
462                        ControlResponsePayload::Error { error, .. } => {
463                            return Err(Error::Protocol(format!(
464                                "Tool approval initialization failed: {}",
465                                error
466                            )));
467                        }
468                        _ => {
469                            // Different request_id, keep waiting
470                            continue;
471                        }
472                    }
473                }
474                Ok(_) => {
475                    // Got a different message type (system, etc.), keep waiting
476                    continue;
477                }
478                Err(e) => {
479                    return Err(e.into());
480                }
481            }
482        }
483    }
484
485    /// Send a control response back to the CLI.
486    ///
487    /// Use this to respond to `ControlRequest` messages received during tool approval.
488    /// The easiest way to create responses is using the helper methods on
489    /// `ToolPermissionRequest`:
490    ///
491    /// # Example
492    ///
493    /// ```no_run
494    /// use claude_codes::{AsyncClient, ClaudeOutput, ControlRequestPayload};
495    ///
496    /// # async fn example(client: &mut AsyncClient) -> Result<(), Box<dyn std::error::Error>> {
497    /// # let output = client.receive().await?;
498    /// if let ClaudeOutput::ControlRequest(req) = output {
499    ///     if let ControlRequestPayload::CanUseTool(perm_req) = &req.request {
500    ///         // Use the ergonomic helpers on ToolPermissionRequest
501    ///         let response = if perm_req.tool_name == "Bash" {
502    ///             perm_req.deny("Bash commands not allowed", &req.request_id)
503    ///         } else {
504    ///             perm_req.allow(&req.request_id)
505    ///         };
506    ///         client.send_control_response(response).await?;
507    ///     }
508    /// }
509    /// # Ok(())
510    /// # }
511    /// ```
512    pub async fn send_control_response(&mut self, response: ControlResponse) -> Result<()> {
513        let message: ControlResponseMessage = response.into();
514        let json_line = Protocol::serialize(&message)?;
515        debug!(
516            "[TOOL_APPROVAL] Sending control response: {}",
517            json_line.trim()
518        );
519
520        self.stdin
521            .write_all(json_line.as_bytes())
522            .await
523            .map_err(Error::Io)?;
524        self.stdin.flush().await.map_err(Error::Io)?;
525        Ok(())
526    }
527
528    /// Check if tool approval protocol is enabled
529    pub fn is_tool_approval_enabled(&self) -> bool {
530        self.tool_approval_enabled
531    }
532}
533
534/// A response stream that yields ClaudeOutput messages
535/// Holds a reference to the client to read from
536pub struct ResponseStream<'a> {
537    client: &'a mut AsyncClient,
538    finished: bool,
539}
540
541impl ResponseStream<'_> {
542    /// Convert to a vector by collecting all responses
543    pub async fn collect(mut self) -> Result<Vec<ClaudeOutput>> {
544        let mut responses = Vec::new();
545
546        while !self.finished {
547            let output = self.client.receive().await?;
548            let is_result = matches!(&output, ClaudeOutput::Result(_));
549            responses.push(output);
550
551            if is_result {
552                self.finished = true;
553                break;
554            }
555        }
556
557        Ok(responses)
558    }
559
560    /// Get the next response
561    pub async fn next(&mut self) -> Option<Result<ClaudeOutput>> {
562        if self.finished {
563            return None;
564        }
565
566        match self.client.receive().await {
567            Ok(output) => {
568                if matches!(&output, ClaudeOutput::Result(_)) {
569                    self.finished = true;
570                }
571                Some(Ok(output))
572            }
573            Err(e) => {
574                self.finished = true;
575                Some(Err(e))
576            }
577        }
578    }
579}
580
581impl Drop for AsyncClient {
582    fn drop(&mut self) {
583        if self.is_alive() {
584            // Try to kill the process
585            if let Err(e) = self.child.start_kill() {
586                error!("Failed to kill Claude process on drop: {}", e);
587            }
588        }
589    }
590}
591
592// Protocol extension methods for asynchronous I/O
593impl Protocol {
594    /// Write a message to an async writer
595    pub async fn write_async<W: AsyncWriteExt + Unpin, T: Serialize>(
596        writer: &mut W,
597        message: &T,
598    ) -> Result<()> {
599        let line = Self::serialize(message)?;
600        debug!("[PROTOCOL] Sending async: {}", line.trim());
601        writer.write_all(line.as_bytes()).await?;
602        writer.flush().await?;
603        Ok(())
604    }
605
606    /// Read a message from an async reader
607    pub async fn read_async<R: AsyncBufReadExt + Unpin, T: for<'de> Deserialize<'de>>(
608        reader: &mut R,
609    ) -> Result<T> {
610        let mut line = String::new();
611        let bytes_read = reader.read_line(&mut line).await?;
612        if bytes_read == 0 {
613            return Err(Error::ConnectionClosed);
614        }
615        debug!("[PROTOCOL] Received async: {}", line.trim());
616        Self::deserialize(&line)
617    }
618}
619
620/// Async stream processor for handling continuous message streams
621pub struct AsyncStreamProcessor<R> {
622    reader: AsyncBufReader<R>,
623}
624
625impl<R: tokio::io::AsyncRead + Unpin> AsyncStreamProcessor<R> {
626    /// Create a new async stream processor
627    pub fn new(reader: R) -> Self {
628        Self {
629            reader: AsyncBufReader::new(reader),
630        }
631    }
632
633    /// Process the next message from the stream
634    pub async fn next_message<T: for<'de> Deserialize<'de>>(&mut self) -> Result<T> {
635        Protocol::read_async(&mut self.reader).await
636    }
637
638    /// Process all messages in the stream
639    pub async fn process_all<T, F, Fut>(&mut self, mut handler: F) -> Result<()>
640    where
641        T: for<'de> Deserialize<'de>,
642        F: FnMut(T) -> Fut,
643        Fut: std::future::Future<Output = Result<()>>,
644    {
645        loop {
646            match self.next_message().await {
647                Ok(message) => handler(message).await?,
648                Err(Error::ConnectionClosed) => break,
649                Err(e) => return Err(e),
650            }
651        }
652        Ok(())
653    }
654}
claude_codes/client_async.rs

claude_codes/
client_async.rs
