gemini_cli_sdk/

client.rs

//! Stateful client for multi-turn Gemini CLI sessions.
//!
//! The [`Client`] ties together transport, permissions, hooks, and translation
//! into a single consumer-facing type. Call [`connect`] once to establish the
//! session, then [`send`] or [`send_content`] for each user turn.
//!
//! # Architecture
//!
//! ```text
//! Client
//!   ├── AnyTransport  ──► GeminiTransport (production) | MockTransport (testing)
//!   ├── TranslationContext — wire SessionUpdate → public Message
//!   ├── HookContext   — lifecycle hooks (UserPromptSubmit, Stop, …)
//!   └── notification_stream — single mpsc receiver taken from transport
//! ```
//!
//! # Example
//!
//! ```rust,no_run
//! use gemini_cli_sdk::{Client, ClientConfig};
//!
//! #[tokio::main]
//! async fn main() -> gemini_cli_sdk::Result<()> {
//!     let config = ClientConfig::builder()
//!         .prompt("Build a REST API")
//!         .build();
//!     let mut client = Client::new(config)?;
//!     let _info = client.connect().await?;
//!     // Consume the stream by pinning it in place.
//!     // (See tokio::pin! or futures::pin_mut! for non-Unpin streams.)
//!     client.close().await?;
//!     Ok(())
//! }
//! ```

use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use futures_core::Stream;
use serde_json::Value;
use tokio::sync::Mutex;

use crate::callback::MessageCallback;
use crate::config::ClientConfig;
use crate::hooks::{self, HookContext, HookDecision, HookEvent, HookInput};
use crate::permissions::{PermissionHandler, ToolInputCache};
use crate::translate::TranslationContext;
use crate::transport::{GeminiTransport, Transport};
use crate::types::content::UserContent;
use crate::types::messages::{Message, SessionInfo};
use crate::wire;
use crate::{Error, Result};

// ── AnyTransport ─────────────────────────────────────────────────────────────

/// Internal enum that wraps either the production [`GeminiTransport`] or the
/// test [`MockTransport`] so both can expose request/notification helpers.
///
/// This sidesteps the object-safety constraint that prevents adding generic
/// `send_request<P, R>` methods to the [`Transport`] trait directly.
pub(crate) enum AnyTransport {
    /// Production subprocess transport.
    Gemini(Arc<GeminiTransport>),
    /// In-memory transport for unit tests.
    #[cfg(feature = "testing")]
    Mock(Arc<crate::testing::MockTransport>),
}

/// Generate a simple delegation method on `AnyTransport` that forwards the call
/// identically to every variant.
macro_rules! delegate_transport {
    // async method with arguments
    (async fn $name:ident(&self $(, $arg:ident : $arg_ty:ty)*) -> $ret:ty) => {
        async fn $name(&self $(, $arg: $arg_ty)*) -> $ret {
            match self {
                AnyTransport::Gemini(t) => t.$name($($arg),*).await,
                #[cfg(feature = "testing")]
                AnyTransport::Mock(t) => t.$name($($arg),*).await,
            }
        }
    };
    // sync method
    (fn $name:ident(&self $(, $arg:ident : $arg_ty:ty)*) -> $ret:ty) => {
        fn $name(&self $(, $arg: $arg_ty)*) -> $ret {
            match self {
                AnyTransport::Gemini(t) => t.$name($($arg),*),
                #[cfg(feature = "testing")]
                AnyTransport::Mock(t) => t.$name($($arg),*),
            }
        }
    };
}

impl AnyTransport {
    // ── Transport trait delegation ────────────────────────────────────────

    delegate_transport!(async fn connect(&self) -> Result<()>);
    delegate_transport!(fn read_messages(&self) -> Pin<Box<dyn Stream<Item = Result<Value>> + Send>>);
    delegate_transport!(async fn interrupt(&self) -> Result<()>);
    delegate_transport!(async fn close(&self) -> Result<Option<i32>>);

    // ── JSON-RPC helpers ──────────────────────────────────────────────────

