gemini-cli-sdk 0.1.0

//! Public message types — the primary types consumers work with.
//!
//! The [`Message`] enum is the top-level type yielded by the streaming API.
//! Each variant carries a typed payload. All structs use `#[serde(default)]`
//! on every field for forward-compatibility with unknown server additions,
//! and `#[serde(flatten)] pub extra: Value` to capture unknown fields without
//! losing them.

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

use super::content::ContentBlock;

// ── Top-level Message enum ──────────────────────────────────────────────

/// Top-level message enum — the primary type consumers work with.
///
/// Every message emitted by the SDK during a session is one of these variants.
/// Consumers typically `match` on this enum to route messages to their
/// appropriate handlers.
///
/// # Example
///
/// ```rust,ignore
/// match message {
///     Message::System(sys) => println!("session: {}", sys.session_id),
///     Message::Assistant(asst) => {
///         if let Some(text) = Message::Assistant(asst).assistant_text() {
///             println!("{}", text);
///         }
///     }
///     Message::Result(r) if r.is_error => eprintln!("error turn"),
///     Message::StreamEvent(ev) => { /* rich streaming data */ }
///     _ => {}
/// }
/// ```
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Message {
    /// Synthesized initialization message, emitted once per session.
    System(SystemMessage),
    /// An assistant response containing one or more content blocks.
    Assistant(AssistantMessage),
    /// A user message (present in session history / resume replay).
    User(UserMessage),
    /// Terminal message for a completed prompt turn.
    Result(ResultMessage),
    /// Rich streaming event for data that does not map to a content block
    /// (tool call lifecycle, plan updates, usage deltas, etc.).
    #[serde(rename = "stream_event")]
    StreamEvent(StreamEvent),
}

impl Message {
    /// Returns the session ID carried by this message, if present.
    ///
    /// All current variants carry a `session_id`. This method returns
    /// `Some(&str)` for every variant; the `Option` wrapper exists for
    /// forward-compatibility with future variants that may omit it.
    pub fn session_id(&self) -> Option<&str> {
        match self {
            Message::System(m) => Some(&m.session_id),
            Message::Assistant(m) => Some(&m.session_id),
            Message::User(m) => Some(&m.session_id),
            Message::Result(m) => Some(&m.session_id),
            Message::StreamEvent(m) => Some(&m.session_id),
        }
    }

    /// Returns `true` if this is a [`ResultMessage`] with `is_error = true`.
    pub fn is_error_result(&self) -> bool {
        matches!(self, Message::Result(r) if r.is_error)
    }

    /// Returns `true` if this is a [`StreamEvent`].
    pub fn is_stream_event(&self) -> bool {
        matches!(self, Message::StreamEvent(_))
    }

    /// Extracts concatenated text from all [`ContentBlock::Text`] blocks in
    /// an [`AssistantMessage`]. Returns `None` for any other variant, and
    /// `None` when the assistant message contains no text blocks.
    pub fn assistant_text(&self) -> Option<String> {
        let Message::Assistant(m) = self else {
            return None;
        };
        let texts: Vec<&str> = m
            .message
            .content
            .iter()
            .filter_map(|b| match b {
                ContentBlock::Text(t) => Some(t.text.as_str()),
                _ => None,
            })
            .collect();
        if texts.is_empty() {
            None
        } else {
            Some(texts.join(""))
        }
    }
}

// ── System Message ──────────────────────────────────────────────────────

/// Initialization message synthesized from `initialize` + `session/new`
/// responses. Emitted exactly once at the start of each session.
///
/// This is not a wire type — it is constructed by the SDK transport layer
/// during the connection handshake and never arrives as a raw JSON-RPC
/// notification.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SystemMessage {
    /// Subtype discriminator (e.g., `"init"`).
    #[serde(default)]
    pub subtype: String,
    /// Unique session identifier assigned by the server.
    #[serde(default)]
    pub session_id: String,
    /// Server's current working directory.
    #[serde(default)]
    pub cwd: String,
    /// Names of built-in tools available in this session.
    #[serde(default)]
    pub tools: Vec<String>,
    /// Status of configured MCP servers.
    #[serde(default)]
    pub mcp_servers: Vec<McpServerStatus>,
    /// Model identifier active for this session.
    #[serde(default)]
    pub model: String,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

/// MCP server connection status reported in [`SystemMessage`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct McpServerStatus {
    /// Server name as declared in the MCP configuration.
    pub name: String,
    /// Connection status (e.g., `"connected"`, `"failed"`).
    #[serde(default)]
    pub status: String,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

// ── Assistant Message ───────────────────────────────────────────────────

