xai_rust/api/

responses.rs

//! Responses API - the primary endpoint for chat interactions.

use futures_util::Stream;
use pin_project_lite::pin_project;
use serde::Serialize;
use std::collections::HashSet;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::time::sleep;

use crate::client::XaiClient;
use crate::config::RetryPolicy;
use crate::models::content::ContentPart;
use crate::models::message::{Message, MessageContent, Role};
use crate::models::response::{OutputItem, Response, ResponseFormat, StreamChunk, TextContent};
use crate::models::tool::{Tool, ToolCall, ToolChoice};
use crate::stream::ResponseStream;
use crate::{Error, Result};

const DEFAULT_DEFERRED_MAX_ATTEMPTS: u32 = 30;
const DEFAULT_DEFERRED_POLL_INTERVAL: Duration = Duration::ZERO;
const DEFAULT_STATEFUL_TOOL_LOOP_MAX_ROUNDS: u32 = 8;

/// API for creating model responses.
#[derive(Debug, Clone)]
pub struct ResponsesApi {
    client: XaiClient,
}

impl ResponsesApi {
    pub(crate) fn new(client: XaiClient) -> Self {
        Self { client }
    }

    /// Create a new response request builder.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use xai_rust::{XaiClient, Role};
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = XaiClient::from_env()?;
    ///
    /// let response = client.responses()
    ///     .create("grok-4")
    ///     .message(Role::System, "You are a helpful assistant.")
    ///     .message(Role::User, "Hello!")
    ///     .temperature(0.7)
    ///     .send()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn create(&self, model: impl Into<String>) -> CreateResponseBuilder {
        CreateResponseBuilder::new(self.client.clone(), model.into())
    }

    /// Create a deferred response poller.
    pub fn deferred(&self, response_id: impl Into<String>) -> DeferredResponsePoller {
        DeferredResponsePoller::new(self.client.clone(), response_id.into())
    }

    /// Create a stateful chat handle for multi-turn interactions.
    ///
    /// This helper keeps a local message list and lets you append messages over time,
    /// then sample or stream with the accumulated context.
    pub fn chat(&self, model: impl Into<String>) -> StatefulChat {
        StatefulChat::new(self.client.clone(), model.into())
    }

    /// Get a previous response by ID.
    pub async fn get(&self, response_id: &str) -> Result<Response> {
        let id = XaiClient::encode_path(response_id);
        let url = format!("{}/responses/{}", self.client.base_url(), id);

        let response = self.client.send(self.client.http().get(&url)).await?;

        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }

        Ok(response.json().await?)
    }

    /// Delete a response by ID.
    pub async fn delete(&self, response_id: &str) -> Result<()> {
        let id = XaiClient::encode_path(response_id);
        let url = format!("{}/responses/{}", self.client.base_url(), id);

        let response = self.client.send(self.client.http().delete(&url)).await?;

        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }

        Ok(())
    }

    /// Poll a response until output is available or attempts are exhausted.
    pub async fn poll_until_ready(&self, response_id: &str, max_attempts: u32) -> Result<Response> {
        self.deferred(response_id.to_string())
            .max_attempts(max_attempts)
            .wait()
            .await
    }
}

/// Helper for polling deferred response completion.
#[derive(Debug, Clone)]
pub struct DeferredResponsePoller {
    client: XaiClient,
    response_id: String,
    max_attempts: u32,
    poll_initial_delay: Duration,
    poll_max_delay: Duration,
}

impl DeferredResponsePoller {
    fn new(client: XaiClient, response_id: String) -> Self {
        Self {
            client,
            response_id,
            max_attempts: DEFAULT_DEFERRED_MAX_ATTEMPTS,
            poll_initial_delay: DEFAULT_DEFERRED_POLL_INTERVAL,
            poll_max_delay: DEFAULT_DEFERRED_POLL_INTERVAL,
        }
    }

    /// Set the maximum polling attempts.
    pub fn max_attempts(mut self, max_attempts: u32) -> Self {
        self.max_attempts = max_attempts.max(1);
        self
    }

    /// Set a fixed polling interval between attempts.
    pub fn poll_interval(mut self, interval: Duration) -> Self {
        self.poll_initial_delay = interval;
        self.poll_max_delay = interval;
        self
    }

    /// Set exponential polling backoff bounds between attempts.
    pub fn poll_backoff(mut self, initial: Duration, max: Duration) -> Self {
        self.poll_initial_delay = initial;
        self.poll_max_delay = max.max(initial);
        self
    }

    fn poll_delay_for(initial: Duration, max: Duration, attempt: u32) -> Duration {
        let max_millis = max.as_millis();
        let initial_millis = initial.as_millis();
        if max_millis == 0 || initial_millis == 0 {
            return Duration::ZERO;
        }

        let factor_shift = attempt.min(16);
        let factor = 1u128 << factor_shift;
        let delayed = initial_millis.saturating_mul(factor).min(max_millis);
        Duration::from_millis(delayed as u64)
    }

    /// Wait for response output to become available.
    pub async fn wait(self) -> Result<Response> {
        let DeferredResponsePoller {
            client,
            response_id,
            max_attempts,
            poll_initial_delay,
            poll_max_delay,
        } = self;

        let api = ResponsesApi::new(client);

        for attempt in 0..max_attempts {
            let response = api.get(&response_id).await?;
            if !response.output.is_empty() {
                return Ok(response);
            }

            if attempt + 1 < max_attempts {
                let delay = Self::poll_delay_for(poll_initial_delay, poll_max_delay, attempt);
                if !delay.is_zero() {
                    sleep(delay).await;
                }
            }
        }

        Err(Error::Timeout)
    }
}

/// Stateful chat helper on top of the Responses API.
#[derive(Debug, Clone)]
pub struct StatefulChat {
    client: XaiClient,
    model: String,
    messages: Vec<Message>,
    pending_tool_calls: Vec<ToolCall>,
}

impl StatefulChat {
    fn new(client: XaiClient, model: String) -> Self {
        Self {
            client,
            model,
            messages: Vec::new(),
            pending_tool_calls: Vec::new(),
        }
    }

    /// Append a message to the local chat state.
    pub fn append(&mut self, role: Role, content: impl Into<MessageContent>) -> &mut Self {
        self.messages.push(Message::new(role, content));
        self
    }

