llm-stack 0.7.0

//! Conversation primitives: messages, content blocks, and responses.
//!
//! A conversation is a sequence of [`ChatMessage`]s, each carrying a
//! [`ChatRole`] and one or more [`ContentBlock`]s. Providers return a
//! [`ChatResponse`] containing the assistant's reply along with token
//! usage and stop-reason metadata.
//!
//! # Content model
//!
//! Every message holds a `Vec<ContentBlock>` rather than a plain string.
//! This uniform representation handles text, images, tool calls, tool
//! results, and reasoning traces without special-casing:
//!
//! ```rust
//! use llm_stack::{ChatMessage, ContentBlock};
//! use llm_stack::chat::ChatRole;
//!
//! // Simple text message
//! let msg = ChatMessage::user("Hello, world!");
//!
//! // Mixed content (text + image) in a single message
//! let mixed = ChatMessage {
//!     role: ChatRole::User,
//!     content: vec![
//!         ContentBlock::Text("What's in this image?".into()),
//!         ContentBlock::Image {
//!             media_type: "image/png".into(),
//!             data: llm_stack::chat::ImageSource::Base64("...".into()),
//!         },
//!     ],
//! };
//! ```

use std::collections::HashMap;
use std::fmt;

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::usage::Usage;

/// The role of a participant in a conversation.
///
/// Most conversations follow the pattern `System → User → Assistant`,
/// with [`ChatRole::Tool`] appearing only when the assistant invokes a
/// tool and the caller feeds back the result.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ChatRole {
    /// Sets behavioral instructions for the assistant. Typically the
    /// first message in a conversation; not all providers surface it as
    /// a discrete message (some use a separate `system` parameter).
    System,
    /// A message from the human user.
    User,
    /// A message generated by the model.
    Assistant,
    /// A tool result fed back to the model after it made a tool call.
    Tool,
}

impl fmt::Display for ChatRole {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::System => f.write_str("system"),
            Self::User => f.write_str("user"),
            Self::Assistant => f.write_str("assistant"),
            Self::Tool => f.write_str("tool"),
        }
    }
}

/// A single message in a conversation.
///
/// Use the convenience constructors for common cases:
///
/// ```rust
/// use llm_stack::ChatMessage;
///
/// let user  = ChatMessage::user("What is 2+2?");
/// let asst  = ChatMessage::assistant("4");
/// let sys   = ChatMessage::system("You are a math tutor.");
/// let tool  = ChatMessage::tool_result("call_123", "4");
/// ```
///
/// For multi-block content (images, multiple tool results), construct
/// the `content` vec directly.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatMessage {
    /// Who sent this message.
    pub role: ChatRole,
    /// The content blocks that make up the message body.
    pub content: Vec<ContentBlock>,
}

