gemini-cli-sdk 0.1.0

//! Wire session notification types and the `SessionUpdate` parsed enum.
//!
//! Gemini CLI pushes session progress as JSON-RPC notifications with the method
//! `session/update`. The `params` object contains `sessionId` and an `update`
//! sub-object with a `sessionUpdate` string discriminator followed by
//! variant-specific fields. `SessionUpdateNotification` captures the raw form;
//! calling `.parse()` converts it into the strongly-typed `SessionUpdate` enum.

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

// ── Raw notification ──────────────────────────────────────────────────────────

/// Raw `session/update` notification params as received from Gemini CLI.
///
/// The Gemini CLI nests the discriminator and variant fields inside an `update`
/// sub-object: `{ "sessionId": "...", "update": { "sessionUpdate": "...", ... } }`.
///
/// A custom deserializer also supports the legacy flat layout where
/// `sessionUpdate` appears directly in the params object alongside `sessionId`.
#[derive(Debug, Clone, Serialize)]
pub(crate) struct SessionUpdateNotification {
    /// Session this update belongs to.
    pub session_id: String,
    /// Discriminator string identifying the update variant.
    pub session_update: String,
    /// All remaining fields from the update body for variant dispatch.
    pub data: Value,
}

/// Nested `update` body used by the current Gemini CLI wire format.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct UpdateBody {
    session_update: String,
    #[serde(flatten)]
    data: Value,
}

/// Supports both the current nested format and the legacy flat format.
impl<'de> Deserialize<'de> for SessionUpdateNotification {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        // Deserialize as a generic map first.
        let mut map: serde_json::Map<String, Value> =
            serde_json::Map::deserialize(deserializer)?;

        let session_id = map
            .remove("sessionId")
            .and_then(|v| v.as_str().map(String::from))
            .unwrap_or_default();

        // Current format: nested `update` object containing `sessionUpdate` + fields.
        if let Some(update_val) = map.remove("update") {
            let body: UpdateBody = serde_json::from_value(update_val)
                .map_err(serde::de::Error::custom)?;
            return Ok(Self {
                session_id,
                session_update: body.session_update,
                data: body.data,
            });
        }

        // Legacy flat format: `sessionUpdate` directly in the params object.
        let session_update = map
            .remove("sessionUpdate")
            .and_then(|v| v.as_str().map(String::from))
            .ok_or_else(|| {
                serde::de::Error::missing_field("update or sessionUpdate")
            })?;

        Ok(Self {
            session_id,
            session_update,
            data: Value::Object(map),
        })
    }
}

// ── Parsed enum ───────────────────────────────────────────────────────────────

/// Strongly-typed representation of a `session/update` notification.
///
/// The variants mirror the `sessionUpdate` discriminator strings that Gemini CLI
/// emits. `Unknown` is the catch-all for any future variants we have not yet
/// modelled — callers should ignore it rather than error.
#[derive(Debug, Clone)]
pub(crate) enum SessionUpdate {
    /// A fragment of the agent's internal reasoning (may be streamed).
    AgentThoughtChunk {
        content: super::content::WireContentBlock,
    },
    /// A fragment of the agent's visible reply (may be streamed).
    AgentMessageChunk {
        content: super::content::WireContentBlock,
    },
    /// A new tool call has been initiated by the agent.
    ToolCall {
        tool_call_id: String,
        title: String,
        kind: String,
        status: String,
        locations: Vec<super::methods::ToolLocation>,
    },
    /// An in-progress tool call has changed status or produced output.
    ToolCallUpdate {
        tool_call_id: String,
        status: String,
        content: Option<Vec<super::content::WireContentBlock>>,
    },
    /// The agent has emitted or updated its plan.
    Plan {
        entries: Vec<WirePlanEntry>,
    },
    /// Token / cost usage information.
    UsageUpdate {
        cost: Option<f64>,
        size: Option<Value>,
        used: Option<Value>,
    },
    /// Session metadata (e.g. auto-generated title) has changed.
    SessionInfoUpdate {
        title: Option<String>,
    },
    /// A fragment of the echoed user message (may be streamed).
    UserMessageChunk {
        content: super::content::WireContentBlock,
    },
    /// A variant this SDK version does not recognise — forward-compat escape hatch.
    Unknown {
        kind: String,
        data: Value,
    },
}

// ── Plan entry ────────────────────────────────────────────────────────────────

/// A single step in a plan emitted by the agent.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct WirePlanEntry {
    /// Human-readable description of this plan step.
    #[serde(default)]
    pub content: String,
    /// Priority hint (`"high"`, `"normal"`, `"low"`, …).
    #[serde(default)]
    pub priority: String,
    /// Current execution status (`"pending"`, `"in_progress"`, `"done"`, …).
    #[serde(default)]
    pub status: String,
    /// Unknown/future keys are preserved.
    #[serde(flatten)]
    pub extra: Value,
}

// ── Parsing ───────────────────────────────────────────────────────────────────