    /// Append a system message.
    pub fn append_system(&mut self, content: impl Into<String>) -> &mut Self {
        self.append(Role::System, content.into())
    }

    /// Append a user message.
    pub fn append_user(&mut self, content: impl Into<MessageContent>) -> &mut Self {
        self.append(Role::User, content)
    }

    /// Append an assistant message.
    pub fn append_assistant(&mut self, content: impl Into<String>) -> &mut Self {
        self.append(Role::Assistant, content.into())
    }

    /// Append a pre-built message.
    pub fn append_message(&mut self, message: Message) -> &mut Self {
        self.messages.push(message);
        self
    }

    /// Append a tool result message.
    pub fn append_tool_result(
        &mut self,
        tool_call_id: impl Into<String>,
        content: impl Into<String>,
    ) -> &mut Self {
        self.append_message(Message::tool(tool_call_id, content))
    }

    /// Get the current local message history.
    pub fn messages(&self) -> &[Message] {
        &self.messages
    }

    /// Get pending tool calls extracted from the latest sampled response.
    pub fn pending_tool_calls(&self) -> &[ToolCall] {
        &self.pending_tool_calls
    }

    /// Take and clear pending tool calls.
    pub fn take_pending_tool_calls(&mut self) -> Vec<ToolCall> {
        std::mem::take(&mut self.pending_tool_calls)
    }

    /// Clear the local message history.
    pub fn clear(&mut self) -> &mut Self {
        self.messages.clear();
        self.pending_tool_calls.clear();
        self
    }

    fn text_content_slice(content: &TextContent) -> &str {
        match content {
            TextContent::Text { text } => text,
            // Preserve refusal semantics in history by carrying refusal text.
            TextContent::Refusal { refusal } => refusal,
        }
    }

    fn merge_output_message_content(content: &[TextContent]) -> Option<String> {
        match content {
            [] => None,
            [single] => {
                let text = Self::text_content_slice(single);
                (!text.is_empty()).then(|| text.to_string())
            }
            _ => {
                let total_len: usize = content
                    .iter()
                    .map(|part| Self::text_content_slice(part).len())
                    .sum();
                if total_len == 0 {
                    return None;
                }

                let mut merged = String::with_capacity(total_len);
                for part in content {
                    merged.push_str(Self::text_content_slice(part));
                }
                Some(merged)
            }
        }
    }

    fn collect_response_semantics(response: &Response) -> (Vec<String>, Vec<ToolCall>) {
        let output_item_count = response.output.len();
        let top_level_call_count = response.tool_calls.as_ref().map_or(0, Vec::len);
        let mut assistant_messages = Vec::with_capacity(output_item_count);
        let mut pending_tool_calls = Vec::with_capacity(top_level_call_count + output_item_count);
        let mut seen_tool_call_ids: HashSet<String> =
            HashSet::with_capacity(top_level_call_count + output_item_count);

        // Prefer explicit top-level tool call payloads when present.
        if let Some(calls) = &response.tool_calls {
            for call in calls {
                if seen_tool_call_ids.insert(call.id.clone()) {
                    pending_tool_calls.push(call.clone());
                }
            }
        }

        for item in &response.output {
            match item {
                OutputItem::Message { content, .. } => {
                    if let Some(merged) = Self::merge_output_message_content(content) {
                        assistant_messages.push(merged);
                    }
                }
                OutputItem::FunctionCall { call } => {
                    if seen_tool_call_ids.insert(call.id.clone()) {
                        pending_tool_calls.push(call.clone());
                    }
                }
                OutputItem::CodeInterpreterCall { id, .. } => {
                    if seen_tool_call_ids.insert(id.clone()) {
                        pending_tool_calls.push(ToolCall {
                            id: id.clone(),
                            call_type: Some("code_interpreter".to_string()),
                            function: None,
                        });
                    }
                }
                OutputItem::WebSearchCall { id, .. } => {
                    if seen_tool_call_ids.insert(id.clone()) {
                        pending_tool_calls.push(ToolCall {
                            id: id.clone(),
                            call_type: Some("web_search".to_string()),
                            function: None,
                        });
                    }
                }
                OutputItem::XSearchCall { id, .. } => {
                    if seen_tool_call_ids.insert(id.clone()) {
                        pending_tool_calls.push(ToolCall {
                            id: id.clone(),
                            call_type: Some("x_search".to_string()),
                            function: None,
                        });
                    }
                }
            }
        }

        (assistant_messages, pending_tool_calls)
    }

    /// Sample a response using the current local message history.
    pub async fn sample(&self) -> Result<Response> {
        self.client
            .responses()
            .create(self.model.clone())
            .messages(self.messages.clone())
            .send()
            .await
    }

    /// Append assistant text from a response into local chat history.
    pub fn append_response_text(&mut self, response: &Response) -> &mut Self {
        let text = response.all_text();
        if !text.is_empty() {
            self.append_assistant(text);
        }
        self
    }

    /// Append response semantics to local history and pending tool calls.
    pub fn append_response_semantics(&mut self, response: &Response) -> &mut Self {
        let (assistant_messages, pending_tool_calls) = Self::collect_response_semantics(response);
        for text in assistant_messages {
            self.append_assistant(text);
        }
        self.pending_tool_calls = pending_tool_calls;
        self
    }

    /// Sample a response and append semantic carryover to local state.
    pub async fn sample_and_append(&mut self) -> Result<Response> {
        let response = self.sample().await?;
        self.append_response_semantics(&response);
        Ok(response)
    }

    /// Run a sampled chat turn with a tool handler loop (default max rounds).
    ///
    /// The loop runs `sample_and_append()` repeatedly. When pending tool calls are
    /// produced, the handler is invoked for each call and tool results are appended
    /// as `tool` messages before the next sampling round.
    pub async fn sample_with_tool_loop<H, Fut>(&mut self, handler: H) -> Result<Response>
    where
        H: FnMut(ToolCall) -> Fut,
        Fut: std::future::Future<Output = Result<String>>,
    {
        self.sample_with_tool_handler(DEFAULT_STATEFUL_TOOL_LOOP_MAX_ROUNDS, handler)
            .await
    }

