supercode_interchange/session/

grok.rs

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

use super::*;

impl Session {
    /// Load Grok's resumable `chat_history.jsonl` transcript.
    ///
    /// The surrounding session directory carries the session id, workspace,
    /// and `summary.json`; [`Self::from_grok_str`] handles the transcript
    /// itself while this path-aware entry point overlays that directory
    /// metadata.
    pub fn from_grok(path: impl AsRef<Path>) -> Result<Session> {
        let path = path.as_ref();
        let mut session = Self::from_grok_str(&std::fs::read_to_string(path)?)?;
        session.capture_grok_path_metadata(path);
        Ok(session)
    }

    /// Parse Grok's line-oriented `chat_history.jsonl` format.
    ///
    /// Conversational records are `user`, `assistant`, and `tool_result`.
    /// `system` is the regenerated base prompt and is retained in
    /// [`SessionMeta::system_prompt`]; encrypted reasoning and backend-only
    /// state remain byte-exact in [`Session::raw`] but are intentionally not
    /// replayed as chat turns.
    pub fn from_grok_str(jsonl: &str) -> Result<Session> {
        let mut meta = SessionMeta::new(SessionSource::Grok);
        let (raw_lines, raw_trailing_newline) = split_lines_verbatim(jsonl);
        let raw: Vec<String> = raw_lines.iter().map(|line| line.to_string()).collect();
        let mut messages = Vec::new();
        let mut parse_error_lines = 0usize;
        let mut tool_names: HashMap<String, 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 (record_index, line) in non_empty_lines(jsonl).enumerate() {
            // What the previous record created belongs to it; this record starts after.
            if record_index > 0 {
                stamp_message_records(&mut messages, raw_record[record_index - 1]);
            }
            let value: Value = match serde_json::from_str(line) {
                Ok(value) => value,
                Err(_) => {
                    parse_error_lines += 1;
                    continue;
                }
            };
            restore_codex_provenance_from_top_level(&value, &mut meta)?;
            // PARITY-23: records the open-union arm below would drop are
            // grok's residue inventory — captured for cross-format hops.
            if let Some(kind) = grok_residue_kind(&value) {
                capture_native_residue(&mut meta, "grok", record_index, line, &value, &kind);
            }
            match value.get("type").and_then(Value::as_str) {
                Some("system") => {
                    if meta.system_prompt.is_none() {
                        meta.system_prompt = value
                            .get("content")
                            .and_then(Value::as_str)
                            .map(str::to_string);
                    }
                }
                Some("user") => {
                    let content = extract_text_content(value.get("content"));
                    let role = if value.get("synthetic_reason").and_then(Value::as_str)
                        == Some("supercode_system_event")
                    {
                        Role::System
                    } else {
                        Role::User
                    };
                    let content = if role == Role::User {
                        match grok_human_user_text(&content) {
                            Some(content) => content,
                            None if value.get(SUPERCODE_GROK_MESSAGE_KEY).is_some() => {
                                String::new()
                            }
                            None => continue,
                        }
                    } else {
                        content
                    };
                    let mut message = ChatMessage {
                        role,
                        content: Some(content),
                        content_parts: None,
                        tool_calls: None,
                        tool_call_id: None,
                        name: None,
                        metadata: Default::default(),
                    };
                    capture_grok_scalar_metadata(
                        &value,
                        &mut message,
                        &["prompt_index", "prior_turn_interrupt", "synthetic_reason"],
                    );
                    restore_grok_message_extension(&value, &mut message);
                    messages.push(message);
                }
                Some("assistant") => {
                    let calls: Vec<ToolCall> = value
                        .get("tool_calls")
                        .and_then(Value::as_array)
                        .into_iter()
                        .flatten()
                        .filter_map(|call| {
                            let id = call.get("id")?.as_str()?.to_string();
                            let name = call.get("name")?.as_str()?.to_string();
                            let arguments = call
                                .get("arguments")
                                .map(value_to_arg_string)
                                .unwrap_or_else(|| "{}".to_string());
                            tool_names.insert(id.clone(), name.clone());
                            Some(function_call(&id, &name, arguments))
                        })
                        .collect();
                    let content = value
                        .get("content")
                        .and_then(Value::as_str)
                        .filter(|content| !content.is_empty())
                        .map(str::to_string);
                    let mut message = ChatMessage {
                        role: Role::Assistant,
                        content,
                        content_parts: None,
                        tool_calls: (!calls.is_empty()).then_some(calls),
                        tool_call_id: None,
                        name: None,
                        metadata: Default::default(),
                    };
                    capture_grok_scalar_metadata(
                        &value,
                        &mut message,
                        &["model_id", "model_fingerprint", "reasoning_effort"],
                    );
                    if let Some(model) = value.get("model_id").and_then(Value::as_str) {
                        meta.model = Some(model.to_string());
                    }
                    restore_grok_message_extension(&value, &mut message);
                    messages.push(message);
                }
                Some("tool_result") => {
                    let id = value
                        .get("tool_call_id")
                        .and_then(Value::as_str)
                        .unwrap_or_default();
                    let content = value
                        .get("content")
                        .map(|value| match value {
                            Value::String(text) => text.clone(),
                            other => extract_text_content(Some(other)),
                        })
                        .unwrap_or_default();
                    let mut message = tool_message(id, content);
                    message.name = tool_names.get(id).cloned();
                    // A tool result's screenshots: `images: [{type: "image", url: <data URI>}]`.
                    let images: Vec<Value> = value
                        .get("images")
                        .and_then(Value::as_array)
                        .into_iter()
                        .flatten()
                        .filter_map(|image| image.get("url").and_then(Value::as_str))
                        .filter(|url| !url.is_empty())
                        .map(|url| serde_json::json!({"type": "image_url", "image_url": {"url": url}}))
                        .collect();
                    if !images.is_empty() {
                        let mut parts = Vec::new();
                        if let Some(text) = message.content.as_deref().filter(|t| !t.is_empty()) {
                            parts.push(serde_json::json!({"type": "text", "text": text}));
                        }
                        parts.extend(images);
                        message.content_parts = Some(parts);
                    }
                    restore_grok_message_extension(&value, &mut message);
                    messages.push(message);
                }
                // `reasoning` contains encrypted chain-of-thought and
                // `backend_tool_call` is execution bookkeeping. Both survive
                // verbatim in raw without being replayed to another model.
                _ => {}
            }
        }

        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(),
        })
    }

    pub(super) fn capture_grok_path_metadata(&mut self, transcript: &Path) {
        let Some(session_dir) = transcript.parent() else {
            return;
        };
        self.meta.session_id = session_dir
            .file_name()
            .and_then(|name| name.to_str())
            .map(str::to_string);
        self.meta.cwd = session_dir
            .parent()
            .and_then(Path::file_name)
            .and_then(|name| name.to_str())
            .and_then(percent_decode_path)
            .map(PathBuf::from);

        let Ok(summary_text) = std::fs::read_to_string(session_dir.join("summary.json")) else {
            return;
        };
        let Ok(summary) = serde_json::from_str::<Value>(&summary_text) else {
            return;
        };
        if let Some(model) = summary.get("current_model_id").and_then(Value::as_str) {
            self.meta.model = Some(model.to_string());
        }
        for (source, target) in [
            ("generated_title", "session_name"),
            ("created_at", "created_at"),
            ("updated_at", "updated_at"),
            ("chat_format_version", "grok_chat_format_version"),
        ] {
            if let Some(value) = summary.get(source) {
                self.meta.lineage.insert(
                    target.to_string(),
                    value
                        .as_str()
                        .map(str::to_string)
                        .unwrap_or_else(|| value.to_string()),
                );
            }
        }
    }
}