    /// Send a typed JSON-RPC request and await the correlated response.
    ///
    /// For `Mock` transports the params are serialised and written to the
    /// captures list; the result is a zero-value default-deserialised response.
    /// Tests that need specific response values should pre-load them via
    /// `ScenarioBuilder` or `MockTransport::push_message`.
    async fn send_request<P, R>(&self, method: &str, params: P) -> Result<R>
    where
        P: serde::Serialize + Send,
        R: serde::de::DeserializeOwned,
    {
        match self {
            AnyTransport::Gemini(t) => t.send_request(method, params).await,
            #[cfg(feature = "testing")]
            AnyTransport::Mock(t) => {
                // Capture the outbound request for assertion in tests.
                let req = serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": method,
                    "params": serde_json::to_value(params)?,
                    "id": 1
                });
                t.write(&serde_json::to_string(&req)?).await?;
                // Return a default-deserialized result. Tests using the mock
                // transport should override this via scenario injection when
                // the actual response fields matter.
                serde_json::from_value(Value::Object(Default::default())).map_err(Error::Json)
            }
        }
    }

    /// Send a JSON-RPC request and return the response receiver without
    /// blocking. The caller must await the receiver concurrently.
    async fn send_request_start<P>(
        &self,
        method: &str,
        params: P,
    ) -> Result<tokio::sync::oneshot::Receiver<crate::jsonrpc::JsonRpcResponse>>
    where
        P: serde::Serialize + Send,
    {
        match self {
            AnyTransport::Gemini(t) => t.send_request_start(method, params).await,
            #[cfg(feature = "testing")]
            AnyTransport::Mock(t) => {
                // Capture the request for test assertions.
                let req = serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": method,
                    "params": serde_json::to_value(params)?,
                    "id": 1
                });
                t.write(&serde_json::to_string(&req)?).await?;
                // Return a pre-resolved receiver with a default prompt result.
                let (tx, rx) = tokio::sync::oneshot::channel();
                let _ = tx.send(crate::jsonrpc::JsonRpcResponse::success(
                    crate::jsonrpc::JsonRpcId::Number(0),
                    serde_json::json!({"stopReason": "end_turn"}),
                ));
                Ok(rx)
            }
        }
    }

    /// Send a JSON-RPC notification (fire-and-forget, no response expected).
    async fn send_notification<P>(&self, method: &str, params: P) -> Result<()>
    where
        P: serde::Serialize + Send,
    {
        match self {
            AnyTransport::Gemini(t) => t.send_notification(method, params).await,
            #[cfg(feature = "testing")]
            AnyTransport::Mock(t) => {
                let notif = serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": method,
                    "params": serde_json::to_value(params)?
                });
                t.write(&serde_json::to_string(&notif)?).await
            }
        }
    }

    /// Register the reverse-request handler on the underlying transport.
    ///
    /// Only meaningful for the `Gemini` variant; `Mock` silently ignores it
    /// because `MockTransport` has no background reader for reverse requests.
    async fn set_reverse_handler(
        &self,
        handler: Arc<dyn crate::transport::ReverseRequestHandler>,
    ) {
        match self {
            AnyTransport::Gemini(t) => t.set_reverse_handler(handler).await,
            #[cfg(feature = "testing")]
            AnyTransport::Mock(_) => {} // no-op — mock has no background reader
        }
    }
}

// ── Client ───────────────────────────────────────────────────────────────────

/// Stateful client for multi-turn Gemini CLI sessions.
///
/// Each `Client` instance manages a single Gemini CLI subprocess and JSON-RPC
/// session. Calls to [`send`] / [`send_content`] return streams of [`Message`]
/// values translated from raw `session/update` notifications.
///
/// # Lifecycle
///
/// 1. Construct with [`Client::new`].
/// 2. Call [`connect`] — this spawns the subprocess and runs the handshake.
/// 3. Call [`send`] one or more times to converse with the model.
/// 4. Call [`close`] when finished. Dropping without closing is safe but may
///    leave the subprocess running briefly until the OS reclaims it.
///
/// # Threading
///
/// `Client` is `Send` but not `Sync`. Share across tasks by wrapping in
/// `Arc<Mutex<Client>>` when concurrent access is required.
///
/// [`connect`]: Client::connect
/// [`send`]: Client::send
/// [`close`]: Client::close
pub struct Client {
    /// Full session configuration — fields are referenced during the connect
    /// handshake and at the start of each prompt turn.
    config: ClientConfig,
    /// Concrete transport implementation, wrapped in an enum for testability.
    transport: AnyTransport,
    /// Session ID assigned by the server after `session/new` or `session/load`.
    session_id: Option<String>,
    /// Stored notification stream — taken exactly once in `connect()` via
    /// `Transport::read_messages()`. Held behind a `Mutex` so the async-stream
    /// closure inside `send_content()` can lock it across yield points.
    #[allow(clippy::type_complexity)]
    notification_stream: Mutex<Option<Pin<Box<dyn Stream<Item = Result<Value>> + Send>>>>,
    /// Accumulated per-turn translation state (text buffer, tool calls, …).
    translation_ctx: Mutex<Option<TranslationContext>>,
    /// Immutable context stamped onto every hook invocation.
    hook_context: Option<HookContext>,
    /// `true` after a successful `connect()`.
    connected: bool,
    /// Guard that prevents concurrent `send_content` calls from silently
    /// hanging on the `notification_stream` Mutex. Set to `true` at the
    /// start of `send_content`, reset to `false` when the stream completes
    /// or in `close()`.
    turn_in_progress: Arc<AtomicBool>,
}