impl SessionUpdateNotification {
    /// Parse this raw notification into a strongly-typed [`SessionUpdate`].
    ///
    /// Parsing is best-effort: any field that cannot be deserialized into its
    /// expected type falls back to a zero/default value rather than returning
    /// an error, keeping streaming sessions resilient against version skew.
    pub fn parse(self) -> SessionUpdate {
        match self.session_update.as_str() {
            "agent_thought_chunk" => {
                let content = extract_content(&self.data, "content");
                SessionUpdate::AgentThoughtChunk { content }
            }
            "agent_message_chunk" => {
                let content = extract_content(&self.data, "content");
                SessionUpdate::AgentMessageChunk { content }
            }
            "tool_call" => SessionUpdate::ToolCall {
                tool_call_id: str_field(&self.data, "toolCallId"),
                title: str_field(&self.data, "title"),
                kind: str_field(&self.data, "kind"),
                status: str_field(&self.data, "status"),
                locations: serde_json::from_value(
                    self.data.get("locations").cloned().unwrap_or_default(),
                )
                .unwrap_or_default(),
            },
            "tool_call_update" => SessionUpdate::ToolCallUpdate {
                tool_call_id: str_field(&self.data, "toolCallId"),
                status: str_field(&self.data, "status"),
                content: self
                    .data
                    .get("content")
                    .and_then(|v| serde_json::from_value(v.clone()).ok()),
            },
            "plan" => {
                let entries = serde_json::from_value(
                    self.data.get("entries").cloned().unwrap_or_default(),
                )
                .unwrap_or_default();
                SessionUpdate::Plan { entries }
            }
            "usage_update" => SessionUpdate::UsageUpdate {
                cost: self.data.get("cost").and_then(|v| v.as_f64()),
                size: self.data.get("size").cloned(),
                used: self.data.get("used").cloned(),
            },
            "session_info_update" => SessionUpdate::SessionInfoUpdate {
                title: self
                    .data
                    .get("title")
                    .and_then(|v| v.as_str())
                    .map(str::to_string),
            },
            "user_message_chunk" => {
                let content = extract_content(&self.data, "content");
                SessionUpdate::UserMessageChunk { content }
            }
            other => SessionUpdate::Unknown {
                kind: other.to_string(),
                data: self.data,
            },
        }
    }
}

// ── Private helpers ───────────────────────────────────────────────────────────

/// Extract a nested `WireContentBlock` from `data[key]`, defaulting on failure.
fn extract_content(data: &Value, key: &str) -> super::content::WireContentBlock {
    data.get(key)
        .and_then(|v| serde_json::from_value(v.clone()).ok())
        .unwrap_or_default()
}

