supercode_interchange/session/

gemini.rs

//! Gemini session codec: loaders, writers and native-record helpers.

use super::*;

impl Session {
    /// Load a Gemini CLI transcript from disk.
    pub fn from_gemini(path: impl AsRef<Path>) -> Result<Session> {
        Self::from_gemini_str(&std::fs::read_to_string(path.as_ref())?)
    }

    /// Parse Gemini CLI's line-oriented session format.
    ///
    /// Gemini stores a header without a `type`, followed by `user` and
    /// `gemini` records. Function calls are embedded in assistant content
    /// parts and function responses in user content parts. Unknown records
    /// remain byte-exact in [`Session::raw`] instead of silently entering the
    /// replay conversation.
    pub fn from_gemini_str(jsonl: &str) -> Result<Session> {
        let mut meta = SessionMeta::new(SessionSource::Gemini);
        let (raw_lines, raw_trailing_newline) = split_lines_verbatim(jsonl);
        let raw = raw_lines.iter().map(|line| line.to_string()).collect();
        let mut messages = Vec::new();
        let mut parse_error_lines = 0usize;
        let mut pending_by_name: HashMap<String, Vec<String>> = HashMap::new();

        // The parse walk skips blank lines; `raw` keeps them. Position `i` of the walk is raw line
        // `raw_record[i]`.
        let raw_record: Vec<usize> = raw_lines
            .iter()
            .enumerate()
            .filter(|(_, line)| !line.trim().is_empty())
            .map(|(index, _)| index)
            .collect();
        for (line_index, line) in non_empty_lines(jsonl).enumerate() {
            // What the previous record created belongs to it; this record starts after.
            if line_index > 0 {
                stamp_message_records(&mut messages, raw_record[line_index - 1]);
            }
            let value: Value = match serde_json::from_str(line) {
                Ok(value) => value,
                Err(_) => {
                    parse_error_lines += 1;
                    continue;
                }
            };
            let kind = value.get("type").and_then(Value::as_str);
            if kind.is_none() {
                if meta.session_id.is_none() {
                    meta.session_id = value
                        .get("sessionId")
                        .and_then(Value::as_str)
                        .map(str::to_string);
                }
                for (source, target) in [
                    ("projectHash", "gemini_project_hash"),
                    ("startTime", "created_at"),
                    ("lastUpdated", "updated_at"),
                    ("kind", "gemini_session_kind"),
                ] {
                    if let Some(raw) = value.get(source) {
                        meta.lineage.insert(
                            target.to_string(),
                            raw.as_str()
                                .map(str::to_string)
                                .unwrap_or_else(|| raw.to_string()),
                        );
                    }
                }
                continue;
            }
            if kind != Some("user") && kind != Some("gemini") {
                continue;
            }

            let timestamp = value.get("timestamp").and_then(Value::as_str);
            let model = value.get("model").and_then(Value::as_str);
            if let Some(model) = model {
                meta.model = Some(model.to_string());
            }
            let content = value.get("content").unwrap_or(&Value::Null);
            let parts = content.as_array();
            let text = match content {
                Value::String(text) => text.clone(),
                Value::Array(parts) => parts
                    .iter()
                    .filter_map(|part| part.get("text").and_then(Value::as_str))
                    .collect::<Vec<_>>()
                    .join(" ")
                    .trim()
                    .to_string(),
                _ => String::new(),
            };

            if kind == Some("gemini") {
                let legacy_calls = parts
                    .into_iter()
                    .flatten()
                    .filter_map(|part| part.get("functionCall"));
                let native_calls = value
                    .get("toolCalls")
                    .and_then(Value::as_array)
                    .into_iter()
                    .flatten();
                let calls = native_calls
                    .chain(legacy_calls)
                    .enumerate()
                    .filter_map(|(call_index, call)| {
                        let name = call.get("name")?.as_str()?.to_string();
                        let id = call
                            .get("id")
                            .and_then(Value::as_str)
                            .map(str::to_string)
                            .unwrap_or_else(|| format!("gemini-{line_index}-{call_index}"));
                        pending_by_name
                            .entry(name.clone())
                            .or_default()
                            .push(id.clone());
                        let arguments = call
                            .get("args")
                            .map(value_to_arg_string)
                            .unwrap_or_else(|| "{}".to_string());
                        Some(function_call(&id, &name, arguments))
                    })
                    .collect::<Vec<_>>();
                let mut message = ChatMessage {
                    role: Role::Assistant,
                    content: (!text.is_empty()).then_some(text),
                    content_parts: None,
                    tool_calls: (!calls.is_empty()).then_some(calls),
                    tool_call_id: None,
                    name: None,
                    metadata: Default::default(),
                };
                if let Some(timestamp) = timestamp {
                    message
                        .metadata
                        .insert("timestamp".into(), timestamp.into());
                }
                if let Some(model) = model {
                    message.metadata.insert("gemini_model".into(), model.into());
                }
                if let Some(thoughts) = value.get("thoughts").filter(|value| !value.is_null()) {
                    message
                        .metadata
                        .insert("gemini_thoughts".into(), thoughts.to_string());
                }
                restore_gemini_message_extension(&value, &mut message);
                if message.content.is_some() || message.tool_calls.is_some() {
                    messages.push(message);
                }
                continue;
            }

            let mut user_parts = Vec::new();
            if let Some(parts) = parts {
                for part in parts {
                    if let Some(response) = part.get("functionResponse") {
                        push_gemini_user_parts(
                            &mut messages,
                            std::mem::take(&mut user_parts),
                            timestamp,
                            &value,
                        );
                        let name = response
                            .get("name")
                            .and_then(Value::as_str)
                            .unwrap_or("tool")
                            .to_string();
                        let explicit_id = response
                            .get("id")
                            .and_then(Value::as_str)
                            .map(str::to_string);
                        if let Some(id) = explicit_id.as_deref() {
                            if let Some(ids) = pending_by_name.get_mut(&name) {
                                if let Some(position) = ids.iter().position(|pending| pending == id)
                                {
                                    ids.remove(position);
                                }
                            }
                        }
                        let id = explicit_id
                            .or_else(|| {
                                pending_by_name
                                    .get_mut(&name)
                                    .and_then(|ids| (!ids.is_empty()).then(|| ids.remove(0)))
                            })
                            .unwrap_or_else(|| format!("gemini-{line_index}-response"));
                        let output = response
                            .get("response")
                            .and_then(|response| response.get("output"))
                            .map(|output| {
                                output
                                    .as_str()
                                    .map(str::to_string)
                                    .unwrap_or_else(|| output.to_string())
                            })
                            .or_else(|| response.get("response").map(Value::to_string))
                            .unwrap_or_default();
                        let mut message = tool_message(&id, output);
                        message.name = Some(name);
                        if let Some(timestamp) = timestamp {
                            message
                                .metadata
                                .insert("timestamp".into(), timestamp.into());
                        }
                        restore_gemini_message_extension(&value, &mut message);
                        messages.push(message);
                        continue;
                    }
                    if let Some(text) = part.get("text").and_then(Value::as_str) {
                        user_parts.push(serde_json::json!({"type": "text", "text": text}));
                        continue;
                    }
                    if let Some(inline) = part.get("inlineData") {
                        let Some(data) = inline.get("data").and_then(Value::as_str) else {
                            continue;
                        };
                        let media_type = inline
                            .get("mimeType")
                            .and_then(Value::as_str)
                            .unwrap_or("application/octet-stream");
                        user_parts.push(serde_json::json!({
                            "type": "image_url",
                            "image_url": {"url": format!("data:{media_type};base64,{data}")},
                        }));
                    }
                }
            } else if !text.is_empty() {
                user_parts.push(serde_json::json!({"type": "text", "text": text}));
            }
            push_gemini_user_parts(&mut messages, user_parts, timestamp, &value);
        }

        if let Some(&last) = raw_record.last() {
            stamp_message_records(&mut messages, last);
        }
        ensure_tool_results_paired(&mut messages);
        meta.message_records = take_message_records(&mut messages);
        let imported_message_count = Some(messages.len());
        Ok(Session {
            meta,
            messages,
            subagents: Vec::new(),
            raw,
            raw_trailing_newline,
            imported_message_count,
            raw_is_verbatim: true,
            parse_error_lines,
            load_residue: Vec::new(),
        })
    }
}