/// RAII guard that resets `turn_in_progress` to `false` on drop, ensuring the
/// flag is cleared even when the stream or function body returns early.
struct TurnGuard(Arc<AtomicBool>);

impl Drop for TurnGuard {
    fn drop(&mut self) {
        self.0.store(false, Ordering::Release);
    }
}

impl Client {
    // ── Constructors ─────────────────────────────────────────────────────────

    /// Build a `Client` from a pre-constructed transport. All public
    /// constructors delegate to this.
    fn from_transport(config: ClientConfig, transport: AnyTransport) -> Self {
        Self {
            config,
            transport,
            session_id: None,
            notification_stream: Mutex::new(None),
            translation_ctx: Mutex::new(None),
            hook_context: None,
            connected: false,
            turn_in_progress: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Create a new client with the given configuration.
    ///
    /// Resolves the `gemini` binary path (via `config.cli_path` or `PATH`)
    /// and constructs a [`GeminiTransport`]. The subprocess is not spawned
    /// until [`connect`] is called.
    ///
    /// # Errors
    ///
    /// Returns [`Error::CliNotFound`] when the binary cannot be located on
    /// `PATH` and `config.cli_path` is `None`.
    ///
    /// [`connect`]: Client::connect
    /// [`Error::CliNotFound`]: crate::Error::CliNotFound
    pub fn new(config: ClientConfig) -> Result<Self> {
        let transport = Arc::new(GeminiTransport::from_config(&config)?);
        Ok(Self::from_transport(config, AnyTransport::Gemini(transport)))
    }

    /// Create a client backed by a caller-supplied [`GeminiTransport`].
    ///
    /// Useful when the caller has already constructed the transport with custom
    /// parameters (e.g. a non-default working directory or extra env vars).
    pub fn with_gemini_transport(config: ClientConfig, transport: Arc<GeminiTransport>) -> Self {
        Self::from_transport(config, AnyTransport::Gemini(transport))
    }

    /// Create a client backed by a [`MockTransport`] for unit testing.
    ///
    /// Only available when the `testing` crate feature is enabled. The mock
    /// transport captures writes and yields pre-loaded messages, making it
    /// straightforward to test connect / send behaviour without spawning a
    /// real subprocess.
    ///
    /// [`MockTransport`]: crate::testing::MockTransport
    #[cfg(feature = "testing")]
    pub fn with_mock_transport(
        config: ClientConfig,
        transport: Arc<crate::testing::MockTransport>,
    ) -> Self {
        Self::from_transport(config, AnyTransport::Mock(transport))
    }

    // ── Accessors ─────────────────────────────────────────────────────────────

    /// Return the session ID assigned by the server, or `None` before
    /// [`connect`] is called.
    ///
    /// [`connect`]: Client::connect
    pub fn session_id(&self) -> Option<&str> {
        self.session_id.as_deref()
    }

    /// Return the `prompt` field from the config.
    ///
    /// Exposed as a convenience for the free-function wrappers in `lib.rs`
    /// that need to send the initial prompt without a separate `send()` call.
    #[inline]
    pub fn prompt(&self) -> &str {
        &self.config.prompt
    }

    /// Return `true` if [`connect`] has been called successfully.
    ///
    /// [`connect`]: Client::connect
    #[inline]
    pub fn is_connected(&self) -> bool {
        self.connected
    }

    // ── connect() ────────────────────────────────────────────────────────────

    /// Connect to the Gemini CLI and establish a session.
    ///
    /// Performs the full initialisation sequence in order:
    ///
    /// 1. Spawn the subprocess via [`Transport::connect`].
    /// 2. Take the notification stream (must happen before any requests).
    /// 3. Register the optional [`PermissionHandler`] as the reverse-request handler.
    /// 4. Send the `initialize` JSON-RPC request and await the result.
    /// 5. Send `session/new` (or `session/load` when `config.resume` is set).
    /// 6. Initialise the [`TranslationContext`] and hook context.
    ///
    /// Returns [`SessionInfo`] describing the established session.
    ///
    /// # Errors
    ///
    /// - [`Error::Config`] — called more than once on the same client.
    /// - [`Error::SpawnFailed`] — the subprocess could not be started.
    /// - [`Error::JsonRpcError`] — the server rejected `initialize` or `session/new`.
    /// - [`Error::NotConnected`] — internal transport error during the handshake.
    ///
    /// [`Transport::connect`]: crate::transport::Transport::connect
    /// [`Error::Config`]: crate::Error::Config
    /// [`Error::SpawnFailed`]: crate::Error::SpawnFailed
    /// [`Error::JsonRpcError`]: crate::Error::JsonRpcError
    /// [`Error::NotConnected`]: crate::Error::NotConnected
    pub async fn connect(&mut self) -> Result<SessionInfo> {
        if self.connected {
            return Err(Error::Config("Already connected".to_string()));
        }
        match self.config.connect_timeout {
            Some(d) => {
                tokio::time::timeout(d, self.connect_inner())
                    .await
                    .map_err(|_| {
                        Error::Timeout(format!(
                            "connect timed out after {:.1}s",
                            d.as_secs_f64()
                        ))
                    })?
            }
            None => self.connect_inner().await,
        }
    }

    async fn connect_inner(&mut self) -> Result<SessionInfo> {
        // ── Step 1: Spawn subprocess ─────────────────────────────────────────
        self.transport.connect().await?;

        // ── Step 2: Take the notification stream ─────────────────────────────
        let stream = self.transport.read_messages();
        *self.notification_stream.lock().await = Some(stream);

        // ── Step 3: Create shared tool input cache ───────────────────────────
        let tool_input_cache: ToolInputCache =
            Arc::new(std::sync::Mutex::new(std::collections::HashMap::new()));

        // ── Step 4: Register permission handler ──────────────────────────────
        if let Some(callback) = self.config.can_use_tool.clone() {
            let handler =
                Arc::new(PermissionHandler::new(callback, Some(Arc::clone(&tool_input_cache))));
            self.transport.set_reverse_handler(handler).await;
        }

        // ── Step 5: initialize request ───────────────────────────────────────
        let init_params = wire::InitializeParams {
            protocol_version: 1,
            client_capabilities: wire::ClientCapabilities::default(),
            client_info: wire::ClientInfo {
                name: "gemini-cli-sdk".to_string(),
                version: env!("CARGO_PKG_VERSION").to_string(),
            },
        };
        let init_result: wire::InitializeResult = self
            .transport
            .send_request(wire::method::INITIALIZE, init_params)
            .await?;

        // ── Step 6: Create or resume session ─────────────────────────────────
        let session_id = if let Some(resume_id) = self.config.resume.clone() {
            let params = wire::SessionLoadParams {
                session_id: resume_id,
                extra: Value::Object(Default::default()),
            };
            let result: wire::SessionLoadResult = self
                .transport
                .send_request(wire::method::SESSION_LOAD, params)
                .await?;
            result.session_id
        } else {
            let cwd = self
                .config
                .cwd
                .clone()
                .map(Ok)
                .unwrap_or_else(|| {
                    std::env::current_dir()
                        .map_err(|e| Error::Config(format!("cannot determine cwd: {e}")))
                })?
                .to_string_lossy()
                .to_string();
            let mcp_wire = crate::mcp::mcp_servers_to_wire(&self.config.mcp_servers);
            let params = wire::SessionNewParams {
                cwd,
                mcp_servers: mcp_wire,
                extra: Value::Object(Default::default()),
            };
            let result: wire::SessionNewResult = self
                .transport
                .send_request(wire::method::SESSION_NEW, params)
                .await?;
            result.session_id
        };

        self.session_id = Some(session_id.clone());

        // ── Step 7: Initialise translation and hook contexts ─────────────────
        let model = self
            .config
            .model
            .clone()
            .unwrap_or_else(|| "gemini-2.5-pro".to_string());

        *self.translation_ctx.lock().await =
            Some(TranslationContext::new_with_cache(session_id.clone(), model.clone(), tool_input_cache));

        let cwd_str = self
            .config
            .cwd
            .clone()
            .map(Ok)
            .unwrap_or_else(|| {
                std::env::current_dir()
                    .map_err(|e| Error::Config(format!("cannot determine cwd: {e}")))
            })?
            .to_string_lossy()
            .to_string();

        self.hook_context = Some(HookContext {
            session_id: session_id.clone(),
            cwd: cwd_str,
        });

        self.connected = true;

        let tools = init_result.agent_capabilities.tools.unwrap_or_default();
        Ok(SessionInfo {
            session_id,
            model,
            tools,
            extra: init_result.extra,
        })
    }

    // ── send() / send_content() ───────────────────────────────────────────────

    /// Send a plain-text prompt and return a stream of translated [`Message`]
    /// values.
    ///
    /// Convenience wrapper around [`send_content`] that constructs a single
    /// `UserContent::Text` block from the provided string slice.
    ///
    /// # Errors
    ///
    /// Returns [`Error::NotConnected`] if [`connect`] has not been called.
    /// Returns [`Error::Config`] if a `UserPromptSubmit` hook blocks the send.
    ///
    /// [`send_content`]: Client::send_content
    /// [`connect`]: Client::connect
    /// [`Error::NotConnected`]: crate::Error::NotConnected
    /// [`Error::Config`]: crate::Error::Config
    pub async fn send(
        &self,
        message: &str,
    ) -> Result<impl Stream<Item = Result<Message>> + '_> {
        self.send_content(vec![UserContent::text(message)]).await
    }

    /// Send structured content and return a stream of translated [`Message`]
    /// values.
    ///
    /// Accepts any mix of [`UserContent`] variants (text, base-64 image, URL
    /// image). The content is serialised to `WireContentBlock` format and sent
    /// to the CLI as a `session/prompt` JSON-RPC notification. The returned
    /// stream drains `session/update` notifications from the shared channel
    /// until the transport closes or the caller drops the stream.
    ///
    /// # Backpressure
    ///
    /// The notification channel is shared across all `send_content` calls on
    /// the same client. Only one stream should be active at a time; concurrent
    /// polling will contend on the `notification_stream` Mutex and may
    /// interleave results.
    ///
    /// # Errors
    ///
    /// Returns [`Error::NotConnected`] if [`connect`] has not been called.
    /// Returns [`Error::Config`] if a `UserPromptSubmit` hook blocks the send.
    ///
    /// [`connect`]: Client::connect
    /// [`Error::NotConnected`]: crate::Error::NotConnected
    /// [`Error::Config`]: crate::Error::Config
    pub async fn send_content(
        &self,
        content: Vec<UserContent>,
    ) -> Result<impl Stream<Item = Result<Message>> + '_> {
        if !self.connected {
            return Err(Error::NotConnected);
        }

        // Prevent concurrent turns: a second call would block indefinitely on
        // the notification_stream Mutex. Fail fast with a descriptive error.
        if self
            .turn_in_progress
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            return Err(Error::TurnInProgress);
        }
        let turn_guard = TurnGuard(Arc::clone(&self.turn_in_progress));