/// Extract a string field from a `Value` map, defaulting to `""` on absence/failure.
fn str_field(data: &Value, key: &str) -> String {
    data.get(key)
        .and_then(|v| v.as_str())
        .unwrap_or_default()
        .to_string()
}

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

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

    /// Build a notification using the **current nested** format:
    /// `{ "sessionId": "...", "update": { "sessionUpdate": "...", ...extra } }`
    fn make_notification(session_update: &str, extra: Value) -> SessionUpdateNotification {
        let mut update_obj = serde_json::Map::new();
        update_obj.insert("sessionUpdate".to_string(), json!(session_update));
        if let Value::Object(map) = extra {
            update_obj.extend(map);
        }
        let raw = json!({
            "sessionId": "sess-test",
            "update": Value::Object(update_obj),
        });
        serde_json::from_value(raw).expect("construct test notification")
    }

    /// Build a notification using the **legacy flat** format:
    /// `{ "sessionId": "...", "sessionUpdate": "...", ...extra }`
    fn make_notification_flat(session_update: &str, extra: Value) -> SessionUpdateNotification {
        let mut obj = serde_json::Map::new();
        obj.insert("sessionId".to_string(), json!("sess-test"));
        obj.insert("sessionUpdate".to_string(), json!(session_update));
        if let Value::Object(map) = extra {
            obj.extend(map);
        }
        serde_json::from_value(Value::Object(obj)).expect("construct test notification (flat)")
    }

    #[test]
    fn test_parse_thought_chunk() {
        let notif = make_notification(
            "agent_thought_chunk",
            json!({ "content": { "type": "text", "text": "thinking..." } }),
        );
        assert_eq!(notif.session_id, "sess-test");

        match notif.parse() {
            SessionUpdate::AgentThoughtChunk { content } => {
                assert_eq!(content.content_type, "text");
                assert_eq!(content.text.as_deref(), Some("thinking..."));
            }
            other => panic!("expected AgentThoughtChunk, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_tool_call() {
        let notif = make_notification(
            "tool_call",
            json!({
                "toolCallId": "tc-42",
                "title":      "Read file",
                "kind":       "read_file",
                "status":     "in_progress",
                "locations":  [{ "uri": "file:///src/main.rs" }]
            }),
        );

        match notif.parse() {
            SessionUpdate::ToolCall {
                tool_call_id,
                title,
                kind,
                status,
                locations,
            } => {
                assert_eq!(tool_call_id, "tc-42");
                assert_eq!(title, "Read file");
                assert_eq!(kind, "read_file");
                assert_eq!(status, "in_progress");
                assert_eq!(locations.len(), 1);
                assert_eq!(locations[0].uri, "file:///src/main.rs");
            }
            other => panic!("expected ToolCall, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_tool_call_update() {
        let notif = make_notification(
            "tool_call_update",
            json!({
                "toolCallId": "tc-42",
                "status": "done",
                "content": [{ "type": "text", "text": "result" }]
            }),
        );

        match notif.parse() {
            SessionUpdate::ToolCallUpdate {
                tool_call_id,
                status,
                content,
            } => {
                assert_eq!(tool_call_id, "tc-42");
                assert_eq!(status, "done");
                let blocks = content.expect("content should be present");
                assert_eq!(blocks.len(), 1);
                assert_eq!(blocks[0].text.as_deref(), Some("result"));
            }
            other => panic!("expected ToolCallUpdate, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_plan() {
        let notif = make_notification(
            "plan",
            json!({
                "entries": [
                    { "content": "Step 1", "priority": "high",   "status": "done" },
                    { "content": "Step 2", "priority": "normal",  "status": "pending" }
                ]
            }),
        );

        match notif.parse() {
            SessionUpdate::Plan { entries } => {
                assert_eq!(entries.len(), 2);
                assert_eq!(entries[0].content, "Step 1");
                assert_eq!(entries[0].status, "done");
                assert_eq!(entries[1].priority, "normal");
            }
            other => panic!("expected Plan, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_usage_update() {
        let notif = make_notification(
            "usage_update",
            json!({ "cost": 0.0042, "size": { "tokens": 512 } }),
        );

        match notif.parse() {
            SessionUpdate::UsageUpdate { cost, size, used } => {
                let c = cost.expect("cost should be present");
                assert!((c - 0.0042_f64).abs() < f64::EPSILON);
                assert!(size.is_some());
                assert!(used.is_none());
            }
            other => panic!("expected UsageUpdate, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_session_info_update() {
        let notif = make_notification(
            "session_info_update",
            json!({ "title": "My conversation" }),
        );

        match notif.parse() {
            SessionUpdate::SessionInfoUpdate { title } => {
                assert_eq!(title.as_deref(), Some("My conversation"));
            }
            other => panic!("expected SessionInfoUpdate, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_unknown() {
        let notif = make_notification(
            "some_future_variant",
            json!({ "arbitraryField": 99 }),
        );

        match notif.parse() {
            SessionUpdate::Unknown { kind, data } => {
                assert_eq!(kind, "some_future_variant");
                assert_eq!(data["arbitraryField"], 99);
            }
            other => panic!("expected Unknown, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_agent_message_chunk() {
        let notif = make_notification(
            "agent_message_chunk",
            json!({ "content": { "type": "text", "text": "Hello!" } }),
        );

        match notif.parse() {
            SessionUpdate::AgentMessageChunk { content } => {
                assert_eq!(content.text.as_deref(), Some("Hello!"));
            }
            other => panic!("expected AgentMessageChunk, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_gracefully_handles_missing_content_field() {
        // If `content` is absent the parser must return a default block, not panic.
        let notif = make_notification("agent_thought_chunk", json!({}));
        match notif.parse() {
            SessionUpdate::AgentThoughtChunk { content } => {
                assert_eq!(content.content_type, "");
            }
            other => panic!("expected AgentThoughtChunk, got {other:?}"),
        }
    }

    #[test]
    fn test_legacy_flat_format_still_works() {
        // Older Gemini CLI versions sent sessionUpdate flat alongside sessionId.
        let notif = make_notification_flat(
            "agent_message_chunk",
            json!({ "content": { "type": "text", "text": "flat format" } }),
        );
        assert_eq!(notif.session_id, "sess-test");
        match notif.parse() {
            SessionUpdate::AgentMessageChunk { content } => {
                assert_eq!(content.text.as_deref(), Some("flat format"));
            }
            other => panic!("expected AgentMessageChunk, got {other:?}"),
        }
    }

    #[test]
    fn test_nested_format_real_cli_shape() {
        // Exact shape observed from Gemini CLI v2025+.
        let raw = json!({
            "sessionId": "7ea33c5a-xxxx",
            "update": {
                "sessionUpdate": "agent_message_chunk",
                "content": { "type": "text", "text": "PONG" }
            }
        });
        let notif: SessionUpdateNotification =
            serde_json::from_value(raw).expect("deserialize nested format");
        assert_eq!(notif.session_id, "7ea33c5a-xxxx");
        assert_eq!(notif.session_update, "agent_message_chunk");
        match notif.parse() {
            SessionUpdate::AgentMessageChunk { content } => {
                assert_eq!(content.text.as_deref(), Some("PONG"));
            }
            other => panic!("expected AgentMessageChunk, got {other:?}"),
        }
    }
}