/// Record types grok's loader consumes into the canonical model; anything
/// else on a grok transcript is residue (the loader's open-union arm is the
/// authoritative inventory, per the PARITY-23 design doc).
const GROK_CONSUMED_TYPES: [&str; 4] = ["system", "user", "assistant", "tool_result"];

pub(super) fn grok_residue_kind(record: &Value) -> Option<String> {
    record.as_object()?;
    match record.get("type").and_then(Value::as_str) {
        Some(kind) if !GROK_CONSUMED_TYPES.contains(&kind) => Some(kind.to_string()),
        Some(_) => None,
        None => Some("untyped".to_string()),
    }
}

// ---- Grok -------------------------------------------------------------

fn capture_grok_scalar_metadata(value: &Value, message: &mut ChatMessage, keys: &[&str]) {
    for key in keys {
        if let Some(value) = value.get(*key) {
            message.metadata.insert(
                format!("grok_{key}"),
                value
                    .as_str()
                    .map(str::to_string)
                    .unwrap_or_else(|| value.to_string()),
            );
        }
    }
}

fn grok_human_user_text(raw: &str) -> Option<String> {
    let text = raw.trim();
    if text.is_empty() || text.starts_with("<user_info>") || text.starts_with("<system-reminder>") {
        return None;
    }
    let unwrapped = text
        .strip_prefix("<user_query>")
        .and_then(|value| value.strip_suffix("</user_query>"))
        .map(str::trim)
        .unwrap_or(text);
    (!unwrapped.is_empty()).then(|| unwrapped.to_string())
}