/// An assistant response containing one or more typed content blocks.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AssistantMessage {
    /// The inner message payload with role, content, and model metadata.
    pub message: AssistantMessageInner,
    /// Session this message belongs to.
    #[serde(default)]
    pub session_id: String,
}

/// Inner content of an [`AssistantMessage`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AssistantMessageInner {
    /// Always `"assistant"`.
    #[serde(default)]
    pub role: String,
    /// Typed content blocks (text, tool use, thinking, image, etc.).
    #[serde(default)]
    pub content: Vec<ContentBlock>,
    /// Model that generated this response.
    #[serde(default)]
    pub model: String,
    /// Reason the model stopped generating (e.g., `"end_turn"`, `"tool_use"`).
    #[serde(default)]
    pub stop_reason: String,
    /// Stop sequence that triggered termination, if any.
    #[serde(default)]
    pub stop_sequence: Option<String>,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

// ── User Message ────────────────────────────────────────────────────────

/// A user message — present in session history and resume replay only.
///
/// Outbound user messages are sent via the transport layer, not constructed
/// as `Message::User` variants.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct UserMessage {
    /// The inner message payload.
    pub message: UserMessageInner,
    /// Session this message belongs to.
    #[serde(default)]
    pub session_id: String,
}

/// Inner content of a [`UserMessage`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct UserMessageInner {
    /// Always `"user"`.
    #[serde(default)]
    pub role: String,
    /// Content blocks in this user message.
    #[serde(default)]
    pub content: Vec<ContentBlock>,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

// ── Result Message ──────────────────────────────────────────────────────

/// Terminal message for a completed prompt turn.
///
/// Emitted exactly once per `send_message` call, after all streaming content
/// has been delivered. Check [`ResultMessage::is_error`] to determine success.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResultMessage {
    /// Subtype discriminator (e.g., `"success"`, `"error"`).
    #[serde(default)]
    pub subtype: String,
    /// `true` if the turn ended in an error condition.
    #[serde(default)]
    pub is_error: bool,
    /// Wall-clock duration of the full turn in milliseconds.
    #[serde(default)]
    pub duration_ms: f64,
    /// Time spent in API calls in milliseconds.
    #[serde(default)]
    pub duration_api_ms: f64,
    /// Number of agentic turns taken.
    #[serde(default)]
    pub num_turns: u32,
    /// Session this result belongs to.
    #[serde(default)]
    pub session_id: String,
    /// Token usage statistics for this turn.
    #[serde(default)]
    pub usage: Usage,
    /// Final stop reason from the model.
    #[serde(default)]
    pub stop_reason: String,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

// ── Usage ───────────────────────────────────────────────────────────────

/// Token usage statistics for a completed turn.
///
/// All fields default to `0`. Gemini-specific fields (e.g., `thought_tokens`)
/// are included alongside the standard Anthropic-compatible fields.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Usage {
    /// Tokens consumed from the input (prompt) context.
    #[serde(default)]
    pub input_tokens: u32,
    /// Tokens generated in the response.
    #[serde(default)]
    pub output_tokens: u32,
    /// Prompt cache read tokens (prompt cache hit).
    #[serde(default)]
    pub cache_read_input_tokens: u32,
    /// Prompt cache write tokens (prompt cache miss, storing new cache).
    #[serde(default)]
    pub cache_creation_input_tokens: u32,
    /// Gemini-specific: tokens consumed by thinking/reasoning steps.
    #[serde(default)]
    pub thought_tokens: u32,
}

// ── Stream Event ────────────────────────────────────────────────────────

/// A rich streaming event for data that does not map 1:1 to content blocks.
///
/// Used for tool call lifecycle events, plan updates, usage deltas,
/// permission requests, and any other structured event the wire format
/// emits between content blocks. The `data` field carries the raw JSON
/// payload — consumers should match on `event_type` to deserialize it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StreamEvent {
    /// Discriminator identifying what `data` contains.
    pub event_type: String,
    /// Raw event payload — structure depends on `event_type`.
    #[serde(default)]
    pub data: Value,
    /// Session this event belongs to.
    #[serde(default)]
    pub session_id: String,
}

// ── Session Info ────────────────────────────────────────────────────────

/// Metadata about the established session, returned from `Client::connect()`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SessionInfo {
    /// Unique session identifier.
    pub session_id: String,
    /// Model active for this session.
    #[serde(default)]
    pub model: String,
    /// Names of tools available in this session.
    #[serde(default)]
    pub tools: Vec<String>,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

// ── Plan Entry (Gemini-specific) ────────────────────────────────────────