const SUPERCODE_GEMINI_MESSAGE_KEY: &str = "_supercode_gemini_message";

fn set_gemini_message_extension(value: &mut Value, message: &ChatMessage) {
    value[SUPERCODE_GEMINI_MESSAGE_KEY] = serde_json::json!({
        "schema": 1,
        "role": message.role,
        "content": message.content,
        "content_parts": message.content_parts,
        "tool_calls": message.tool_calls,
        "tool_call_id": message.tool_call_id,
        "name": message.name,
        "metadata": message.metadata,
    });
}

pub(super) fn restore_gemini_message_extension(value: &Value, message: &mut ChatMessage) {
    let Some(extension) = value.get(SUPERCODE_GEMINI_MESSAGE_KEY) else {
        return;
    };
    if extension.get("schema").and_then(Value::as_u64) != Some(1) {
        return;
    }
    if let Some(role) = extension
        .get("role")
        .and_then(|value| serde_json::from_value(value.clone()).ok())
    {
        message.role = role;
    }
    message.content = extension
        .get("content")
        .and_then(Value::as_str)
        .map(str::to_string);
    message.content_parts = extension
        .get("content_parts")
        .and_then(|value| serde_json::from_value(value.clone()).ok());
    message.tool_calls = extension
        .get("tool_calls")
        .and_then(|value| serde_json::from_value(value.clone()).ok());
    message.tool_call_id = extension
        .get("tool_call_id")
        .and_then(Value::as_str)
        .map(str::to_string);
    message.name = extension
        .get("name")
        .and_then(Value::as_str)
        .map(str::to_string);
    message.metadata.clear();
    if let Some(metadata) = extension.get("metadata").and_then(Value::as_object) {
        for (key, value) in metadata {
            if let Some(value) = value.as_str() {
                message.metadata.insert(key.clone(), value.to_string());
            }
        }
    }
}