/// Portable extension for messages whose canonical fields cannot be expressed
/// by the target's stock schema. It was introduced for Grok and retains that
/// on-disk key for compatibility. Gemini has the same need: Claude Code and
/// Codex have no native slot for a tool-result name or Gemini-only metadata.
/// Their readers tolerate unknown namespaced fields, so forwarding this
/// adapter-owned envelope keeps those cross-format hops reversible without
/// pretending the stock schemas represent the fields directly.
const SUPERCODE_GROK_MESSAGE_KEY: &str = "_supercode_grok_message";

fn grok_message_extension(source: SessionSource, message: &ChatMessage) -> Option<Value> {
    let metadata = message
        .metadata
        .iter()
        .filter(|(key, _)| {
            key.starts_with("gemini_") || key.starts_with("grok_") || key.starts_with("goose_")
        })
        .map(|(key, value)| (key.clone(), Value::String(value.clone())))
        .collect::<serde_json::Map<_, _>>();

    // `meta.source` changes after every reload. Keying portability only on
    // the immediate source therefore made Grok metadata survive one hop but
    // disappear on A -> B -> C translations. Once Grok-owned fields are
    // present, keep forwarding them regardless of the current container.
    let has_portable_fields =
        !metadata.is_empty() || message.metadata.contains_key("codex_custom_tool_call_ids");
    (matches!(
        source,
        SessionSource::Gemini | SessionSource::Grok | SessionSource::Goose
    ) || has_portable_fields
        || message.content_parts.is_some())
    .then(|| {
        serde_json::json!({
            "schema": 2,
            "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 set_grok_target_message_extension(value: &mut Value, message: &ChatMessage) {
    value[SUPERCODE_GROK_MESSAGE_KEY] = serde_json::json!({
        "schema": 2,
        "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 set_grok_message_extension(
    value: &mut Value,
    source: SessionSource,
    message: &ChatMessage,
) {
    if let Some(extension) = grok_message_extension(source, message) {
        value[SUPERCODE_GROK_MESSAGE_KEY] = extension;
    }
}

pub(super) fn restore_grok_message_extension(value: &Value, message: &mut ChatMessage) {
    let Some(extension) = value.get(SUPERCODE_GROK_MESSAGE_KEY) else {
        return;
    };
    // Codex temporarily marks a text assistant item so immediately-following
    // function-call items can merge back into the same canonical turn. The
    // portable envelope must not erase that loader-private marker before the
    // merge happens; `from_codex_str` removes it before returning.
    let codex_open_turn = message.metadata.get("__codex_open_turn").cloned();
    let codex_turn_id = message.metadata.get("turn_id").cloned();
    let extension_has_turn_id = extension
        .get("metadata")
        .and_then(Value::as_object)
        .is_some_and(|metadata| metadata.contains_key("turn_id"));
    if extension.get("schema").and_then(Value::as_u64) == Some(2) {
        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());
        // Tool calls are shared native structure in every supported format.
        // Keep the loader's reconstruction instead of restoring this copy:
        // Codex stores a combined text+tool turn across multiple records, so
        // eagerly restoring calls on its text record would duplicate them
        // when the following function-call records merge.
        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());
            }
        }
    }
    if let Some(name) = extension.get("name").and_then(Value::as_str) {
        message.name = Some(name.to_string());
    }
    if let Some(marker) = codex_open_turn {
        message
            .metadata
            .insert("__codex_open_turn".to_string(), marker);
    }
    if let Some(turn_id) = codex_turn_id {
        message.metadata.insert("turn_id".to_string(), turn_id);
        if !extension_has_turn_id {
            message.metadata.insert(
                "__grok_remove_synthetic_turn_id".to_string(),
                "true".to_string(),
            );
        }
    }
}

pub(super) fn restore_single_grok_message(value: &Value, messages: &mut [ChatMessage]) {
    if let [message] = messages {
        restore_grok_message_extension(value, message);
    }
}

impl Session {
    // ---- Grok writers -----------------------------------------------

    /// Synthesize Grok's resumable `chat_history.jsonl` transcript.
    pub(super) fn to_grok_jsonl(&self) -> String {
        let mut out = String::new();
        if let Some(prompt) = self
            .meta
            .system_prompt
            .as_deref()
            .filter(|prompt| !prompt.is_empty())
        {
            push_jsonl(
                &mut out,
                &serde_json::json!({
                    "type": "system",
                    "content": prompt,
                }),
            );
        }
        self.write_grok_records(&mut out, &self.messages);
        // PARITY-23: grok-source residue restored from a foreign hop is
        // NATIVE here again — re-emit the exact source records (relative
        // order preserved) instead of wrapping them in an envelope.
        if self.meta.native_residue_source.as_deref() == Some("grok") {
            let mut records: Vec<&Value> = self.meta.native_residue.iter().collect();
            records.sort_by_key(|entry| {
                entry
                    .get("record_index")
                    .and_then(Value::as_u64)
                    .unwrap_or(u64::MAX)
            });
            for entry in records {
                if let Some(raw) = entry.get("raw").and_then(Value::as_str) {
                    out.push_str(raw);
                    out.push('\n');
                }
            }
        } else if let Some(extension) = native_residue_envelope(&self.meta) {
            if out.is_empty() {
                push_jsonl(
                    &mut out,
                    &serde_json::json!({"type": "system", "content": ""}),
                );
            }
            inject_first_jsonl_top_level(
                &mut out,
                SUPERCODE_NATIVE_RESIDUE_SUMMARY_KEY,
                native_residue_summary(&extension),
            );
            inject_first_jsonl_top_level(&mut out, SUPERCODE_NATIVE_RESIDUE_KEY, extension);
        }
        out
    }

    fn write_grok_records(&self, out: &mut String, messages: &[ChatMessage]) {
        for message in messages {
            if is_replay_excluded(message) {
                continue;
            }
            let mut value = match message.role {
                Role::System => serde_json::json!({
                    "type": "user",
                    "content": [{"type": "text", "text": message.content.clone().unwrap_or_default()}],
                    "synthetic_reason": "supercode_system_event",
                }),
                Role::User => {
                    let mut value = serde_json::json!({
                        "type": "user",
                        "content": [{"type": "text", "text": message.content.clone().unwrap_or_default()}],
                    });
                    if let Some(object) = value.as_object_mut() {
                        for (metadata, field) in [
                            ("grok_prompt_index", "prompt_index"),
                            ("grok_prior_turn_interrupt", "prior_turn_interrupt"),
                            ("grok_synthetic_reason", "synthetic_reason"),
                        ] {
                            if let Some(raw) = message.metadata.get(metadata) {
                                object.insert(
                                    field.to_string(),
                                    serde_json::from_str(raw)
                                        .unwrap_or_else(|_| Value::String(raw.clone())),
                                );
                            }
                        }
                    }
                    value
                }
                Role::Assistant => {
                    let calls = message
                        .tool_calls()
                        .iter()
                        .map(|call| {
                            serde_json::json!({
                                "id": call.id,
                                "name": call.function.name,
                                "arguments": call.function.arguments,
                            })
                        })
                        .collect::<Vec<_>>();
                    let mut value = serde_json::json!({
                        "type": "assistant",
                        "content": message.content.clone().unwrap_or_default(),
                        "tool_calls": calls,
                        "model_id": message.metadata.get("grok_model_id")
                            .or(self.meta.model.as_ref())
                            .cloned()
                            .unwrap_or_else(|| "unknown".to_string()),
                    });
                    if let Some(object) = value.as_object_mut() {
                        for (metadata, field) in [
                            ("grok_model_fingerprint", "model_fingerprint"),
                            ("grok_reasoning_effort", "reasoning_effort"),
                        ] {
                            if let Some(raw) = message.metadata.get(metadata) {
                                object.insert(field.to_string(), Value::String(raw.clone()));
                            }
                        }
                    }
                    value
                }
                Role::Tool => {
                    let mut value = serde_json::json!({
                        "type": "tool_result",
                        "tool_call_id": message.tool_call_id.clone().unwrap_or_default(),
                        "content": message.content.clone().unwrap_or_default(),
                    });
                    let images: Vec<Value> = message
                        .content_parts
                        .iter()
                        .flatten()
                        .filter(|part| {
                            part.get("type").and_then(Value::as_str) == Some("image_url")
                        })
                        .filter_map(|part| {
                            part.get("image_url")
                                .and_then(|u| u.get("url"))
                                .and_then(Value::as_str)
                        })
                        .map(|url| serde_json::json!({"type": "image", "url": url}))
                        .collect();
                    if !images.is_empty() {
                        value["images"] = Value::Array(images);
                    }
                    value
                }
            };
            set_grok_target_message_extension(&mut value, message);
            push_jsonl(out, &value);
        }
    }

    /// Replay a Grok imported prefix verbatim, then append newly-created
    /// canonical turns. Grok stores the session id in the directory name,
    /// not in transcript records, so there is no in-file id to rewrite.
    pub(super) fn to_grok_jsonl_spliced(&self) -> String {
        let (raw_prefix_len, message_prefix_len) = self.spliced_prefix_lens();
        if raw_prefix_len == 0 {
            return self.to_grok_jsonl();
        }
        let mut out = String::new();
        for line in &self.raw[..raw_prefix_len] {
            out.push_str(line);
            out.push('\n');
        }
        self.write_grok_records(&mut out, &self.messages[message_prefix_len..]);
        out
    }
}