/// A single entry in a Gemini agentic plan.
///
/// Gemini CLI emits structured plans during multi-step reasoning. Each entry
/// represents one planned action. Plans arrive via [`StreamEvent`] with
/// `event_type = "plan_update"`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PlanEntry {
    /// Human-readable description of this plan step.
    #[serde(default)]
    pub content: String,
    /// Execution priority (e.g., `"high"`, `"normal"`, `"low"`).
    #[serde(default)]
    pub priority: String,
    /// Current execution status (e.g., `"pending"`, `"done"`, `"failed"`).
    #[serde(default)]
    pub status: String,
    /// Unknown fields preserved for forward-compatibility.
    #[serde(flatten)]
    pub extra: Value,
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::content::TextBlock;
    use serde_json::json;

    // Helper: build a minimal SystemMessage with a known session_id.
    fn system_msg(session_id: &str) -> Message {
        Message::System(SystemMessage {
            subtype: "init".to_owned(),
            session_id: session_id.to_owned(),
            cwd: "/tmp".to_owned(),
            tools: vec![],
            mcp_servers: vec![],
            model: "gemini-2.5-pro".to_owned(),
            extra: Value::Object(Default::default()),
        })
    }

    // Helper: build a ResultMessage.
    fn result_msg(session_id: &str, is_error: bool) -> Message {
        Message::Result(ResultMessage {
            subtype: if is_error { "error" } else { "success" }.to_owned(),
            is_error,
            duration_ms: 123.4,
            duration_api_ms: 100.0,
            num_turns: 1,
            session_id: session_id.to_owned(),
            usage: Usage::default(),
            stop_reason: "end_turn".to_owned(),
            extra: Value::Object(Default::default()),
        })
    }

    // Helper: build an AssistantMessage with given content blocks.
    fn assistant_msg(session_id: &str, content: Vec<ContentBlock>) -> Message {
        Message::Assistant(AssistantMessage {
            message: AssistantMessageInner {
                role: "assistant".to_owned(),
                content,
                model: "gemini-2.5-pro".to_owned(),
                stop_reason: "end_turn".to_owned(),
                stop_sequence: None,
                extra: Value::Object(Default::default()),
            },
            session_id: session_id.to_owned(),
        })
    }

    // Helper: build a StreamEvent message.
    fn stream_event_msg(session_id: &str) -> Message {
        Message::StreamEvent(StreamEvent {
            event_type: "tool_call_start".to_owned(),
            data: json!({ "tool": "bash" }),
            session_id: session_id.to_owned(),
        })
    }

    // ── session_id ─────────────────────────────────────────────────────

    #[test]
    fn test_message_system_session_id() {
        let msg = system_msg("sess-abc");
        assert_eq!(msg.session_id(), Some("sess-abc"));
    }

    #[test]
    fn test_message_result_session_id() {
        let msg = result_msg("sess-xyz", false);
        assert_eq!(msg.session_id(), Some("sess-xyz"));
    }

    #[test]
    fn test_message_stream_event_session_id() {
        let msg = stream_event_msg("sess-ev");
        assert_eq!(msg.session_id(), Some("sess-ev"));
    }

    // ── is_error_result ────────────────────────────────────────────────

    #[test]
    fn test_message_is_error_result_true() {
        let msg = result_msg("s1", true);
        assert!(msg.is_error_result(), "is_error=true must return true");
    }

    #[test]
    fn test_message_is_error_result_false() {
        let msg = result_msg("s1", false);
        assert!(!msg.is_error_result(), "is_error=false must return false");
    }

    #[test]
    fn test_message_is_error_result_non_result_variant() {
        let msg = system_msg("s1");
        assert!(!msg.is_error_result(), "non-Result variant must return false");
    }

    // ── is_stream_event ────────────────────────────────────────────────

    #[test]
    fn test_message_is_stream_event() {
        let msg = stream_event_msg("s1");
        assert!(msg.is_stream_event());
    }

    #[test]
    fn test_message_is_stream_event_false_for_system() {
        let msg = system_msg("s1");
        assert!(!msg.is_stream_event());
    }

    // ── assistant_text ─────────────────────────────────────────────────

    #[test]
    fn test_message_assistant_text_single() {
        let content = vec![ContentBlock::Text(TextBlock::new("hello world"))];
        let msg = assistant_msg("s1", content);
        assert_eq!(msg.assistant_text(), Some("hello world".to_owned()));
    }

    #[test]
    fn test_message_assistant_text_multiple_blocks_concatenated() {
        let content = vec![
            ContentBlock::Text(TextBlock::new("foo")),
            ContentBlock::Text(TextBlock::new("bar")),
        ];
        let msg = assistant_msg("s1", content);
        assert_eq!(msg.assistant_text(), Some("foobar".to_owned()));
    }

    #[test]
    fn test_message_assistant_text_empty() {
        let msg = assistant_msg("s1", vec![]);
        assert_eq!(
            msg.assistant_text(),
            None,
            "no content blocks must yield None"
        );
    }

    #[test]
    fn test_message_assistant_text_non_text_blocks_only() {
        // A thinking block alone — no text → None.
        use crate::types::content::ThinkingBlock;
        let content = vec![ContentBlock::Thinking(ThinkingBlock::new("reasoning..."))];
        let msg = assistant_msg("s1", content);
        assert_eq!(
            msg.assistant_text(),
            None,
            "no Text blocks must yield None"
        );
    }

    #[test]
    fn test_message_assistant_text_non_assistant_variant() {
        let msg = system_msg("s1");
        assert_eq!(msg.assistant_text(), None);
    }

    // ── Usage::default ─────────────────────────────────────────────────

    #[test]
    fn test_usage_default() {
        let usage = Usage::default();
        assert_eq!(usage.input_tokens, 0);
        assert_eq!(usage.output_tokens, 0);
        assert_eq!(usage.cache_read_input_tokens, 0);
        assert_eq!(usage.cache_creation_input_tokens, 0);
        assert_eq!(usage.thought_tokens, 0);
    }

    // ── Serde roundtrip ────────────────────────────────────────────────

    #[test]
    fn test_message_serde_roundtrip_system() {
        let original = Message::System(SystemMessage {
            subtype: "init".to_owned(),
            session_id: "sess-roundtrip".to_owned(),
            cwd: "/workspace".to_owned(),
            tools: vec!["bash".to_owned(), "read_file".to_owned()],
            mcp_servers: vec![McpServerStatus {
                name: "filesystem".to_owned(),
                status: "connected".to_owned(),
                extra: Value::Object(Default::default()),
            }],
            model: "gemini-2.5-pro".to_owned(),
            extra: Value::Object(Default::default()),
        });

        let json = serde_json::to_string(&original).expect("serialize");
        let recovered: Message = serde_json::from_str(&json).expect("deserialize");

        assert_eq!(original, recovered);
    }

    #[test]
    fn test_message_serde_roundtrip_result() {
        let original = Message::Result(ResultMessage {
            subtype: "success".to_owned(),
            is_error: false,
            duration_ms: 450.75,
            duration_api_ms: 400.0,
            num_turns: 3,
            session_id: "sess-rt2".to_owned(),
            usage: Usage {
                input_tokens: 512,
                output_tokens: 128,
                cache_read_input_tokens: 64,
                cache_creation_input_tokens: 32,
                thought_tokens: 256,
            },
            stop_reason: "end_turn".to_owned(),
            extra: Value::Object(Default::default()),
        });

        let json = serde_json::to_string(&original).expect("serialize");
        let recovered: Message = serde_json::from_str(&json).expect("deserialize");

        assert_eq!(original, recovered);
    }

    #[test]
    fn test_message_serde_roundtrip_stream_event() {
        let original = Message::StreamEvent(StreamEvent {
            event_type: "plan_update".to_owned(),
            data: json!({ "step": 1, "action": "read_file" }),
            session_id: "sess-rt3".to_owned(),
        });

        let json = serde_json::to_string(&original).expect("serialize");
        let recovered: Message = serde_json::from_str(&json).expect("deserialize");

        assert_eq!(original, recovered);
    }

    // ── Plan Entry ─────────────────────────────────────────────────────

    #[test]
    fn test_plan_entry_defaults() {
        let entry: PlanEntry =
            serde_json::from_str("{}").expect("empty object must deserialize via defaults");
        assert!(entry.content.is_empty());
        assert!(entry.priority.is_empty());
        assert!(entry.status.is_empty());
    }

    #[test]
    fn test_plan_entry_roundtrip() {
        let original = PlanEntry {
            content: "Analyze the repository structure".to_owned(),
            priority: "high".to_owned(),
            status: "pending".to_owned(),
            extra: Value::Object(Default::default()),
        };

        let json = serde_json::to_string(&original).expect("serialize");
        let recovered: PlanEntry = serde_json::from_str(&json).expect("deserialize");

        assert_eq!(original, recovered);
    }

    // ── Usage thought_tokens (Gemini-specific) ─────────────────────────

    #[test]
    fn test_usage_thought_tokens_roundtrip() {
        let usage = Usage {
            input_tokens: 100,
            output_tokens: 50,
            cache_read_input_tokens: 0,
            cache_creation_input_tokens: 0,
            thought_tokens: 300,
        };
        let json = serde_json::to_string(&usage).expect("serialize");
        let recovered: Usage = serde_json::from_str(&json).expect("deserialize");
        assert_eq!(recovered.thought_tokens, 300);
    }
}