claude-agent-sdk-rust/src/client.rs

//! [`Client`] — bidirectional, multi-turn session with Claude Code.
//!
//! Mirrors the Python SDK's `ClaudeSDKClient`. Use [`Client::connect`] to
//! spawn the CLI, [`Client::send`] to push user messages, and
//! [`Client::messages`] to receive a stream of [`Message`]s.
//!
//! ```no_run
//! use claude_agent_sdk::{Client, ClaudeAgentOptions, Message};
//! use tokio_stream::StreamExt;
//!
//! # async fn run() -> claude_agent_sdk::Result<()> {
//! let mut client = Client::new(ClaudeAgentOptions::new()).await?;
//! client.connect().await?;
//! client.send("Hello!").await?;
//! let mut stream = client.messages();
//! while let Some(msg) = stream.next().await {
//!     let msg = msg?;
//!     if msg.is_result() {
//!         break;
//!     }
//! }
//! client.disconnect().await?;
//! # Ok(()) }
//! ```
//!
//! # v0.1 limitations
//!
//! - No control protocol — `interrupt()`, `set_permission_mode()`,
//!   `set_model()`, and other live-mutation methods are deferred to v0.2.
//! - The Python SDK sends an `initialize` control request before the first
//!   user message; this SDK skips that step because v0.1 doesn't surface
//!   hooks or agents-via-initialize. Newer CLI versions still operate
//!   correctly without it.

use futures_core::Stream;
use serde_json::json;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;

use crate::errors::{Error, Result};
use crate::messages::Message;
use crate::options::ClaudeAgentOptions;
use crate::transport::{SubprocessTransport, TransportHandle, TransportWriter};

/// Buffer depth on the inbound message channel — picked to match typical CLI
/// burst size for one turn.
const MESSAGE_CHANNEL_BUFFER: usize = 64;

/// A bidirectional client for the Claude CLI.
///
/// Lifecycle: [`new`](Self::new) → [`connect`](Self::connect) → one or more
/// [`send`](Self::send) calls interleaved with consumption of
/// [`messages`](Self::messages) → [`disconnect`](Self::disconnect).
///
/// `Drop` does NOT kill the subprocess — call `disconnect()` for a clean
/// shutdown. (Killing on drop would leak a tokio task; the `TransportHandle`
/// drop still aborts the child process if the runtime allows it.)
pub struct Client {
    options: ClaudeAgentOptions,
    writer: Option<TransportWriter>,
    handle: Option<TransportHandle>,
    reader_task: Option<tokio::task::JoinHandle<()>>,
    rx: Option<mpsc::Receiver<Result<Message>>>,
    session_id: String,
    disconnected: bool,
}

impl Client {
    /// Build a client from options. Resolves the CLI binary path eagerly.
    pub async fn new(options: ClaudeAgentOptions) -> Result<Self> {
        // Eager path resolution: catches CliNotFound before spawn.
        let _ = SubprocessTransport::new(options.clone())?;
        Ok(Self {
            options,
            writer: None,
            handle: None,
            reader_task: None,
            rx: None,
            session_id: "default".into(),
            disconnected: false,
        })
    }