impl Session {
    // ---- Gemini writers ---------------------------------------------

    pub(super) fn to_gemini_jsonl(&self) -> String {
        let mut out = String::new();
        self.write_gemini_header(&mut out, self.meta.session_id.as_deref());
        self.write_gemini_records(&mut out, &self.messages);
        push_jsonl(
            &mut out,
            &serde_json::json!({
                "$set": {"lastUpdated": SYNTH_TS}
            }),
        );
        out
    }

    fn write_gemini_header(&self, out: &mut String, session_id: Option<&str>) {
        push_jsonl(
            out,
            &serde_json::json!({
                "sessionId": session_id.unwrap_or("supercode-gemini-session"),
                "projectHash": self.meta.lineage.get("gemini_project_hash")
                    .cloned().unwrap_or_else(|| "supercode".to_string()),
                "startTime": self.meta.lineage.get("created_at")
                    .cloned().unwrap_or_else(|| SYNTH_TS.to_string()),
                "lastUpdated": self.meta.lineage.get("updated_at")
                    .cloned().unwrap_or_else(|| SYNTH_TS.to_string()),
                "kind": self.meta.lineage.get("gemini_session_kind")
                    .cloned().unwrap_or_else(|| "main".to_string()),
            }),
        );
    }

    fn write_gemini_records(&self, out: &mut String, messages: &[ChatMessage]) {
        let mut call_names = HashMap::new();
        for (index, message) in messages.iter().enumerate() {
            if is_replay_excluded(message) {
                continue;
            }
            let timestamp = message
                .metadata
                .get("timestamp")
                .cloned()
                .unwrap_or_else(|| SYNTH_TS.to_string());
            match message.role {
                Role::System | Role::User => {
                    let mut parts = Vec::new();
                    let text = message.content.clone().or_else(|| {
                        message.content_parts.as_ref().and_then(|parts| {
                            let text = parts
                                .iter()
                                .filter_map(|part| part.get("text").and_then(Value::as_str))
                                .collect::<Vec<_>>()
                                .join(" ");
                            (!text.is_empty()).then_some(text)
                        })
                    });
                    if let Some(text) = text {
                        let text = if message.role == Role::System {
                            format!("[System] {text}")
                        } else {
                            text
                        };
                        parts.push(serde_json::json!({"text": text}));
                    }
                    if let Some(content_parts) = &message.content_parts {
                        for part in content_parts {
                            let Some(url) = part
                                .get("image_url")
                                .and_then(|value| value.get("url"))
                                .and_then(Value::as_str)
                            else {
                                continue;
                            };
                            let Some(rest) = url.strip_prefix("data:") else {
                                continue;
                            };
                            let Some((media_type, data)) = rest.split_once(";base64,") else {
                                continue;
                            };
                            parts.push(serde_json::json!({
                                "inlineData": {"mimeType": media_type, "data": data}
                            }));
                        }
                    }
                    if !parts.is_empty() {
                        let mut value = serde_json::json!({
                            "id": format!("supercode-user-{index}"),
                            "timestamp": timestamp,
                            "type": "user",
                            "content": parts,
                        });
                        set_gemini_message_extension(&mut value, message);
                        push_jsonl(out, &value);
                    }
                }
                Role::Assistant => {
                    let mut tool_calls = Vec::new();
                    for call in message.tool_calls() {
                        call_names.insert(call.id.clone(), call.function.name.clone());
                        let args = serde_json::from_str::<Value>(&call.function.arguments)
                            .unwrap_or_else(|_| Value::String(call.function.arguments.clone()));
                        tool_calls.push(serde_json::json!({
                            "id": call.id,
                            "name": call.function.name,
                            "args": args,
                        }));
                    }
                    let mut value = serde_json::json!({
                        "id": format!("supercode-gemini-{index}"),
                        "timestamp": timestamp,
                        "type": "gemini",
                        "content": message.content.clone().unwrap_or_default(),
                        "model": message.metadata.get("gemini_model")
                            .or(self.meta.model.as_ref())
                            .cloned().unwrap_or_else(|| "unknown".to_string()),
                    });
                    if !tool_calls.is_empty() {
                        value["toolCalls"] = Value::Array(tool_calls);
                    }
                    if let Some(thoughts) = message.metadata.get("gemini_thoughts") {
                        value["thoughts"] = serde_json::from_str(thoughts)
                            .unwrap_or_else(|_| Value::String(thoughts.clone()));
                    }
                    set_gemini_message_extension(&mut value, message);
                    push_jsonl(out, &value);
                }
                Role::Tool => {
                    let id = message.tool_call_id.clone().unwrap_or_default();
                    let name = message
                        .name
                        .clone()
                        .or_else(|| call_names.get(&id).cloned())
                        .unwrap_or_else(|| "tool".to_string());
                    let output = message.content.clone().unwrap_or_else(|| {
                        message
                            .content_parts
                            .as_ref()
                            .map(|parts| Value::Array(parts.clone()))
                            .map(|value| value.to_string())
                            .unwrap_or_default()
                    });
                    let mut value = serde_json::json!({
                        "id": format!("supercode-tool-{index}"),
                        "timestamp": timestamp,
                        "type": "user",
                        "content": [{
                            "functionResponse": {
                                "id": id,
                                "name": name,
                                "response": {"output": output}
                            }
                        }],
                    });
                    set_gemini_message_extension(&mut value, message);
                    push_jsonl(out, &value);
                }
            }
        }
    }