impl ChatMessage {
    /// Creates a message with a single [`ContentBlock::Text`] block.
    pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
        Self {
            role,
            content: vec![ContentBlock::Text(text.into())],
        }
    }

    /// Shorthand for a [`ChatRole::User`] text message.
    pub fn user(text: impl Into<String>) -> Self {
        Self::text(ChatRole::User, text)
    }

    /// Shorthand for a [`ChatRole::Assistant`] text message.
    pub fn assistant(text: impl Into<String>) -> Self {
        Self::text(ChatRole::Assistant, text)
    }

    /// Shorthand for a [`ChatRole::System`] text message.
    pub fn system(text: impl Into<String>) -> Self {
        Self::text(ChatRole::System, text)
    }

    /// Creates a [`ChatRole::Tool`] message with a single successful
    /// [`ToolResult`].
    ///
    /// Use [`ChatMessage::tool_error`] when the tool invocation failed.
    pub fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: ChatRole::Tool,
            content: vec![ContentBlock::ToolResult(ToolResult {
                tool_call_id: tool_call_id.into(),
                content: content.into(),
                is_error: false,
            })],
        }
    }

    /// Creates a [`ChatRole::Tool`] message indicating the tool call
    /// failed. The `content` should describe the error.
    pub fn tool_error(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: ChatRole::Tool,
            content: vec![ContentBlock::ToolResult(ToolResult {
                tool_call_id: tool_call_id.into(),
                content: content.into(),
                is_error: true,
            })],
        }
    }

    /// Returns `true` if the message has zero content blocks.
    ///
    /// This checks the block count only — a message containing a single
    /// `ContentBlock::Text("")` returns `false`. Use this to detect
    /// structurally empty messages (no blocks at all).
    pub fn is_empty(&self) -> bool {
        self.content.is_empty()
    }

    /// Serializes the message to a stable JSON format.
    ///
    /// This format is documented and versioned for persistence use cases
    /// (conversation history, checkpointing, etc.). The format is guaranteed
    /// to be backward compatible within the same major version.
    ///
    /// # Stability Guarantees
    ///
    /// - Field names are stable within major versions
    /// - New optional fields may be added in minor versions
    /// - The output is deterministic for the same input
    ///
    /// # Example
    ///
    /// ```rust
    /// use llm_stack::ChatMessage;
    ///
    /// let msg = ChatMessage::user("Hello, world!");
    /// let json = msg.to_json().expect("serialization should succeed");
    /// assert!(json["content"][0]["text"].as_str() == Some("Hello, world!"));
    /// ```
    pub fn to_json(&self) -> Result<Value, serde_json::Error> {
        serde_json::to_value(self)
    }

    /// Deserializes a message from the stable JSON format.
    ///
    /// Accepts JSON produced by [`to_json`](Self::to_json). This method is
    /// the inverse of `to_json` and maintains roundtrip fidelity.
    ///
    /// # Example
    ///
    /// ```rust
    /// use llm_stack::ChatMessage;
    /// use serde_json::json;
    ///
    /// let json = json!({
    ///     "role": "User",
    ///     "content": [{"text": "Hello!"}]
    /// });
    /// let msg = ChatMessage::from_json(&json).expect("valid message");
    /// assert_eq!(msg.role, llm_stack::chat::ChatRole::User);
    /// ```
    pub fn from_json(value: &Value) -> Result<Self, serde_json::Error> {
        serde_json::from_value(value.clone())
    }

    /// Deserializes a message from the stable JSON format, consuming the value.
    ///
    /// This is more efficient than [`from_json`](Self::from_json) when you have
    /// an owned `Value` that you don't need afterward, as it avoids cloning.
    ///
    /// # Example
    ///
    /// ```rust
    /// use llm_stack::ChatMessage;
    /// use serde_json::json;
    ///
    /// let json = json!({
    ///     "role": "User",
    ///     "content": [{"text": "Hello!"}]
    /// });
    /// let msg = ChatMessage::from_json_owned(json).expect("valid message");
    /// assert_eq!(msg.role, llm_stack::chat::ChatRole::User);
    /// ```
    pub fn from_json_owned(value: Value) -> Result<Self, serde_json::Error> {
        serde_json::from_value(value)
    }
}

/// An individual piece of content within a [`ChatMessage`].
///
/// Messages are composed of one or more blocks, allowing a single
/// message to carry text alongside images, tool invocations, or
/// reasoning traces.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum ContentBlock {
    /// Plain text content.
    Text(String),
    /// An inline image, either base64-encoded or referenced by URL.
    Image {
        /// MIME type, e.g. `"image/png"` or `"image/jpeg"`.
        media_type: String,
        /// The image payload.
        data: ImageSource,
    },
    /// A tool invocation requested by the assistant.
    ToolCall(ToolCall),
    /// The result of executing a tool, sent back to the model.
    ToolResult(ToolResult),
    /// An internal reasoning step (chain-of-thought) emitted by the
    /// model. Not all providers support this.
    Reasoning {
        /// The reasoning text.
        content: String,
    },
}

/// How an image is supplied to the model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ImageSource {
    /// Raw image bytes encoded as a base64 string.
    Base64(String),
    /// A publicly accessible URL pointing to the image.
    ///
    /// Stores a parsed [`url::Url`] that is guaranteed to be valid.
    /// Use [`ImageSource::from_url`] to construct from a string.
    Url(url::Url),
}

impl ImageSource {
    /// Creates a [`Url`](Self::Url) variant after parsing and
    /// normalizing the URL.
    ///
    /// Returns the parse error if the string is not a valid URL.
    pub fn from_url(url: impl AsRef<str>) -> Result<Self, url::ParseError> {
        let parsed = url::Url::parse(url.as_ref())?;
        Ok(Self::Url(parsed))
    }
}