    /// Override the session ID stamped on outgoing user messages. Defaults
    /// to `"default"`.
    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
        self.session_id = id.into();
        self
    }

    /// Spawn the CLI subprocess and start the reader task. Idempotent — a
    /// second call after a successful connect is a no-op.
    pub async fn connect(&mut self) -> Result<()> {
        if self.writer.is_some() {
            return Ok(());
        }
        let transport = SubprocessTransport::new(self.options.clone())?;
        let (mut reader, writer, handle) = transport.connect().await?;

        let (tx, rx) = mpsc::channel::<Result<Message>>(MESSAGE_CHANNEL_BUFFER);
        let reader_task = tokio::spawn(async move {
            loop {
                match reader.read_frame().await {
                    Ok(Some(value)) => {
                        let parsed: Result<Message> = serde_json::from_value(value.clone())
                            .map_err(|e| {
                                Error::parse_with_data(
                                    format!("could not parse CLI frame: {e}"),
                                    value,
                                )
                            });
                        if tx.send(parsed).await.is_err() {
                            break;
                        }
                    }
                    Ok(None) => break,
                    Err(e) => {
                        let _ = tx.send(Err(e)).await;
                        break;
                    }
                }
            }
        });

        self.writer = Some(writer);
        self.handle = Some(handle);
        self.reader_task = Some(reader_task);
        self.rx = Some(rx);
        Ok(())
    }

    /// Send a plain user message. Each call appends one JSON frame:
    ///
    /// ```json
    /// {"type": "user", "message": {"role": "user", "content": "<prompt>"},
    ///  "parent_tool_use_id": null, "session_id": "<id>"}
    /// ```
    ///
    /// The client must be connected — call [`connect`](Self::connect) first.
    pub async fn send(&self, prompt: impl Into<String>) -> Result<()> {
        let writer = self
            .writer
            .as_ref()
            .ok_or_else(|| Error::conn("not connected; call connect() first"))?;
        let frame = json!({
            "type": "user",
            "message": {"role": "user", "content": prompt.into()},
            "parent_tool_use_id": serde_json::Value::Null,
            "session_id": self.session_id,
        });
        writer.write_frame(&frame.to_string()).await
    }

    /// Send a raw user-message JSON frame. Use this when you need to push
    /// structured content blocks (tool results, etc.). The caller is
    /// responsible for shape — the frame is forwarded verbatim modulo a
    /// trailing newline and a stamped `session_id` if absent.
    pub async fn send_raw(&self, mut frame: serde_json::Value) -> Result<()> {
        let writer = self
            .writer
            .as_ref()
            .ok_or_else(|| Error::conn("not connected; call connect() first"))?;
        if let serde_json::Value::Object(map) = &mut frame {
            map.entry("session_id".to_string())
                .or_insert(serde_json::Value::String(self.session_id.clone()));
        }
        writer.write_frame(&frame.to_string()).await
    }

    /// Take the receiver as a [`Stream`]. May only be called once per
    /// `connect()` — subsequent calls return an empty stream.
    pub fn messages(&mut self) -> impl Stream<Item = Result<Message>> + Send + 'static + use<> {
        match self.rx.take() {
            Some(rx) => ReceiverStream::new(rx),
            None => {
                let (_, rx) = mpsc::channel(1);
                ReceiverStream::new(rx)
            }
        }
    }

    /// Close stdin and wait for the subprocess to exit. Idempotent.
    ///
    /// Returns the subprocess exit error as [`Error::Process`] when the CLI
    /// exited non-zero (matches the Python SDK).
    pub async fn disconnect(&mut self) -> Result<()> {
        if self.disconnected {
            return Ok(());
        }
        self.disconnected = true;
        // Drop the receiver first so the reader sees a closed channel and
        // exits when stdout EOFs.
        drop(self.rx.take());
        // Close stdin via the writer to let the CLI shut down cleanly.
        if let Some(writer) = self.writer.take() {
            writer.end_input().await;
        }
        // Now await the reader task — it exits on stdout EOF (since stdin is
        // closed, the CLI will reach end-of-input and exit, which closes
        // stdout from its end).
        if let Some(task) = self.reader_task.take() {
            let _ = task.await;
        }
        // Finally, reap the subprocess.
        if let Some(handle) = self.handle.take() {
            handle.close().await?;
        }
        Ok(())
    }
}

impl Drop for Client {
    fn drop(&mut self) {
        // The TransportHandle's Drop will start_kill() the child if it's
        // still alive; the reader task will then exit on its own. We don't
        // explicitly tear down here because async cleanup on Drop is not
        // ergonomic in Rust — callers should use disconnect() for clean
        // shutdown.
    }
}