        let session_id = self
            .session_id
            .as_ref()
            .ok_or(Error::NotConnected)?
            .clone();

        // ── Fire UserPromptSubmit hook ────────────────────────────────────────
        if let Some(ctx) = &self.hook_context {
            let prompt_text = content.iter().find_map(|c| match c {
                UserContent::Text { text } => Some(text.clone()),
                _ => None,
            });
            let hook_input = HookInput {
                event: HookEvent::UserPromptSubmit,
                tool_name: None,
                tool_input: None,
                tool_output: None,
                prompt: prompt_text,
                session_id: session_id.clone(),
                extra: Value::Object(Default::default()),
            };
            let output = hooks::execute_hooks(
                &self.config.hooks,
                hook_input,
                ctx,
                self.config.default_hook_timeout,
            )
            .await;
            if output.decision == HookDecision::Block {
                return Err(Error::Config(
                    output
                        .message
                        .unwrap_or_else(|| "Blocked by hook".to_string()),
                ));
            }
        }

        // ── Reset translation context for the new turn ────────────────────────
        {
            let mut ctx_guard = self.translation_ctx.lock().await;
            if let Some(ctx) = ctx_guard.as_mut() {
                ctx.reset_turn();
            }
        }

        // ── Convert content to wire format ────────────────────────────────────
        let wire_content: Vec<wire::WireContentBlock> = content
            .iter()
            .map(crate::translate::user_content_to_wire)
            .collect();