/// A tool invocation requested by the assistant.
///
/// After receiving a `ToolCall`, the caller should execute the named
/// tool with the given `arguments`, then feed the result back as a
/// [`ChatMessage::tool_result`] (or [`ChatMessage::tool_error`]).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolCall {
    /// Provider-assigned identifier that links this call to its result.
    pub id: String,
    /// The name of the tool to invoke.
    pub name: String,
    /// JSON arguments to pass to the tool.
    pub arguments: Value,
}

impl std::fmt::Display for ToolCall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}({})", self.name, self.id)
    }
}

/// The result of executing a tool, returned to the model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolResult {
    /// The [`ToolCall::id`] this result corresponds to.
    pub tool_call_id: String,
    /// The tool's output (or error message if `is_error` is true).
    pub content: String,
    /// Whether the tool invocation failed.
    pub is_error: bool,
}

impl std::fmt::Display for ToolResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_error {
            write!(f, "err:{} ({})", self.tool_call_id, self.content)
        } else {
            write!(f, "ok:{}", self.tool_call_id)
        }
    }
}

/// A complete response from a model.
///
/// Returned by [`Provider::generate`](crate::Provider::generate). For
/// streaming responses, accumulate [`StreamEvent`](crate::StreamEvent)s
/// instead.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatResponse {
    /// The content blocks produced by the model.
    pub content: Vec<ContentBlock>,
    /// Token counts for this request/response pair.
    pub usage: Usage,
    /// Why the model stopped generating.
    pub stop_reason: StopReason,
    /// The model identifier that actually served the request (may differ
    /// from the requested model if the provider performed routing).
    pub model: String,
    /// Provider-specific metadata (e.g., request IDs, cache info).
    ///
    /// Contents are provider-specific. See each provider crate's docs for
    /// the keys it populates. Common keys include `"request_id"`.
    pub metadata: HashMap<String, Value>,
}

impl ChatResponse {
    /// Creates an empty response with no content.
    ///
    /// Useful for timeout or error cases where no LLM response was received.
    pub fn empty() -> Self {
        Self {
            content: Vec::new(),
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: String::new(),
            metadata: HashMap::new(),
        }
    }

    /// Returns the text of the first [`ContentBlock::Text`] block, if any.
    ///
    /// This is a convenience for the common case where you only care
    /// about the model's text output. For responses that may contain
    /// multiple text blocks, iterate `content` directly.
    pub fn text(&self) -> Option<&str> {
        self.content.iter().find_map(|b| match b {
            ContentBlock::Text(t) => Some(t.as_str()),
            _ => None,
        })
    }

    /// Returns references to all [`ToolCall`]s in the response.
    ///
    /// Returns an empty `Vec` when the response contains no tool calls.
    /// This is the primary accessor for implementing tool-use loops.
    ///
    /// For allocation-free iteration, use [`tool_calls_iter`](Self::tool_calls_iter).
    pub fn tool_calls(&self) -> Vec<&ToolCall> {
        self.content
            .iter()
            .filter_map(|b| match b {
                ContentBlock::ToolCall(tc) => Some(tc),
                _ => None,
            })
            .collect()
    }

    /// Returns an iterator over all [`ToolCall`]s in the response.
    ///
    /// Prefer this over [`tool_calls`](Self::tool_calls) to avoid allocation.
    pub fn tool_calls_iter(&self) -> impl Iterator<Item = &ToolCall> {
        self.content.iter().filter_map(|b| match b {
            ContentBlock::ToolCall(tc) => Some(tc),
            _ => None,
        })
    }