    /// Run a sampled chat turn with a tool handler loop and explicit max rounds.
    pub async fn sample_with_tool_handler<H, Fut>(
        &mut self,
        max_rounds: u32,
        mut handler: H,
    ) -> Result<Response>
    where
        H: FnMut(ToolCall) -> Fut,
        Fut: std::future::Future<Output = Result<String>>,
    {
        let rounds = max_rounds.max(1);

        for _ in 0..rounds {
            let response = self.sample_and_append().await?;
            let pending = self.take_pending_tool_calls();
            if pending.is_empty() {
                return Ok(response);
            }

            for call in pending {
                let tool_call_id = call.id.clone();
                let tool_result = handler(call).await?;
                self.append_tool_result(tool_call_id, tool_result);
            }
        }

        Err(Error::Config(format!(
            "stateful chat tool loop exceeded max rounds ({rounds})"
        )))
    }

    /// Stream a response using the current local message history.
    pub async fn stream(&self) -> Result<StatefulChatStream> {
        let stream = self
            .client
            .responses()
            .create(self.model.clone())
            .messages(self.messages.clone())
            .stream()
            .await?;

        Ok(StatefulChatStream::new(stream))
    }
}

pin_project! {
    /// Stream wrapper that accumulates text deltas as chunks arrive.
    pub struct StatefulChatStream {
        #[pin]
        inner: ResponseStream,
        accumulated_text: String,
    }
}

/// A streamed chunk paired with accumulated text so far.
#[derive(Debug, Clone)]
pub struct AccumulatedChunk {
    /// The raw stream chunk from the API.
    pub chunk: StreamChunk,
    /// The full accumulated text after applying this chunk.
    pub accumulated_text: String,
}

impl StatefulChatStream {
    fn new(inner: ResponseStream) -> Self {
        Self {
            inner,
            accumulated_text: String::new(),
        }
    }

    /// Get accumulated text so far.
    pub fn accumulated_text(&self) -> &str {
        &self.accumulated_text
    }

    /// Consume the stream wrapper and return accumulated text.
    pub fn into_accumulated_text(self) -> String {
        self.accumulated_text
    }

    /// Read the next chunk along with an accumulated text snapshot.
    pub async fn next_with_accumulated(&mut self) -> Option<Result<AccumulatedChunk>> {
        use futures_util::future::poll_fn;

        poll_fn(|cx| Pin::new(&mut *self).poll_next(cx))
            .await
            .map(|result| {
                result.map(|chunk| AccumulatedChunk {
                    chunk,
                    accumulated_text: self.accumulated_text.clone(),
                })
            })
    }
}

impl Stream for StatefulChatStream {
    type Item = Result<StreamChunk>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        match this.inner.as_mut().poll_next(cx) {
            Poll::Ready(Some(Ok(chunk))) => {
                this.accumulated_text.push_str(chunk.delta());
                Poll::Ready(Some(Ok(chunk)))
            }
            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// Builder for creating a response request.
#[derive(Debug)]
pub struct CreateResponseBuilder {
    client: XaiClient,
    request: CreateResponseRequest,
    retry_policy_override: Option<RetryPolicy>,
}

#[derive(Debug, Clone, Serialize)]
struct CreateResponseRequest {
    model: String,
    input: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<Tool>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_choice: Option<ToolChoice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    response_format: Option<ResponseFormat>,
    #[serde(skip_serializing_if = "Option::is_none")]
    include: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    store: Option<bool>,
}

impl CreateResponseBuilder {
    fn new(client: XaiClient, model: String) -> Self {
        Self {
            client,
            request: CreateResponseRequest {
                model,
                input: Vec::new(),
                tools: None,
                tool_choice: None,
                temperature: None,
                top_p: None,
                max_tokens: None,
                stream: None,
                response_format: None,
                include: None,
                store: None,
            },
            retry_policy_override: None,
        }
    }

    fn retry_policy_mut(&mut self) -> &mut RetryPolicy {
        self.retry_policy_override
            .get_or_insert_with(|| self.client.retry_policy())
    }

    /// Add a message to the conversation.
    pub fn message(mut self, role: Role, content: impl Into<MessageContent>) -> Self {
        self.request.input.push(Message::new(role, content));
        self
    }

    /// Add a system message.
    pub fn system(self, content: impl Into<String>) -> Self {
        self.message(Role::System, content.into())
    }

    /// Add a user message.
    pub fn user(self, content: impl Into<MessageContent>) -> Self {
        self.message(Role::User, content)
    }

    /// Add an assistant message.
    pub fn assistant(self, content: impl Into<String>) -> Self {
        self.message(Role::Assistant, content.into())
    }

    /// Add a user message with an image.
    pub fn user_with_image(
        mut self,
        text: impl Into<String>,
        image_url: impl Into<String>,
    ) -> Self {
        let parts = vec![ContentPart::text(text), ContentPart::image_url(image_url)];
        self.request
            .input
            .push(Message::new(Role::User, MessageContent::Parts(parts)));
        self
    }

    /// Add pre-built messages.
    pub fn messages(mut self, messages: Vec<Message>) -> Self {
        self.request.input.extend(messages);
        self
    }

    /// Add a tool.
    pub fn tool(mut self, tool: Tool) -> Self {
        self.request.tools.get_or_insert_with(Vec::new).push(tool);
        self
    }

    /// Add multiple tools.
    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
        self.request
            .tools
            .get_or_insert_with(Vec::new)
            .extend(tools);
        self
    }

    /// Set the tool choice.
    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
        self.request.tool_choice = Some(choice);
        self
    }

    /// Set the temperature (0.0 - 2.0).
    pub fn temperature(mut self, temperature: f32) -> Self {
        self.request.temperature = Some(temperature.clamp(0.0, 2.0));
        self
    }

    /// Set top_p (0.0 - 1.0).
    pub fn top_p(mut self, top_p: f32) -> Self {
        self.request.top_p = Some(top_p.clamp(0.0, 1.0));
        self
    }