    pub(super) fn to_gemini_jsonl_spliced(&self, session_id: Option<&str>) -> String {
        let (raw_prefix_len, message_prefix_len) = self.spliced_prefix_lens();
        if raw_prefix_len == 0 {
            let mut out = String::new();
            self.write_gemini_header(&mut out, session_id.or(self.meta.session_id.as_deref()));
            self.write_gemini_records(&mut out, &self.messages);
            return out;
        }
        let mut out = String::new();
        for (index, line) in self.raw[..raw_prefix_len].iter().enumerate() {
            if index == 0 && session_id.is_some() {
                if let Ok(mut value) = serde_json::from_str::<Value>(line) {
                    if value.get("type").is_none() && value.get("sessionId").is_some() {
                        value["sessionId"] =
                            Value::String(session_id.unwrap_or_default().to_string());
                        push_jsonl(&mut out, &value);
                        continue;
                    }
                }
            }
            out.push_str(line);
            out.push('\n');
        }
        self.write_gemini_records(&mut out, &self.messages[message_prefix_len..]);
        out
    }
}

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

    #[test]
    fn gemini_user_parts_preserve_text_media_and_response_order() {
        let session = Session::from_gemini_str(
            r#"{"sessionId":"11111111-1111-4111-8111-111111111111","projectHash":"p"}
{"type":"gemini","content":"calling","toolCalls":[{"id":"a","name":"read","args":{}},{"id":"b","name":"read","args":{}}]}
{"type":"user","content":[{"text":"before"},{"functionResponse":{"id":"b","name":"read","response":{"output":"B"}}},{"inlineData":{"mimeType":"image/png","data":"YQ=="}},{"functionResponse":{"name":"read","response":{"output":"A"}}},{"text":"after"}]}
"#,
        )
        .unwrap();

        assert_eq!(session.messages.len(), 6);
        assert_eq!(
            session.messages[1].content_parts.as_ref().unwrap()[0]["text"],
            "before"
        );
        assert_eq!(session.messages[2].tool_call_id.as_deref(), Some("b"));
        assert!(
            session.messages[3].content_parts.as_ref().unwrap()[0]["image_url"]["url"]
                .as_str()
                .unwrap()
                .starts_with("data:image/png;base64,")
        );
        assert_eq!(session.messages[4].tool_call_id.as_deref(), Some("a"));
        assert_eq!(
            session.messages[5].content_parts.as_ref().unwrap()[0]["text"],
            "after"
        );
    }
}