    /// Consumes the response and returns all [`ToolCall`]s.
    ///
    /// This is more efficient than [`tool_calls`](Self::tool_calls) when you
    /// need owned `ToolCall` values and won't use the response afterward.
    /// Non-tool-call content blocks are discarded.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use llm_stack::{ChatResponse, ToolCall, ContentBlock};
    /// # use llm_stack::chat::StopReason;
    /// # use llm_stack::usage::Usage;
    /// # use serde_json::json;
    /// let response = ChatResponse {
    ///     content: vec![
    ///         ContentBlock::Text("Let me help".into()),
    ///         ContentBlock::ToolCall(ToolCall {
    ///             id: "call_1".into(),
    ///             name: "search".into(),
    ///             arguments: json!({"query": "rust"}),
    ///         }),
    ///     ],
    ///     stop_reason: StopReason::ToolUse,
    ///     usage: Usage::default(),
    ///     model: "test".into(),
    ///     metadata: Default::default(),
    /// };
    ///
    /// let calls = response.into_tool_calls();
    /// assert_eq!(calls.len(), 1);
    /// assert_eq!(calls[0].name, "search");
    /// ```
    pub fn into_tool_calls(self) -> Vec<ToolCall> {
        self.content
            .into_iter()
            .filter_map(|b| match b {
                ContentBlock::ToolCall(tc) => Some(tc),
                _ => None,
            })
            .collect()
    }

    /// Consumes the response content and partitions it into tool calls and other blocks.
    ///
    /// Returns `(tool_calls, other_content)` where `other_content` contains all
    /// non-tool-call blocks (`Text`, `Image`, `ToolResult`, etc.) suitable for building
    /// an assistant message in a tool loop.
    ///
    /// This is more efficient than calling both [`tool_calls`](Self::tool_calls)
    /// and filtering content separately, as it processes the content in a single pass.
    pub fn partition_content(self) -> (Vec<ToolCall>, Vec<ContentBlock>) {
        let mut tool_calls = Vec::new();
        let mut other = Vec::new();

        for block in self.content {
            match block {
                ContentBlock::ToolCall(tc) => tool_calls.push(tc),
                // Filter out ToolResult as they shouldn't be in assistant messages
                ContentBlock::ToolResult(_) => {}
                other_block => other.push(other_block),
            }
        }

        (tool_calls, other)
    }
}

/// The reason the model stopped producing output.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[non_exhaustive]
pub enum StopReason {
    /// The model finished its response naturally.
    EndTurn,
    /// The model wants to invoke one or more tools.
    ToolUse,
    /// The response hit the `max_tokens` limit.
    MaxTokens,
    /// The model emitted a stop sequence.
    StopSequence,
}