    /// Set the maximum number of tokens to generate.
    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.request.max_tokens = Some(max_tokens);
        self
    }

    /// Override retry attempts for this request only.
    pub fn max_retries(mut self, max_retries: u32) -> Self {
        self.retry_policy_mut().max_retries = max_retries;
        self
    }

    /// Disable retries for this request only.
    pub fn disable_retries(self) -> Self {
        self.max_retries(0)
    }

    /// Override retry backoff for this request only.
    pub fn retry_backoff(mut self, initial: Duration, max: Duration) -> Self {
        let policy = self.retry_policy_mut();
        policy.initial_backoff = initial;
        policy.max_backoff = max.max(initial);
        self
    }

    /// Override retry jitter factor (0.0 to 1.0) for this request only.
    pub fn retry_jitter(mut self, factor: f64) -> Self {
        self.retry_policy_mut().jitter_factor = factor.clamp(0.0, 1.0);
        self
    }

    /// Set the response format.
    pub fn response_format(mut self, format: ResponseFormat) -> Self {
        self.request.response_format = Some(format);
        self
    }

    /// Request JSON output.
    pub fn json_output(self) -> Self {
        self.response_format(ResponseFormat::json_object())
    }

    /// Include additional fields in the response.
    pub fn include(mut self, fields: Vec<String>) -> Self {
        self.request.include = Some(fields);
        self
    }

    /// Include inline citations.
    pub fn with_inline_citations(mut self) -> Self {
        self.request
            .include
            .get_or_insert_with(Vec::new)
            .push("inline_citations".to_string());
        self
    }

    /// Include verbose streaming output.
    pub fn with_verbose_streaming(mut self) -> Self {
        self.request
            .include
            .get_or_insert_with(Vec::new)
            .push("verbose_streaming".to_string());
        self
    }

    /// Store the response for later retrieval.
    pub fn store(mut self, store: bool) -> Self {
        self.request.store = Some(store);
        self
    }

    /// Send the request and get a response.
    pub async fn send(self) -> Result<Response> {
        let url = format!("{}/responses", self.client.base_url());

        let response = self
            .client
            .send_with_retry_policy(
                self.client.http().post(&url).json(&self.request),
                self.retry_policy_override,
            )
            .await?;

        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }

        Ok(response.json().await?)
    }

    /// Send the request and stream the response.
    pub async fn stream(mut self) -> Result<ResponseStream> {
        self.request.stream = Some(true);

        let url = format!("{}/responses", self.client.base_url());

        let response = self
            .client
            .send_with_retry_policy(
                self.client.http().post(&url).json(&self.request),
                self.retry_policy_override,
            )
            .await?;

        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }

        Ok(ResponseStream::new(response.bytes_stream()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use futures_util::{stream, StreamExt};
    use serde_json::json;
    use std::sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    };
    use wiremock::{
        matchers::{body_partial_json, method, path},
        Mock, MockServer, ResponseTemplate,
    };

    #[test]
    fn user_with_image_puts_text_before_image() {
        let client = XaiClient::new("test-key").unwrap();
        let api = ResponsesApi::new(client);
        let builder = api
            .create("grok-4")
            .user_with_image("describe this", "https://example.com/image.jpg");

        assert_eq!(builder.request.input.len(), 1);
        let msg = &builder.request.input[0];
        assert!(matches!(msg.role, Role::User));

        match &msg.content {
            MessageContent::Parts(parts) => {
                assert_eq!(parts.len(), 2);
                assert!(matches!(
                    &parts[0],
                    ContentPart::Text { text } if text == "describe this"
                ));
                assert!(matches!(
                    &parts[1],
                    ContentPart::ImageUrl { image_url } if image_url.url == "https://example.com/image.jpg"
                ));
            }
            _ => panic!("Expected multipart message content"),
        }
    }

    #[tokio::test]
    async fn create_builder_retry_override_enables_retries_when_client_retries_disabled() {
        let server = MockServer::start().await;
        let call_count = Arc::new(AtomicUsize::new(0));
        let responder_count = Arc::clone(&call_count);

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(move |_req: &wiremock::Request| {
                let count = responder_count.fetch_add(1, Ordering::SeqCst);
                if count == 0 {
                    ResponseTemplate::new(503).set_body_json(json!({
                        "error": {"message": "temporary", "type": "server_error"}
                    }))
                } else {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_retry",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{"type": "text", "text": "retry worked"}]
                        }],
                        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
                    }))
                }
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .disable_retries()
            .build()
            .unwrap();

        let response = client
            .responses()
            .create("grok-4")
            .user("hello")
            .max_retries(1)
            .retry_backoff(Duration::ZERO, Duration::ZERO)
            .send()
            .await
            .unwrap();

        assert_eq!(response.output_text().as_deref(), Some("retry worked"));
        assert_eq!(call_count.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn create_builder_disable_retries_overrides_client_retry_policy() {
        let server = MockServer::start().await;
        let call_count = Arc::new(AtomicUsize::new(0));
        let responder_count = Arc::clone(&call_count);

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(move |_req: &wiremock::Request| {
                responder_count.fetch_add(1, Ordering::SeqCst);
                ResponseTemplate::new(503).set_body_json(json!({
                    "error": {"message": "still unavailable", "type": "server_error"}
                }))
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .max_retries(2)
            .retry_backoff(Duration::ZERO, Duration::ZERO)
            .build()
            .unwrap();

        let err = client
            .responses()
            .create("grok-4")
            .user("hello")
            .disable_retries()
            .send()
            .await
            .unwrap_err();

        assert!(matches!(err, Error::Api { status: 503, .. }));
        assert_eq!(call_count.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn responses_api_get_propagates_api_error() {
        let server = MockServer::start().await;

        Mock::given(method("GET"))
            .and(path("/responses/missing"))
            .respond_with(ResponseTemplate::new(404).set_body_json(json!({
                "error": {"message": "response not found"}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = client.responses().get("missing").await.unwrap_err();
        match err {
            Error::Api {
                status, message, ..
            } => {
                assert_eq!(status, 404);
                assert_eq!(message, "response not found");
            }
            _ => panic!("expected Error::Api"),
        }
    }

    #[tokio::test]
    async fn responses_api_delete_propagates_api_error() {
        let server = MockServer::start().await;

        Mock::given(method("DELETE"))
            .and(path("/responses/missing"))
            .respond_with(ResponseTemplate::new(404).set_body_json(json!({
                "error": {"message": "response not found"}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = client.responses().delete("missing").await.unwrap_err();
        match err {
            Error::Api {
                status, message, ..
            } => {
                assert_eq!(status, 404);
                assert_eq!(message, "response not found");
            }
            _ => panic!("expected Error::Api"),
        }
    }

    #[tokio::test]
    async fn create_response_builder_send_propagates_api_error() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .and(body_partial_json(json!({
                "model": "grok-4",
                "input": [{"role": "user", "content": "error"}]
            })))
            .respond_with(ResponseTemplate::new(503).set_body_json(json!({
                "error": {"message": "service unavailable"}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = client
            .responses()
            .create("grok-4")
            .user("error")
            .send()
            .await
            .unwrap_err();

        match err {
            Error::Api {
                status, message, ..
            } => {
                assert_eq!(status, 503);
                assert_eq!(message, "service unavailable");
            }
            _ => panic!("expected Error::Api"),
        }
    }

    #[tokio::test]
    async fn create_response_builder_stream_propagates_api_error() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .and(body_partial_json(json!({
                "model": "grok-4",
                "input": [{"role": "user", "content": "stream error"}]
            })))
            .respond_with(ResponseTemplate::new(503).set_body_json(json!({
                "error": {"message": "stream unavailable"}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = match client
            .responses()
            .create("grok-4")
            .user("stream error")
            .stream()
            .await
        {
            Ok(_) => panic!("expected stream creation to fail"),
            Err(err) => err,
        };

        match err {
            Error::Api {
                status, message, ..
            } => {
                assert_eq!(status, 503);
                assert_eq!(message, "stream unavailable");
            }
            _ => panic!("expected Error::Api"),
        }
    }

    #[tokio::test]
    async fn create_response_builder_stream_returns_stream() {
        let server = MockServer::start().await;

        let payload = concat!(
            "data: {\"type\":\"response.output_item.delta\",\"delta\":{\"content\":[{\"type\":\"text\",\"text\":\"done\"}]}}\n\n",
            "data: {\"type\":\"response.done\",\"response\":{\"id\":\"resp_stream\",\"model\":\"grok-4\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"done\"}]}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}}\n\n",
            "data: [DONE]\n\n"
        );

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_string(payload))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut stream = client
            .responses()
            .create("grok-4")
            .user("stream response")
            .stream()
            .await
            .unwrap();

        let first = stream.next().await.unwrap().unwrap();
        assert!(!first.done);
        assert_eq!(first.delta(), "done");
        let done = stream.next().await.unwrap().unwrap();
        assert!(done.done);
        assert!(done.response.is_some());
        assert_eq!(done.response.unwrap().id, "resp_stream");
    }

    #[tokio::test]
    async fn create_get_poll_delete_roundtrip() {
        let server = MockServer::start().await;
        let post_count = Arc::new(AtomicUsize::new(0));
        let get_count = Arc::new(AtomicUsize::new(0));
        let delete_count = Arc::new(AtomicUsize::new(0));

        let post_count_for_responder = Arc::clone(&post_count);
        Mock::given(method("POST"))
            .and(path("/responses"))
            .and(body_partial_json(json!( {
                "model": "grok-4",
                "input": [
                    {"role": "system", "content": "Roundtrip test system"},
                    {"role": "user", "content": "What is 1+1?"}
                ]
            })))
            .respond_with(move |_req: &wiremock::Request| {
                post_count_for_responder.fetch_add(1, Ordering::SeqCst);
                ResponseTemplate::new(200).set_body_json(json!({
                    "id": "resp_roundtrip",
                    "model": "grok-4",
                    "output": [],
                    "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}
                }))
            })
            .mount(&server)
            .await;

        let get_count_for_responder = Arc::clone(&get_count);
        Mock::given(method("GET"))
            .and(path("/responses/resp_roundtrip"))
            .respond_with(move |_req: &wiremock::Request| {
                let count = get_count_for_responder.fetch_add(1, Ordering::SeqCst);
                if count == 0 {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_roundtrip",
                        "model": "grok-4",
                        "output": [],
                        "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}
                    }))
                } else {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_roundtrip",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{"type": "text", "text": "hello"}]
                        }],
                        "usage": {"prompt_tokens": 11, "completion_tokens": 1, "total_tokens": 12}
                    }))
                }
            })
            .mount(&server)
            .await;

        let delete_count_for_responder = Arc::clone(&delete_count);
        Mock::given(method("DELETE"))
            .and(path("/responses/resp_roundtrip"))
            .respond_with(move |_req: &wiremock::Request| {
                delete_count_for_responder.fetch_add(1, Ordering::SeqCst);
                ResponseTemplate::new(200).set_body_json(json!({"id": "resp_roundtrip"}))
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let create_response = client
            .responses()
            .create("grok-4")
            .system("Roundtrip test system")
            .user("What is 1+1?")
            .send()
            .await
            .unwrap();
        assert_eq!(create_response.id, "resp_roundtrip");
        assert!(create_response.output_text().is_none());

        let get_response = client.responses().get("resp_roundtrip").await.unwrap();
        assert_eq!(get_response.id, "resp_roundtrip");
        assert!(get_response.output_text().is_none());

        let polled_response = client
            .responses()
            .poll_until_ready("resp_roundtrip", 2)
            .await
            .unwrap();
        assert_eq!(polled_response.output_text().as_deref(), Some("hello"));

        client.responses().delete("resp_roundtrip").await.unwrap();

        assert_eq!(post_count.load(Ordering::SeqCst), 1);
        assert_eq!(get_count.load(Ordering::SeqCst), 2);
        assert_eq!(delete_count.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn stateful_chat_append_and_sample_sends_all_messages() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .and(body_partial_json(json!({
                "model": "grok-4",
                "input": [
                    {"role": "system", "content": "You are helpful."},
                    {"role": "user", "content": "Hello"}
                ]
            })))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_123",
                "model": "grok-4",
                "output": [{
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "text", "text": "Hi there!"}]
                }],
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 3,
                    "total_tokens": 13
                }
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.").append_user("Hello");

        assert_eq!(chat.messages().len(), 2);

        let response = chat.sample().await.unwrap();
        assert_eq!(response.output_text().as_deref(), Some("Hi there!"));
    }

    #[tokio::test]
    async fn stateful_chat_sample_and_append_updates_local_history() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_124",
                "model": "grok-4",
                "output": [{
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "text", "text": "History updated"}]
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.").append_user("Hello");

        let response = chat.sample_and_append().await.unwrap();
        assert_eq!(response.output_text().as_deref(), Some("History updated"));
        assert_eq!(chat.messages().len(), 3);
        assert!(chat.pending_tool_calls().is_empty());

        let last = chat.messages().last().unwrap();
        assert!(matches!(last.role, Role::Assistant));
        assert_eq!(last.content.as_text(), Some("History updated"));
    }

    #[tokio::test]
    async fn stateful_chat_sample_and_append_skips_empty_text() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_125",
                "model": "grok-4",
                "output": [],
                "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.").append_user("Hello");

        let response = chat.sample_and_append().await.unwrap();
        assert!(response.output_text().is_none());
        assert_eq!(chat.messages().len(), 2);
        assert!(chat.pending_tool_calls().is_empty());
    }

    #[tokio::test]
    async fn stateful_chat_append_response_text_appends_non_empty_and_skips_empty() {
        let mut chat = StatefulChat::new(XaiClient::new("test-key").unwrap(), "grok-4".to_string());
        chat.append_system("You are helpful.").append_user("Hello");

        let non_empty = Response {
            id: "resp_non_empty".to_string(),
            model: "grok-4".to_string(),
            output: vec![OutputItem::Message {
                role: Role::Assistant,
                content: vec![
                    TextContent::Text {
                        text: "Hello ".to_string(),
                    },
                    TextContent::Text {
                        text: "again".to_string(),
                    },
                ],
            }],
            usage: Default::default(),
            citations: None,
            inline_citations: None,
            server_side_tool_usage: None,
            tool_calls: None,
            system_fingerprint: None,
        };

        let empty = Response {
            id: "resp_empty".to_string(),
            model: "grok-4".to_string(),
            output: vec![OutputItem::Message {
                role: Role::Assistant,
                content: vec![TextContent::Refusal {
                    refusal: "policy block".to_string(),
                }],
            }],
            usage: Default::default(),
            citations: None,
            inline_citations: None,
            server_side_tool_usage: None,
            tool_calls: None,
            system_fingerprint: None,
        };

        chat.append_response_text(&non_empty);
        chat.append_response_text(&empty);

        assert_eq!(chat.messages().len(), 3);
        assert_eq!(
            chat.messages().last().unwrap().content.as_text(),
            Some("Hello again")
        );
    }

    #[tokio::test]
    async fn stateful_chat_sample_and_append_carries_refusal_text() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_126",
                "model": "grok-4",
                "output": [{
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "refusal", "refusal": "I can't help with that."}]
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.")
            .append_user("Do something unsafe");

        let response = chat.sample_and_append().await.unwrap();
        assert!(response.output_text().is_none());
        assert_eq!(chat.messages().len(), 3);
        assert!(chat.pending_tool_calls().is_empty());
        assert_eq!(
            chat.messages().last().unwrap().content.as_text(),
            Some("I can't help with that.")
        );
    }

    #[tokio::test]
    async fn stateful_chat_sample_and_append_captures_tool_calls() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_127",
                "model": "grok-4",
                "output": [{
                    "type": "function_call",
                    "id": "call_weather",
                    "function": {
                        "name": "get_weather",
                        "arguments": "{\"location\":\"Paris\"}"
                    }
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.")
            .append_user("What's weather in Paris?");

        let _ = chat.sample_and_append().await.unwrap();

        // No assistant text message to append, only a pending tool call to resolve.
        assert_eq!(chat.messages().len(), 2);
        assert_eq!(chat.pending_tool_calls().len(), 1);
        assert_eq!(chat.pending_tool_calls()[0].id, "call_weather");
        assert_eq!(
            chat.pending_tool_calls()[0]
                .function
                .as_ref()
                .map(|f| f.name.as_str()),
            Some("get_weather")
        );

        let pending = chat.take_pending_tool_calls();
        assert_eq!(pending.len(), 1);
        assert!(chat.pending_tool_calls().is_empty());

        chat.append_tool_result("call_weather", r#"{"temperature": 72}"#);
        assert_eq!(chat.messages().len(), 3);
        assert!(matches!(chat.messages().last().unwrap().role, Role::Tool));
    }

    #[tokio::test]
    async fn stateful_chat_sample_with_tool_handler_resolves_tool_loop() {
        let server = MockServer::start().await;
        let call_count = Arc::new(AtomicUsize::new(0));
        let responder_count = Arc::clone(&call_count);

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(move |_req: &wiremock::Request| {
                let count = responder_count.fetch_add(1, Ordering::SeqCst);
                if count == 0 {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_tool_1",
                        "model": "grok-4",
                        "output": [{
                            "type": "function_call",
                            "id": "call_weather",
                            "function": {
                                "name": "get_weather",
                                "arguments": "{\"location\":\"Paris\"}"
                            }
                        }],
                        "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
                    }))
                } else {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_tool_2",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{"type": "text", "text": "Weather is 72F"}]
                        }],
                        "usage": {"prompt_tokens": 12, "completion_tokens": 3, "total_tokens": 15}
                    }))
                }
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.")
            .append_user("What's the weather in Paris?");

        let response = chat
            .sample_with_tool_handler(3, |call| async move {
                if call.id == "call_weather" {
                    Ok(r#"{"temperature": 72}"#.to_string())
                } else {
                    Ok("{}".to_string())
                }
            })
            .await
            .unwrap();

        assert_eq!(response.output_text().as_deref(), Some("Weather is 72F"));
        assert_eq!(call_count.load(Ordering::SeqCst), 2);
        assert!(chat.pending_tool_calls().is_empty());
        assert_eq!(chat.messages().len(), 4);
        assert!(matches!(chat.messages()[2].role, Role::Tool));
        assert!(matches!(chat.messages()[3].role, Role::Assistant));
    }

    #[tokio::test]
    async fn stateful_chat_sample_with_tool_handler_errors_when_rounds_exhausted() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_tool_loop",
                "model": "grok-4",
                "output": [{
                    "type": "function_call",
                    "id": "call_loop",
                    "function": {
                        "name": "looping_tool",
                        "arguments": "{}"
                    }
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.")
            .append_user("Trigger loop");

        let err = chat
            .sample_with_tool_handler(1, |_call| async move { Ok("{}".to_string()) })
            .await
            .unwrap_err();

        match err {
            Error::Config(message) => {
                assert!(message.contains("stateful chat tool loop exceeded max rounds (1)"))
            }
            _ => panic!("expected Error::Config"),
        }
    }

    #[tokio::test]
    async fn stateful_chat_sample_with_tool_handler_propagates_handler_error() {
        let server = MockServer::start().await;

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_tool_error",
                "model": "grok-4",
                "output": [{
                    "type": "function_call",
                    "id": "call_fail",
                    "function": {
                        "name": "failing_tool",
                        "arguments": "{}"
                    }
                }],
                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.")
            .append_user("Trigger tool error");

        let err = chat
            .sample_with_tool_handler(3, |_call| async move {
                Err(Error::Config("tool handler failed".to_string()))
            })
            .await
            .unwrap_err();

        match err {
            Error::Config(message) => assert!(message.contains("tool handler failed")),
            _ => panic!("expected Error::Config"),
        }
    }

    #[tokio::test]
    async fn stateful_chat_stream_accumulates_deltas() {
        let payload = concat!(
            "data: {\"type\":\"response.output_item.delta\",\"delta\":{\"content\":[{\"type\":\"text\",\"text\":\"Hel\"}]}}\n\n",
            "data: {\"type\":\"response.output_item.delta\",\"delta\":{\"content\":[{\"type\":\"text\",\"text\":\"lo\"}]}}\n\n",
            "data: [DONE]\n\n"
        );

        let chunks: Vec<std::result::Result<Bytes, reqwest::Error>> =
            vec![Ok(Bytes::from(payload.to_string()))];
        let raw_stream = stream::iter(chunks);
        let response_stream = ResponseStream::new(raw_stream);
        let mut stream = StatefulChatStream::new(response_stream);

        let first = stream.next().await.unwrap().unwrap();
        assert_eq!(first.delta(), "Hel");
        assert_eq!(stream.accumulated_text(), "Hel");

        let second = stream.next().await.unwrap().unwrap();
        assert_eq!(second.delta(), "lo");
        assert_eq!(stream.accumulated_text(), "Hello");

        let done = stream.next().await.unwrap().unwrap();
        assert!(done.done);
        assert_eq!(stream.accumulated_text(), "Hello");
    }

    #[tokio::test]
    async fn stateful_chat_stream_next_with_accumulated_returns_snapshot() {
        let payload = concat!(
            "data: {\"type\":\"response.output_item.delta\",\"delta\":{\"content\":[{\"type\":\"text\",\"text\":\"Hel\"}]}}\n\n",
            "data: {\"type\":\"response.output_item.delta\",\"delta\":{\"content\":[{\"type\":\"text\",\"text\":\"lo\"}]}}\n\n",
            "data: [DONE]\n\n"
        );

        let chunks: Vec<std::result::Result<Bytes, reqwest::Error>> =
            vec![Ok(Bytes::from(payload.to_string()))];
        let raw_stream = stream::iter(chunks);
        let response_stream = ResponseStream::new(raw_stream);
        let mut stream = StatefulChatStream::new(response_stream);

        let first = stream.next_with_accumulated().await.unwrap().unwrap();
        assert_eq!(first.chunk.delta(), "Hel");
        assert_eq!(first.accumulated_text, "Hel");

        let second = stream.next_with_accumulated().await.unwrap().unwrap();
        assert_eq!(second.chunk.delta(), "lo");
        assert_eq!(second.accumulated_text, "Hello");

        let done = stream.next_with_accumulated().await.unwrap().unwrap();
        assert!(done.chunk.done);
        assert_eq!(done.accumulated_text, "Hello");
    }

    #[tokio::test]
    async fn deferred_poller_returns_when_response_has_output() {
        let server = MockServer::start().await;
        let call_count = Arc::new(AtomicUsize::new(0));
        let call_count_for_responder = Arc::clone(&call_count);

        Mock::given(method("GET"))
            .and(path("/responses/resp_deferred"))
            .respond_with(move |_req: &wiremock::Request| {
                let current = call_count_for_responder.fetch_add(1, Ordering::SeqCst);
                if current == 0 {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_deferred",
                        "model": "grok-4",
                        "output": [],
                        "usage": {"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}
                    }))
                } else {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_deferred",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{"type": "text", "text": "ready"}]
                        }],
                        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
                    }))
                }
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let response = client
            .responses()
            .deferred("resp_deferred")
            .max_attempts(3)
            .wait()
            .await
            .unwrap();

        assert_eq!(response.output_text().as_deref(), Some("ready"));
        assert_eq!(call_count.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn deferred_poller_times_out_when_output_never_arrives() {
        let server = MockServer::start().await;

        Mock::given(method("GET"))
            .and(path("/responses/resp_timeout"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "id": "resp_timeout",
                "model": "grok-4",
                "output": [],
                "usage": {"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}
            })))
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let err = client
            .responses()
            .deferred("resp_timeout")
            .max_attempts(2)
            .wait()
            .await
            .unwrap_err();

        assert!(matches!(err, Error::Timeout));
    }

    #[test]
    fn deferred_poller_poll_delay_returns_zero_when_interval_is_zero() {
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(Duration::ZERO, Duration::from_millis(200), 0),
            Duration::ZERO
        );
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(Duration::from_millis(200), Duration::ZERO, 0),
            Duration::ZERO
        );
    }

    #[tokio::test]
    async fn stateful_chat_sample_with_tool_handler_minimum_round_is_one() {
        let server = MockServer::start().await;
        let call_count = Arc::new(AtomicUsize::new(0));
        let count_for_responses = Arc::clone(&call_count);
        let handler_invocations = Arc::new(AtomicUsize::new(0));
        let handler_invocations_for_handler = Arc::clone(&handler_invocations);

        Mock::given(method("POST"))
            .and(path("/responses"))
            .respond_with(move |_req: &wiremock::Request| {
                let current = count_for_responses.fetch_add(1, Ordering::SeqCst);
                if current == 0 {
                    ResponseTemplate::new(200).set_body_json(json!( {
                        "id": "resp_tool_1",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{ "type": "text", "text": "Recovered" }]
                        }],
                        "usage": {"prompt_tokens": 12, "completion_tokens": 3, "total_tokens": 15}
                    }))
                } else {
                    panic!("unexpected second request")
                }
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let mut chat = client.responses().chat("grok-4");
        chat.append_system("You are helpful.").append_user("Hello");

        let response = chat
            .sample_with_tool_handler(0, move |_call| {
                let handler_invocations_for_handler = Arc::clone(&handler_invocations_for_handler);
                async move {
                    handler_invocations_for_handler.fetch_add(1, Ordering::SeqCst);
                    Ok(r#"{"value":42}"#.to_string())
                }
            })
            .await
            .unwrap();

        assert_eq!(response.output_text().as_deref(), Some("Recovered"));
        assert_eq!(call_count.load(Ordering::SeqCst), 1);
        assert_eq!(handler_invocations.load(Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn deferred_poller_wait_uses_poll_delay_between_attempts() {
        let server = MockServer::start().await;
        let request_count = Arc::new(AtomicUsize::new(0));
        let count_for_handler = Arc::clone(&request_count);

        Mock::given(method("GET"))
            .and(path("/responses/resp_poll_delay"))
            .respond_with(move |_req: &wiremock::Request| {
                let current = count_for_handler.fetch_add(1, Ordering::SeqCst);
                if current == 0 {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_poll_delay",
                        "model": "grok-4",
                        "output": [],
                        "usage": {"prompt_tokens": 4, "completion_tokens": 0, "total_tokens": 4}
                    }))
                } else {
                    ResponseTemplate::new(200).set_body_json(json!({
                        "id": "resp_poll_delay",
                        "model": "grok-4",
                        "output": [{
                            "type": "message",
                            "role": "assistant",
                            "content": [{"type": "text", "text": "ready"}]
                        }],
                        "usage": {"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}
                    }))
                }
            })
            .mount(&server)
            .await;

        let client = XaiClient::builder()
            .api_key("test-key")
            .base_url(server.uri())
            .build()
            .unwrap();

        let started_at = std::time::Instant::now();
        let response = client
            .responses()
            .deferred("resp_poll_delay")
            .poll_interval(Duration::from_millis(80))
            .max_attempts(2)
            .wait()
            .await
            .unwrap();
        let elapsed = started_at.elapsed();

        assert_eq!(response.output_text().as_deref(), Some("ready"));
        assert_eq!(request_count.load(Ordering::SeqCst), 2);
        assert!(
            elapsed >= Duration::from_millis(60),
            "expected exponential backoff delay to be observed"
        );
        assert!(
            elapsed < Duration::from_millis(500),
            "expected delay test to stay bounded"
        );
    }

    #[test]
    fn stateful_chat_merge_output_message_content_and_collect_response_semantics() {
        let text_output = vec![
            TextContent::Text {
                text: "Hello ".to_string(),
            },
            TextContent::Refusal {
                refusal: "Blocked".to_string(),
            },
        ];
        assert_eq!(
            StatefulChat::merge_output_message_content(&text_output),
            Some("Hello Blocked".to_string())
        );
        assert_eq!(StatefulChat::merge_output_message_content(&[]), None);

        let shared_tool_call_id = "shared-call".to_string();
        let response = Response {
            id: "resp_semantics".to_string(),
            model: "grok-4".to_string(),
            output: vec![
                OutputItem::Message {
                    role: Role::Assistant,
                    content: vec![
                        TextContent::Text {
                            text: "Part 1 ".to_string(),
                        },
                        TextContent::Text {
                            text: "Part 2".to_string(),
                        },
                    ],
                },
                OutputItem::FunctionCall {
                    call: ToolCall {
                        id: shared_tool_call_id.clone(),
                        call_type: Some("function".to_string()),
                        function: None,
                    },
                },
                OutputItem::FunctionCall {
                    call: ToolCall {
                        id: "function_call".to_string(),
                        call_type: Some("function".to_string()),
                        function: None,
                    },
                },
                OutputItem::CodeInterpreterCall {
                    id: "ci_call".to_string(),
                    code: None,
                    outputs: None,
                },
                OutputItem::WebSearchCall {
                    id: "web_call".to_string(),
                    results: None,
                },
                OutputItem::XSearchCall {
                    id: "x_call".to_string(),
                    results: None,
                },
            ],
            usage: Default::default(),
            citations: None,
            inline_citations: None,
            server_side_tool_usage: None,
            tool_calls: Some(vec![ToolCall {
                id: shared_tool_call_id.clone(),
                call_type: Some("function".to_string()),
                function: None,
            }]),
            system_fingerprint: None,
        };

        let (assistant_messages, pending_tool_calls) =
            StatefulChat::collect_response_semantics(&response);
        assert_eq!(assistant_messages, vec!["Part 1 Part 2".to_string()]);
        assert_eq!(pending_tool_calls.len(), 5);
        assert_eq!(pending_tool_calls[0].id, shared_tool_call_id);
        assert_eq!(pending_tool_calls[1].id, "function_call");
        assert_eq!(pending_tool_calls[1].call_type.as_deref(), Some("function"));
        assert_eq!(pending_tool_calls[2].id, "ci_call");
        assert_eq!(
            pending_tool_calls[2].call_type.as_deref(),
            Some("code_interpreter")
        );
        assert_eq!(pending_tool_calls[3].id, "web_call");
        assert_eq!(
            pending_tool_calls[3].call_type.as_deref(),
            Some("web_search")
        );
        assert_eq!(pending_tool_calls[4].id, "x_call");
        assert_eq!(pending_tool_calls[4].call_type.as_deref(), Some("x_search"));
    }

    #[test]
    fn deferred_poller_poll_interval_sets_fixed_delay() {
        let poller = DeferredResponsePoller::new(
            XaiClient::new("test-key").unwrap(),
            "resp_123".to_string(),
        )
        .poll_interval(Duration::from_millis(250));

        assert_eq!(poller.poll_initial_delay, Duration::from_millis(250));
        assert_eq!(poller.poll_max_delay, Duration::from_millis(250));
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                0
            ),
            Duration::from_millis(250)
        );
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                5
            ),
            Duration::from_millis(250)
        );
    }

    #[test]
    fn deferred_poller_poll_backoff_is_exponential_and_capped() {
        let poller = DeferredResponsePoller::new(
            XaiClient::new("test-key").unwrap(),
            "resp_123".to_string(),
        )
        .poll_backoff(Duration::from_millis(100), Duration::from_millis(300));

        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                0
            ),
            Duration::from_millis(100)
        );
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                1
            ),
            Duration::from_millis(200)
        );
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                2
            ),
            Duration::from_millis(300)
        );
        assert_eq!(
            DeferredResponsePoller::poll_delay_for(
                poller.poll_initial_delay,
                poller.poll_max_delay,
                3
            ),
            Duration::from_millis(300)
        );
    }
}