        // ── Send session/prompt as a request ──────────────────────────────────
        //
        // `session/prompt` is sent as a JSON-RPC request (with `id`) so that the
        // Gemini CLI sends back a correlated response containing `stopReason`
        // when the turn completes. The notification stream delivers `session/update`
        // events (text deltas, tool calls, etc.) concurrently, while the request
        // response signals the turn boundary.
        let prompt_params = wire::SessionPromptParams {
            session_id: session_id.clone(),
            prompt: wire_content,
            extra: Value::Object(Default::default()),
        };
        let prompt_response_rx = self
            .transport
            .send_request_start(wire::method::SESSION_PROMPT, prompt_params)
            .await?;

        // ── Return notification-draining stream ───────────────────────────────
        //
        // Borrow `self` for the stream's lifetime. The Mutex is locked inside
        // the generator so it can be released between yields, avoiding a held
        // lock across await points in caller code.
        let translation_ctx = &self.translation_ctx;
        let notification_stream = &self.notification_stream;
        let callback: Option<MessageCallback> = self.config.message_callback.clone();

        Ok(async_stream::stream! {
            // Move the guard into the stream so it is dropped (clearing the flag)
            // when the stream is dropped or completes, regardless of the exit path.
            let _turn_guard = turn_guard;
            use tokio_stream::StreamExt as _;

            // Lock the notification stream for the duration of this turn.
            // A second concurrent call will block here until this stream is
            // dropped.
            let mut ns_guard = notification_stream.lock().await;
            let stream = match ns_guard.as_mut() {
                Some(s) => s,
                None => {
                    yield Err(Error::NotConnected);
                    return;
                }
            };

            // Fuse the prompt response oneshot so we can poll it alongside
            // the notification stream without consuming it on the first poll.
            let mut prompt_done = prompt_response_rx;
            let mut turn_finished = false;

            #[allow(unused_assignments)] // `turn_finished = true` precedes a `break`
            loop {
                tokio::select! {
                    biased;

                    // Poll notifications first — drain all pending updates before
                    // checking if the turn is complete.  With `biased`, this branch
                    // is checked before prompt_done, ensuring mock transports (which
                    // resolve the oneshot immediately) still deliver notifications.
                    maybe_notif = stream.next() => {
                        match maybe_notif {
                            None => break, // channel closed — subprocess exited
                            Some(Err(e)) => {
                                yield Err(e);
                                break;
                            }
                            Some(Ok(value)) => {
                                // Filter: only process session/update notifications.
                                let method = value
                                    .get("method")
                                    .and_then(|m| m.as_str())
                                    .unwrap_or("");
                                tracing::debug!(method, "notification received");
                                if method != wire::method::SESSION_UPDATE {
                                    continue;
                                }

                                let params = match value.get("params") {
                                    Some(p) => p.clone(),
                                    None => continue,
                                };

                                let notif: wire::SessionUpdateNotification =
                                    match serde_json::from_value(params) {
                                        Ok(n) => n,
                                        Err(e) => {
                                            tracing::warn!(
                                                error = %e,
                                                "client: failed to parse session/update params — skipping"
                                            );
                                            continue;
                                        }
                                    };

                                // Parse the discriminator and translate to public messages.
                                let update = notif.parse();
                                tracing::debug!(?update, "parsed session update");

                                let mut ctx_guard = translation_ctx.lock().await;
                                if let Some(ctx) = ctx_guard.as_mut() {
                                    let messages = ctx.translate(update);
                                    tracing::debug!(count = messages.len(), "translated to messages");
                                    // Drop the ctx lock before yielding so callers that
                                    // inspect TranslationContext are not blocked.
                                    drop(ctx_guard);

                                    for msg in messages {
                                        // Invoke optional side-effect callback.
                                        if let Some(cb) = &callback {
                                            cb(msg.clone()).await;
                                        }
                                        yield Ok(msg);
                                    }
                                }
                            }
                        }
                    }

                    // Poll prompt response — when it arrives, the turn is done.
                    resp = &mut prompt_done, if !turn_finished => {
                        turn_finished = true;

                        // Translate the prompt result into a Message::Result.
                        match resp {
                            Ok(response) => {
                                match response.into_result() {
                                    Ok(result_value) => {
                                        let prompt_result: wire::SessionPromptResult =
                                            serde_json::from_value(result_value)
                                                .unwrap_or_else(|e| {
                                                    tracing::warn!(
                                                        error = %e,
                                                        "failed to parse SessionPromptResult, using default"
                                                    );
                                                    Default::default()
                                                });

                                        let result_msg = Message::Result(crate::types::messages::ResultMessage {
                                            subtype: "success".to_string(),
                                            is_error: false,
                                            duration_ms: 0.0,
                                            duration_api_ms: 0.0,
                                            num_turns: 1,
                                            session_id: session_id.clone(),
                                            usage: crate::types::messages::Usage::default(),
                                            stop_reason: prompt_result.stop_reason,
                                            extra: prompt_result.extra,
                                        });

                                        if let Some(cb) = &callback {
                                            cb(result_msg.clone()).await;
                                        }
                                        yield Ok(result_msg);
                                    }
                                    Err(err) => {
                                        let error_msg = Message::Result(crate::types::messages::ResultMessage {
                                            subtype: "error".to_string(),
                                            is_error: true,
                                            duration_ms: 0.0,
                                            duration_api_ms: 0.0,
                                            num_turns: 1,
                                            session_id: session_id.clone(),
                                            usage: crate::types::messages::Usage::default(),
                                            stop_reason: format!(
                                                "JSON-RPC error {}: {}",
                                                err.code, err.message
                                            ),
                                            extra: serde_json::json!({
                                                "code": err.code,
                                                "message": err.message,
                                                "data": err.data,
                                            }),
                                        });

                                        if let Some(cb) = &callback {
                                            cb(error_msg.clone()).await;
                                        }
                                        yield Ok(error_msg);
                                    }
                                }
                            }
                            Err(_) => {
                                // Response channel dropped — treat as transport error.
                                yield Err(Error::Transport(
                                    "Prompt response channel closed unexpectedly".to_string()
                                ));
                            }
                        }

                        // Reset the translation context for the next turn.
                        let mut ctx_guard = translation_ctx.lock().await;
                        if let Some(ctx) = ctx_guard.as_mut() {
                            ctx.reset_turn();
                        }

                        break;
                    }
                }
            }

            // Turn complete — _turn_guard is dropped here, clearing the flag.
        })
    }

    // ── interrupt() ──────────────────────────────────────────────────────────

    /// Interrupt the current in-progress prompt turn.
    ///
    /// Sends a `session/cancel` notification to the CLI (best-effort) and then
    /// delivers a process-level interrupt signal via the transport (SIGINT on
    /// Unix, CTRL_C_EVENT on Windows).
    ///
    /// # Errors
    ///
    /// The `session/cancel` notification errors are silently ignored. Transport
    /// interrupt errors are propagated.
    pub async fn interrupt(&self) -> Result<()> {
        if let Some(session_id) = &self.session_id {
            let params = wire::SessionCancelParams {
                session_id: session_id.clone(),
            };
            // Best-effort: a cancel notification failure must not prevent the
            // process-level interrupt below from being delivered.
            let _ = self
                .transport
                .send_notification(wire::method::SESSION_CANCEL, params)
                .await;
        }
        self.transport.interrupt().await
    }

    // ── close() ──────────────────────────────────────────────────────────────

    /// Close the client and terminate the CLI subprocess.
    ///
    /// Fires the `Stop` lifecycle hook before closing the transport. The
    /// call is idempotent — invoking it on an already-closed client is safe.
    ///
    /// # Errors
    ///
    /// Propagates errors from the transport's `close()` implementation. Hook
    /// errors are silently ignored.
    pub async fn close(&mut self) -> Result<()> {
        // ── Fire Stop hook ────────────────────────────────────────────────────
        if let Some(ctx) = &self.hook_context {
            let hook_input = HookInput {
                event: HookEvent::Stop,
                tool_name: None,
                tool_input: None,
                tool_output: None,
                prompt: None,
                session_id: self.session_id.clone().unwrap_or_default(),
                extra: Value::Object(Default::default()),
            };
            // Ignore errors — we are shutting down regardless.
            let _ = hooks::execute_hooks(
                &self.config.hooks,
                hook_input,
                ctx,
                self.config.default_hook_timeout,
            )
            .await;
        }

        self.connected = false;
        self.turn_in_progress.store(false, Ordering::Release);
        self.transport.close().await?;
        Ok(())
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // ── Helpers ───────────────────────────────────────────────────────────────

    /// Build a `GeminiTransport` that points to a nonexistent binary so tests
    /// never accidentally spawn a real subprocess.
    fn make_fake_transport() -> Arc<GeminiTransport> {
        Arc::new(GeminiTransport::new(
            std::path::PathBuf::from("/nonexistent/gemini"),
            vec!["--experimental-acp".to_string()],
            std::path::PathBuf::from("/tmp"),
            std::collections::HashMap::new(),
            None,
            None,
        ))
    }

    fn minimal_config() -> ClientConfig {
        ClientConfig::builder().prompt("test prompt").build()
    }

    // ── test_client_session_id ────────────────────────────────────────────────

    /// Before `connect()`, `session_id()` must return `None`.
    #[test]
    fn test_client_session_id() {
        let client = Client::with_gemini_transport(minimal_config(), make_fake_transport());
        assert!(
            client.session_id().is_none(),
            "session_id must be None before connect() is called"
        );
    }

    // ── test_client_not_connected_error ───────────────────────────────────────

    /// Calling `send()` before `connect()` must return `Err(Error::NotConnected)`.
    #[tokio::test]
    async fn test_client_not_connected_error() {
        let client = Client::with_gemini_transport(minimal_config(), make_fake_transport());
        let result = client.send("hello").await;
        assert!(result.is_err(), "send() before connect() must fail");
        let err = result.err().expect("expected an error");
        assert!(
            matches!(err, Error::NotConnected),
            "error must be Error::NotConnected, got: {err:?}"
        );
    }

    // ── test_client_send_content_not_connected ────────────────────────────────

    /// `send_content()` before connect must also return `Error::NotConnected`.
    #[tokio::test]
    async fn test_client_send_content_not_connected() {
        let client = Client::with_gemini_transport(minimal_config(), make_fake_transport());
        let result = client.send_content(vec![UserContent::text("hi")]).await;
        let err = result.err().expect("expected an error");
        assert!(
            matches!(err, Error::NotConnected),
            "send_content before connect must return Error::NotConnected, got: {err:?}"
        );
    }

    // ── test_client_prompt_accessor ───────────────────────────────────────────

    /// `prompt()` must reflect the value set in the config.
    #[test]
    fn test_client_prompt_accessor() {
        let config = ClientConfig::builder().prompt("my test prompt").build();
        let client = Client::with_gemini_transport(config, make_fake_transport());
        assert_eq!(client.prompt(), "my test prompt");
    }

    // ── test_client_is_connected_default ─────────────────────────────────────

    /// A freshly constructed client must not report itself as connected.
    #[test]
    fn test_client_is_connected_default() {
        let client = Client::with_gemini_transport(minimal_config(), make_fake_transport());
        assert!(
            !client.is_connected(),
            "is_connected must be false before connect()"
        );
    }

    // ── test_client_double_connect_error ─────────────────────────────────────

    /// Calling `connect()` after a successful connection must return
    /// `Err(Error::Config)`. We simulate the connected state by setting the
    /// field directly — no real subprocess is involved.
    #[tokio::test]
    async fn test_client_double_connect_error() {
        let mut client =
            Client::with_gemini_transport(minimal_config(), make_fake_transport());
        // Bypass the real connect sequence by directly marking as connected.
        client.connected = true;
        let result = client.connect().await;
        assert!(result.is_err());
        assert!(
            matches!(result.unwrap_err(), Error::Config(_)),
            "second connect must return Error::Config"
        );
    }

    // ── test_client_interrupt_before_connect ──────────────────────────────────

    /// `interrupt()` before connect must not panic and must return `Ok(())`.
    /// When there is no subprocess, the platform-level signal is a no-op.
    #[tokio::test]
    async fn test_client_interrupt_before_connect() {
        let client = Client::with_gemini_transport(minimal_config(), make_fake_transport());
        // No session_id, no subprocess — must not panic.
        let result = client.interrupt().await;
        assert!(
            result.is_ok(),
            "interrupt before connect must not return an error"
        );
    }

    // ── test_client_mock_transport_constructor ────────────────────────────────

    #[cfg(feature = "testing")]
    #[test]
    fn test_client_mock_transport_constructor() {
        use crate::testing::MockTransport;
        let transport = Arc::new(MockTransport::new(vec![]));
        let client = Client::with_mock_transport(minimal_config(), transport);
        assert!(!client.is_connected());
        assert!(client.session_id().is_none());
    }
}