impl fmt::Display for StopReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EndTurn => f.write_str("end_turn"),
            Self::ToolUse => f.write_str("tool_use"),
            Self::MaxTokens => f.write_str("max_tokens"),
            Self::StopSequence => f.write_str("stop_sequence"),
        }
    }
}

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

    #[test]
    fn test_chat_role_copy_hash() {
        use std::collections::HashMap;
        let mut map = HashMap::new();
        let role = ChatRole::User;
        let role_copy = role; // Copy
        map.insert(role, "user");
        map.insert(role_copy, "user_copy");
        assert_eq!(map.len(), 1);
    }

    #[test]
    fn test_chat_role_all_variants() {
        let variants = [
            ChatRole::System,
            ChatRole::User,
            ChatRole::Assistant,
            ChatRole::Tool,
        ];
        for v in &variants {
            let debug = format!("{v:?}");
            assert!(!debug.is_empty());
        }
    }

    #[test]
    fn test_chat_role_serde_roundtrip() {
        let role = ChatRole::Assistant;
        let json = serde_json::to_string(&role).unwrap();
        let back: ChatRole = serde_json::from_str(&json).unwrap();
        assert_eq!(role, back);
    }

    // --- ChatMessage constructors ---

    #[test]
    fn test_user_constructor() {
        let msg = ChatMessage::user("hello");
        assert_eq!(msg.role, ChatRole::User);
        assert_eq!(msg.content, vec![ContentBlock::Text("hello".into())]);
    }

    #[test]
    fn test_assistant_constructor() {
        let msg = ChatMessage::assistant("hi");
        assert_eq!(msg.role, ChatRole::Assistant);
        assert_eq!(msg.content, vec![ContentBlock::Text("hi".into())]);
    }

    #[test]
    fn test_system_constructor() {
        let msg = ChatMessage::system("be nice");
        assert_eq!(msg.role, ChatRole::System);
    }

    #[test]
    fn test_tool_result_constructor() {
        let msg = ChatMessage::tool_result("tc_1", "42");
        assert_eq!(msg.role, ChatRole::Tool);
        assert!(matches!(
            &msg.content[0],
            ContentBlock::ToolResult(tr)
                if tr.tool_call_id == "tc_1" && tr.content == "42" && !tr.is_error
        ));
    }

    #[test]
    fn test_tool_error_constructor() {
        let msg = ChatMessage::tool_error("tc_1", "something broke");
        assert!(matches!(
            &msg.content[0],
            ContentBlock::ToolResult(tr) if tr.is_error
        ));
    }

    // --- ChatMessage clone/eq/serde ---

    #[test]
    fn test_message_text_clone_eq() {
        let msg = ChatMessage::user("hello");
        assert_eq!(msg, msg.clone());
    }

    #[test]
    fn test_message_serde_roundtrip() {
        let msg = ChatMessage::user("hello");
        let json = serde_json::to_string(&msg).unwrap();
        let back: ChatMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(msg, back);
    }

    #[test]
    fn test_message_tool_use() {
        let msg = ChatMessage {
            role: ChatRole::Assistant,
            content: vec![
                ContentBlock::ToolCall(ToolCall {
                    id: "1".into(),
                    name: "calc".into(),
                    arguments: serde_json::json!({"a": 1}),
                }),
                ContentBlock::ToolCall(ToolCall {
                    id: "2".into(),
                    name: "search".into(),
                    arguments: serde_json::json!({"q": "rust"}),
                }),
            ],
        };
        assert_eq!(msg.content.len(), 2);
        assert_eq!(msg, msg.clone());
    }

    #[test]
    fn test_message_tool_result() {
        let msg = ChatMessage::tool_result("1", "42");
        assert!(matches!(
            &msg.content[0],
            ContentBlock::ToolResult(tr) if tr.content == "42" && !tr.is_error
        ));
    }

    #[test]
    fn test_message_mixed_content() {
        let msg = ChatMessage {
            role: ChatRole::User,
            content: vec![
                ContentBlock::Text("look at this".into()),
                ContentBlock::Image {
                    media_type: "image/png".into(),
                    data: ImageSource::Base64("abc123".into()),
                },
                ContentBlock::ToolCall(ToolCall {
                    id: "1".into(),
                    name: "analyze".into(),
                    arguments: serde_json::json!({}),
                }),
            ],
        };
        assert_eq!(msg.content.len(), 3);
    }

    // --- ContentBlock tests ---

    #[test]
    fn test_content_block_image_base64() {
        let block = ContentBlock::Image {
            media_type: "image/jpeg".into(),
            data: ImageSource::Base64("data...".into()),
        };
        assert_eq!(block, block.clone());
    }

    #[test]
    fn test_content_block_image_url() {
        let block = ContentBlock::Image {
            media_type: "image/png".into(),
            data: ImageSource::from_url("https://example.com/img.png").unwrap(),
        };
        assert_eq!(block, block.clone());
    }

    #[test]
    fn test_image_source_from_url_valid() {
        let src = ImageSource::from_url("https://example.com/img.png");
        assert!(src.is_ok());
        let url = url::Url::parse("https://example.com/img.png").unwrap();
        assert_eq!(src.unwrap(), ImageSource::Url(url));
    }

    #[test]
    fn test_image_source_from_url_normalizes() {
        // URL is parsed and normalized
        let src = ImageSource::from_url("HTTP://EXAMPLE.COM").unwrap();
        assert!(matches!(
            &src,
            ImageSource::Url(u) if u.as_str() == "http://example.com/"
        ));
    }

    #[test]
    fn test_image_source_from_url_invalid() {
        let err = ImageSource::from_url("not a url");
        assert!(err.is_err());
        let _parse_err: url::ParseError = err.unwrap_err();

        assert!(ImageSource::from_url("").is_err());
    }

    #[test]
    fn test_content_block_reasoning() {
        let block = ContentBlock::Reasoning {
            content: "thinking step by step".into(),
        };
        assert_eq!(block, block.clone());
    }

    #[test]
    fn test_tool_call_json_arguments() {
        let call = ToolCall {
            id: "tc_1".into(),
            name: "search".into(),
            arguments: serde_json::json!({
                "query": "rust async",
                "filters": {"lang": "en", "limit": 10}
            }),
        };
        assert_eq!(call, call.clone());
    }

    #[test]
    fn test_tool_result_error_flag() {
        let ok = ToolResult {
            tool_call_id: "1".into(),
            content: "result".into(),
            is_error: false,
        };
        let err = ToolResult {
            tool_call_id: "1".into(),
            content: "result".into(),
            is_error: true,
        };
        assert_ne!(ok, err);
    }

    // --- ChatResponse tests ---

    #[test]
    fn test_chat_response_metadata() {
        let mut metadata = HashMap::new();
        metadata.insert("cost".into(), serde_json::json!({"usd": 0.01}));
        let resp = ChatResponse {
            content: vec![ContentBlock::Text("hi".into())],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test-model".into(),
            metadata,
        };
        assert!(resp.metadata.contains_key("cost"));
    }

    #[test]
    fn test_chat_response_serde_roundtrip() {
        let resp = ChatResponse {
            content: vec![ContentBlock::Text("hi".into())],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test-model".into(),
            metadata: HashMap::new(),
        };
        let json = serde_json::to_string(&resp).unwrap();
        let back: ChatResponse = serde_json::from_str(&json).unwrap();
        assert_eq!(resp, back);
    }

    #[test]
    fn test_chat_response_empty_content() {
        let resp = ChatResponse {
            content: vec![],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        assert!(resp.content.is_empty());
    }

    // --- StopReason tests ---

    #[test]
    fn test_stop_reason_all_variants() {
        let variants = [
            StopReason::EndTurn,
            StopReason::ToolUse,
            StopReason::MaxTokens,
            StopReason::StopSequence,
        ];
        for v in &variants {
            assert_eq!(*v, *v);
        }
    }

    #[test]
    fn test_stop_reason_serde_roundtrip() {
        let sr = StopReason::MaxTokens;
        let json = serde_json::to_string(&sr).unwrap();
        let back: StopReason = serde_json::from_str(&json).unwrap();
        assert_eq!(sr, back);
    }

    #[test]
    fn test_stop_reason_eq_hash() {
        use std::collections::HashMap;
        let mut map = HashMap::new();
        map.insert(StopReason::EndTurn, "end");
        map.insert(StopReason::ToolUse, "tool");
        assert_eq!(map[&StopReason::EndTurn], "end");
        assert_eq!(map[&StopReason::ToolUse], "tool");
    }

    // --- Display impls ---

    #[test]
    fn test_chat_role_display() {
        assert_eq!(ChatRole::System.to_string(), "system");
        assert_eq!(ChatRole::User.to_string(), "user");
        assert_eq!(ChatRole::Assistant.to_string(), "assistant");
        assert_eq!(ChatRole::Tool.to_string(), "tool");
    }

    #[test]
    fn test_stop_reason_display() {
        assert_eq!(StopReason::EndTurn.to_string(), "end_turn");
        assert_eq!(StopReason::ToolUse.to_string(), "tool_use");
        assert_eq!(StopReason::MaxTokens.to_string(), "max_tokens");
        assert_eq!(StopReason::StopSequence.to_string(), "stop_sequence");
    }

    // --- ChatResponse::text() ---

    #[test]
    fn test_chat_response_text_returns_first() {
        let resp = ChatResponse {
            content: vec![
                ContentBlock::Reasoning {
                    content: "thinking...".into(),
                },
                ContentBlock::Text("first".into()),
                ContentBlock::Text("second".into()),
            ],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        assert_eq!(resp.text(), Some("first"));
    }

    #[test]
    fn test_chat_response_text_none_when_no_text_blocks() {
        let resp = ChatResponse {
            content: vec![ContentBlock::Reasoning {
                content: "thinking".into(),
            }],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        assert_eq!(resp.text(), None);
    }

    #[test]
    fn test_chat_response_text_none_when_empty() {
        let resp = ChatResponse {
            content: vec![],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        assert_eq!(resp.text(), None);
    }

    // --- ChatResponse::tool_calls() ---

    #[test]
    fn test_chat_response_tool_calls() {
        let resp = ChatResponse {
            content: vec![
                ContentBlock::Text("Let me search.".into()),
                ContentBlock::ToolCall(ToolCall {
                    id: "1".into(),
                    name: "search".into(),
                    arguments: serde_json::json!({"q": "rust"}),
                }),
                ContentBlock::ToolCall(ToolCall {
                    id: "2".into(),
                    name: "calc".into(),
                    arguments: serde_json::json!({"expr": "2+2"}),
                }),
            ],
            usage: Usage::default(),
            stop_reason: StopReason::ToolUse,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        let calls = resp.tool_calls();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "search");
        assert_eq!(calls[1].name, "calc");
    }

    #[test]
    fn test_chat_response_tool_calls_empty_when_text_only() {
        let resp = ChatResponse {
            content: vec![ContentBlock::Text("hello".into())],
            usage: Usage::default(),
            stop_reason: StopReason::EndTurn,
            model: "test".into(),
            metadata: HashMap::new(),
        };
        assert!(resp.tool_calls().is_empty());
    }

    // --- is_empty ---

    #[test]
    fn test_message_is_empty() {
        let empty = ChatMessage {
            role: ChatRole::User,
            content: vec![],
        };
        assert!(empty.is_empty());
        assert!(!ChatMessage::user("hi").is_empty());
    }

    // --- ContentBlock serde wire format ---

    #[test]
    fn test_content_block_serde_text() {
        let block = ContentBlock::Text("hello".into());
        let val = serde_json::to_value(&block).unwrap();
        assert_eq!(val, serde_json::json!({"text": "hello"}));
        let back: ContentBlock = serde_json::from_value(val).unwrap();
        assert_eq!(back, block);
    }

    #[test]
    fn test_content_block_serde_image() {
        let block = ContentBlock::Image {
            media_type: "image/png".into(),
            data: ImageSource::Base64("abc".into()),
        };
        let val = serde_json::to_value(&block).unwrap();
        assert_eq!(
            val,
            serde_json::json!({"image": {"media_type": "image/png", "data": {"Base64": "abc"}}})
        );
        let back: ContentBlock = serde_json::from_value(val).unwrap();
        assert_eq!(back, block);
    }

    #[test]
    fn test_content_block_serde_tool_call() {
        let block = ContentBlock::ToolCall(ToolCall {
            id: "tc_1".into(),
            name: "search".into(),
            arguments: serde_json::json!({"q": "rust"}),
        });
        let val = serde_json::to_value(&block).unwrap();
        assert_eq!(
            val,
            serde_json::json!({"tool_call": {"id": "tc_1", "name": "search", "arguments": {"q": "rust"}}})
        );
        let back: ContentBlock = serde_json::from_value(val).unwrap();
        assert_eq!(back, block);
    }

    #[test]
    fn test_content_block_serde_tool_result() {
        let block = ContentBlock::ToolResult(ToolResult {
            tool_call_id: "tc_1".into(),
            content: "42".into(),
            is_error: false,
        });
        let val = serde_json::to_value(&block).unwrap();
        assert_eq!(
            val,
            serde_json::json!({"tool_result": {"tool_call_id": "tc_1", "content": "42", "is_error": false}})
        );
        let back: ContentBlock = serde_json::from_value(val).unwrap();
        assert_eq!(back, block);
    }

    #[test]
    fn test_content_block_serde_reasoning() {
        let block = ContentBlock::Reasoning {
            content: "thinking".into(),
        };
        let val = serde_json::to_value(&block).unwrap();
        assert_eq!(
            val,
            serde_json::json!({"reasoning": {"content": "thinking"}})
        );
        let back: ContentBlock = serde_json::from_value(val).unwrap();
        assert_eq!(back, block);
    }

    // --- Semantic validity: constructors produce valid role/content combinations ---

    #[test]
    fn test_user_constructor_produces_text_only() {
        let msg = ChatMessage::user("hello");
        assert_eq!(msg.role, ChatRole::User);
        assert!(
            msg.content
                .iter()
                .all(|b| matches!(b, ContentBlock::Text(_)))
        );
    }

    #[test]
    fn test_assistant_constructor_produces_text_only() {
        let msg = ChatMessage::assistant("hi");
        assert_eq!(msg.role, ChatRole::Assistant);
        assert!(
            msg.content
                .iter()
                .all(|b| matches!(b, ContentBlock::Text(_)))
        );
    }

    #[test]
    fn test_system_constructor_produces_text_only() {
        let msg = ChatMessage::system("be nice");
        assert_eq!(msg.role, ChatRole::System);
        assert!(
            msg.content
                .iter()
                .all(|b| matches!(b, ContentBlock::Text(_)))
        );
    }

    #[test]
    fn test_tool_result_constructor_produces_tool_result_only() {
        let msg = ChatMessage::tool_result("tc_1", "42");
        assert_eq!(msg.role, ChatRole::Tool);
        assert!(
            msg.content
                .iter()
                .all(|b| matches!(b, ContentBlock::ToolResult(_)))
        );
    }

    #[test]
    fn test_tool_error_constructor_produces_tool_result_only() {
        let msg = ChatMessage::tool_error("tc_1", "boom");
        assert_eq!(msg.role, ChatRole::Tool);
        assert!(
            msg.content
                .iter()
                .all(|b| matches!(b, ContentBlock::ToolResult(r) if r.is_error))
        );
    }

    #[test]
    fn test_assistant_tool_calls_is_valid_combination() {
        // Assistants can produce tool calls — this is a valid state
        let msg = ChatMessage {
            role: ChatRole::Assistant,
            content: vec![
                ContentBlock::Text("Let me search for that.".into()),
                ContentBlock::ToolCall(ToolCall {
                    id: "1".into(),
                    name: "search".into(),
                    arguments: serde_json::json!({"q": "rust"}),
                }),
            ],
        };
        assert_eq!(msg.role, ChatRole::Assistant);
        assert_eq!(msg.content.len(), 2);
    }

    #[test]
    fn test_user_with_image_is_valid_combination() {
        // Users can send text + images — this is a valid state
        let msg = ChatMessage {
            role: ChatRole::User,
            content: vec![
                ContentBlock::Text("What's this?".into()),
                ContentBlock::Image {
                    media_type: "image/png".into(),
                    data: ImageSource::Base64("...".into()),
                },
            ],
        };
        assert_eq!(msg.role, ChatRole::User);
        assert_eq!(msg.content.len(), 2);
    }

    #[test]
    fn test_assistant_with_reasoning_is_valid_combination() {
        // Assistants can produce reasoning + text — valid state
        let msg = ChatMessage {
            role: ChatRole::Assistant,
            content: vec![
                ContentBlock::Reasoning {
                    content: "step 1: think about it".into(),
                },
                ContentBlock::Text("The answer is 42.".into()),
            ],
        };
        assert_eq!(msg.role, ChatRole::Assistant);
        assert_eq!(msg.content.len(), 2);
    }

    // --- ChatMessage serialization stability tests ---

    #[test]
    fn test_chat_message_to_json() {
        let msg = ChatMessage::user("Hello, world!");
        let json = msg.to_json().unwrap();
        assert_eq!(json["role"], "User");
        assert_eq!(json["content"][0]["text"], "Hello, world!");
    }

    #[test]
    fn test_chat_message_from_json() {
        let json = serde_json::json!({
            "role": "Assistant",
            "content": [{"text": "Hello!"}]
        });
        let msg = ChatMessage::from_json(&json).unwrap();
        assert_eq!(msg.role, ChatRole::Assistant);
        assert!(matches!(&msg.content[0], ContentBlock::Text(t) if t == "Hello!"));
    }

    #[test]
    fn test_chat_message_json_roundtrip() {
        let original = ChatMessage {
            role: ChatRole::User,
            content: vec![
                ContentBlock::Text("What's this?".into()),
                ContentBlock::Image {
                    media_type: "image/png".into(),
                    data: ImageSource::Base64("abc123".into()),
                },
            ],
        };
        let json = original.to_json().unwrap();
        let restored = ChatMessage::from_json(&json).unwrap();
        assert_eq!(original, restored);
    }

    #[test]
    fn test_chat_message_json_roundtrip_with_tool_result() {
        let original = ChatMessage::tool_result("tc_1", "success");
        let json = original.to_json().unwrap();
        let restored = ChatMessage::from_json(&json).unwrap();
        assert_eq!(original, restored);
    }

    // --- ToolResult metadata tests ---
}