mi6_core/input/codex_session/

entry.rs

//! Codex session entry types and parsing.
//!
//! Codex CLI stores session data in JSONL files at:
//! `~/.codex/sessions/YYYY/MM/DD/rollout-YYYY-MM-DDTHH-MM-SS-<UUID>.jsonl`
//!
//! The format differs from Claude Code transcripts:
//! - Top-level `type` field indicates entry type
//! - `payload` contains type-specific data
//! - Token usage is in nested `event_msg.token_count` events

use std::collections::HashMap;

use chrono::{DateTime, Utc};
use serde::Deserialize;

use crate::model::{Event, EventBuilder, EventType};
use crate::util::truncate_string;

/// Information about a function call, tracked for enriching output events.
#[derive(Debug, Clone)]
pub struct FunctionCallInfo {
    /// The function name (e.g., "shell_command")
    pub name: String,
    /// The function arguments as JSON string
    pub arguments: String,
    /// Whether this call requires user approval (sandbox_permissions: "require_escalated")
    pub requires_approval: bool,
    /// Timestamp of the function call
    pub timestamp: Option<DateTime<Utc>>,
}

/// Map from call_id to function call info.
pub type FunctionCallMap = HashMap<String, FunctionCallInfo>;

/// A single entry in a Codex session JSONL file.
///
/// Codex uses a unified entry format with `type` discriminator.
#[derive(Debug, Clone, Deserialize)]
pub struct CodexSessionEntry {
    /// Entry timestamp (ISO 8601)
    pub timestamp: String,

    /// Entry type: "session_meta", "response_item", "event_msg", "turn_context"
    #[serde(rename = "type")]
    pub entry_type: String,

    /// Type-specific payload
    pub payload: serde_json::Value,
}

/// Session metadata from the first entry.
#[derive(Debug, Clone, Deserialize)]
pub struct SessionMeta {
    /// Session UUID
    pub id: String,
    /// Session start timestamp
    pub timestamp: String,
    /// Working directory
    pub cwd: Option<String>,
    /// CLI version
    pub cli_version: Option<String>,
    /// Model provider (e.g., "openai")
    pub model_provider: Option<String>,
    /// Git context
    pub git: Option<GitInfo>,
}

/// Git repository information.
#[derive(Debug, Clone, Deserialize)]
pub struct GitInfo {
    pub commit_hash: Option<String>,
    pub branch: Option<String>,
    pub repository_url: Option<String>,
}

/// Token usage information from token_count events.
#[derive(Debug, Clone, Deserialize)]
pub struct TokenUsageInfo {
    pub total_token_usage: Option<TokenUsage>,
    pub last_token_usage: Option<TokenUsage>,
    pub model_context_window: Option<i64>,
}

/// Token counts.
#[derive(Debug, Clone, Deserialize)]
pub struct TokenUsage {
    pub input_tokens: Option<i64>,
    pub cached_input_tokens: Option<i64>,
    pub output_tokens: Option<i64>,
    pub reasoning_output_tokens: Option<i64>,
    pub total_tokens: Option<i64>,
}

/// Turn context information.
#[derive(Debug, Clone, Deserialize)]
pub struct TurnContext {
    pub cwd: Option<String>,
    pub model: Option<String>,
    pub approval_policy: Option<String>,
}

impl CodexSessionEntry {
    /// Parse timestamp from the entry.
    fn parse_timestamp(&self) -> Option<DateTime<Utc>> {
        chrono::DateTime::parse_from_rfc3339(&self.timestamp)
            .ok()
            .map(|dt| dt.with_timezone(&Utc))
    }

    /// Convert this entry to Event(s).
    ///
    /// A single entry may produce multiple events depending on type.
    /// The `function_call_map` tracks call_id -> function info for enriching outputs.
    /// The `session_meta` provides session-level context like cwd, git branch.
    /// The `turn_model` tracks the model from the most recent turn_context.
    pub fn into_events(
        self,
        machine_id: &str,
        session_id: &str,
        function_call_map: &mut FunctionCallMap,
        session_meta: &Option<SessionMeta>,
        turn_model: &mut Option<String>,
    ) -> Vec<Event> {
        let timestamp = self.parse_timestamp();

        match self.entry_type.as_str() {
            "session_meta" => {
                // session_meta is handled separately to set session context
                // No events generated here
                vec![]
            }

            "turn_context" => {
                // Update turn model for subsequent events
                if let Ok(ctx) = serde_json::from_value::<TurnContext>(self.payload) {
                    *turn_model = ctx.model;
                }
                vec![]
            }

            "event_msg" => {
                self.parse_event_msg(machine_id, session_id, timestamp, session_meta, turn_model)
            }

            "response_item" => self.parse_response_item(
                machine_id,
                session_id,
                timestamp,
                function_call_map,
                session_meta,
                turn_model,
            ),

            _ => vec![],
        }
    }

    /// Parse event_msg entries (token_count, user_message, agent_message, etc.)
    fn parse_event_msg(
        self,
        machine_id: &str,
        session_id: &str,
        timestamp: Option<DateTime<Utc>>,
        session_meta: &Option<SessionMeta>,
        turn_model: &Option<String>,
    ) -> Vec<Event> {
        let msg_type = self
            .payload
            .get("type")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        match msg_type {
            "token_count" => {
                // Parse token usage and create ApiRequest event
                let info = self.payload.get("info");
                if info.is_none() || info.is_some_and(|v| v.is_null()) {
                    return vec![];
                }

                let info: TokenUsageInfo =
                    match serde_json::from_value(info.cloned().unwrap_or_default()) {
                        Ok(i) => i,
                        Err(_) => return vec![],
                    };

                // Use last_token_usage for per-turn incremental counts
                let usage = match info.last_token_usage {
                    Some(u) => u,
                    None => return vec![],
                };

                // Skip if no meaningful token data
                if usage.input_tokens.unwrap_or(0) == 0 && usage.output_tokens.unwrap_or(0) == 0 {
                    return vec![];
                }

                let mut builder = EventBuilder::new(machine_id, EventType::ApiRequest, session_id)
                    .source("codex_session")
                    .framework("codex")
                    .timestamp_opt(timestamp);

                // Set tokens
                let input = usage.input_tokens.unwrap_or(0);
                let output = usage.output_tokens.unwrap_or(0);
                builder = builder.tokens(input, output);

                // Set cache tokens (Codex uses cached_input_tokens)
                if let Some(cache_read) = usage.cached_input_tokens {
                    builder = builder.tokens_cache_read_opt(Some(cache_read));
                }

                // Set model from turn context
                builder = builder.model_opt(turn_model.clone());

                // Set cwd and git from session meta
                if let Some(meta) = session_meta {
                    builder = builder.cwd_opt(meta.cwd.clone());
                    if let Some(ref git) = meta.git {
                        builder = builder.git_branch_opt(git.branch.clone());
                    }
                }

                // Store extra info (reasoning_output_tokens, model_context_window) in metadata
                let mut extra = serde_json::Map::new();
                if let Some(reasoning) = usage.reasoning_output_tokens {
                    extra.insert(
                        "reasoning_output_tokens".into(),
                        serde_json::json!(reasoning),
                    );
                }
                if let Some(ctx_window) = info.model_context_window {
                    extra.insert("model_context_window".into(), serde_json::json!(ctx_window));
                }
                if !extra.is_empty() {
                    builder = builder.metadata(serde_json::Value::Object(extra).to_string());
                }

                vec![builder.build()]
            }

            "user_message" => {
                // Create UserPromptSubmit event
                let message = self
                    .payload
                    .get("message")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");

                if message.is_empty() {
                    return vec![];
                }

                let mut builder =
                    EventBuilder::new(machine_id, EventType::UserPromptSubmit, session_id)
                        .source("codex_session")
                        .framework("codex")
                        .timestamp_opt(timestamp);

                // Set cwd from session meta
                if let Some(meta) = session_meta {
                    builder = builder.cwd_opt(meta.cwd.clone());
                    if let Some(ref git) = meta.git {
                        builder = builder.git_branch_opt(git.branch.clone());
                    }
                }

                // Store truncated prompt in payload
                let truncated = truncate_string(message, 1000);
                builder = builder.payload(serde_json::json!({"prompt": truncated}).to_string());

                vec![builder.build()]
            }

            // agent_message, agent_reasoning, turn_aborted are informational
            // We don't create separate events for these
            _ => vec![],
        }
    }

    /// Parse response_item entries (function_call, function_call_output, message)
    fn parse_response_item(
        self,
        machine_id: &str,
        session_id: &str,
        timestamp: Option<DateTime<Utc>>,
        function_call_map: &mut FunctionCallMap,
        session_meta: &Option<SessionMeta>,
        turn_model: &Option<String>,
    ) -> Vec<Event> {
        let item_type = self
            .payload
            .get("type")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        match item_type {
            "function_call" => {
                let name = self
                    .payload
                    .get("name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let call_id = self
                    .payload
                    .get("call_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let arguments = self
                    .payload
                    .get("arguments")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();

                if call_id.is_empty() {
                    return vec![];
                }

                // Check if this call requires user approval
                // Codex sets sandbox_permissions: "require_escalated" when approval is needed
                let requires_approval = serde_json::from_str::<serde_json::Value>(&arguments)
                    .ok()
                    .and_then(|args| {
                        args.get("sandbox_permissions")
                            .and_then(|v| v.as_str())
                            .map(|s| s == "require_escalated")
                    })
                    .unwrap_or(false);

                // Track function call for enriching output
                function_call_map.insert(
                    call_id.clone(),
                    FunctionCallInfo {
                        name: name.clone(),
                        arguments: arguments.clone(),
                        requires_approval,
                        timestamp,
                    },
                );

                // Create PreToolUse event
                let mut builder = EventBuilder::new(machine_id, EventType::PreToolUse, session_id)
                    .source("codex_session")
                    .framework("codex")
                    .timestamp_opt(timestamp)
                    .tool(call_id, name)
                    .model_opt(turn_model.clone());

                if let Some(meta) = session_meta {
                    builder = builder.cwd_opt(meta.cwd.clone());
                    if let Some(ref git) = meta.git {
                        builder = builder.git_branch_opt(git.branch.clone());
                    }
                }

                vec![builder.build()]
            }

            "function_call_output" => {
                let call_id = self
                    .payload
                    .get("call_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string();
                let output = self
                    .payload
                    .get("output")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");

                if call_id.is_empty() {
                    return vec![];
                }

                // Create PostToolUse event
                let mut builder = EventBuilder::new(machine_id, EventType::PostToolUse, session_id)
                    .source("codex_session")
                    .framework("codex")
                    .timestamp_opt(timestamp)
                    .tool_use_id(call_id.clone())
                    .model_opt(turn_model.clone());

                // Enrich with function call info
                if let Some(call_info) = function_call_map.get(&call_id) {
                    builder = builder.tool_name(call_info.name.clone());

                    // For shell_command, store command in payload (for GitHub context)
                    if call_info.name == "shell_command" {
                        // Parse arguments JSON to extract command
                        if let Ok(args) =
                            serde_json::from_str::<serde_json::Value>(&call_info.arguments)
                            && let Some(cmd) = args.get("command")
                        {
                            builder =
                                builder.payload(serde_json::json!({ "command": cmd }).to_string());
                        }

                        // Store truncated output in metadata (for debugging)
                        let truncated_output = truncate_string(output, 500);
                        builder = builder.metadata(
                            serde_json::json!({ "output": truncated_output }).to_string(),
                        );
                    }
                }

                if let Some(meta) = session_meta {
                    builder = builder.cwd_opt(meta.cwd.clone());
                    if let Some(ref git) = meta.git {
                        builder = builder.git_branch_opt(git.branch.clone());
                    }
                }

                vec![builder.build()]
            }

            // message, reasoning, ghost_snapshot - no events for these
            _ => vec![],
        }
    }

    /// Get a unique key for deduplication.
    ///
    /// Combines timestamp and entry type since Codex doesn't have per-entry UUIDs.
    pub fn dedup_key(&self) -> String {
        format!("{}:{}", self.timestamp, self.entry_type)
    }
}

/// Extract session ID from a Codex session filename.
///
/// Filename format: `rollout-YYYY-MM-DDTHH-MM-SS-<UUID>.jsonl`
/// Returns the UUID portion.
pub fn extract_session_id_from_filename(filename: &str) -> Option<String> {
    // UUID is the last 36 characters before .jsonl
    // Format: rollout-2026-01-08T17-00-06-019ba044-90fa-7b30-be9c-6b7b601599cd.jsonl
    let without_ext = filename.strip_suffix(".jsonl")?;
    let parts: Vec<&str> = without_ext.rsplitn(6, '-').collect();

    // UUID has 5 parts: 8-4-4-4-12
    if parts.len() >= 5 {
        // Reconstruct UUID from last 5 parts (reversed)
        let uuid = format!(
            "{}-{}-{}-{}-{}",
            parts[4], parts[3], parts[2], parts[1], parts[0]
        );
        // Validate it looks like a UUID
        if uuid.len() == 36 && uuid.chars().filter(|c| *c == '-').count() == 4 {
            return Some(uuid);
        }
    }
    None
}

/// Parse session metadata from a session_meta entry.
pub fn parse_session_meta(entry: &CodexSessionEntry) -> Option<SessionMeta> {
    if entry.entry_type != "session_meta" {
        return None;
    }
    serde_json::from_value(entry.payload.clone()).ok()
}

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

    #[test]
    fn test_extract_session_id_from_filename() {
        assert_eq!(
            extract_session_id_from_filename(
                "rollout-2026-01-08T17-00-06-019ba044-90fa-7b30-be9c-6b7b601599cd.jsonl"
            ),
            Some("019ba044-90fa-7b30-be9c-6b7b601599cd".to_string())
        );

        assert_eq!(
            extract_session_id_from_filename(
                "rollout-2025-11-26T17-55-56-019ac306-3951-72d1-8b5d-447c33ce17f5.jsonl"
            ),
            Some("019ac306-3951-72d1-8b5d-447c33ce17f5".to_string())
        );

        assert_eq!(extract_session_id_from_filename("invalid.jsonl"), None);
        assert_eq!(extract_session_id_from_filename("rollout.jsonl"), None);
    }

    #[test]
    fn test_parse_session_meta() {
        let entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:55:56.451Z".to_string(),
            entry_type: "session_meta".to_string(),
            payload: serde_json::json!({
                "id": "019ac306-3951-72d1-8b5d-447c33ce17f5",
                "timestamp": "2025-11-27T01:55:56.369Z",
                "cwd": "/Users/storm/repos/test",
                "cli_version": "0.63.0",
                "model_provider": "openai",
                "git": {
                    "commit_hash": "abc123",
                    "branch": "main",
                    "repository_url": "git@github.com:test/repo.git"
                }
            }),
        };

        let meta = parse_session_meta(&entry).unwrap();
        assert_eq!(meta.id, "019ac306-3951-72d1-8b5d-447c33ce17f5");
        assert_eq!(meta.cwd, Some("/Users/storm/repos/test".to_string()));
        assert_eq!(meta.cli_version, Some("0.63.0".to_string()));
        assert_eq!(meta.git.as_ref().unwrap().branch, Some("main".to_string()));
    }

    #[test]
    fn test_parse_token_count_event() {
        let entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:10.186Z".to_string(),
            entry_type: "event_msg".to_string(),
            payload: serde_json::json!({
                "type": "token_count",
                "info": {
                    "total_token_usage": {
                        "input_tokens": 3597,
                        "cached_input_tokens": 3072,
                        "output_tokens": 33,
                        "reasoning_output_tokens": 0,
                        "total_tokens": 3630
                    },
                    "last_token_usage": {
                        "input_tokens": 500,
                        "cached_input_tokens": 400,
                        "output_tokens": 33,
                        "reasoning_output_tokens": 10,
                        "total_tokens": 533
                    },
                    "model_context_window": 258400
                }
            }),
        };

        let session_meta = Some(SessionMeta {
            id: "test".to_string(),
            timestamp: "2025-11-27T01:55:56.369Z".to_string(),
            cwd: Some("/test/cwd".to_string()),
            cli_version: None,
            model_provider: None,
            git: Some(GitInfo {
                branch: Some("main".to_string()),
                commit_hash: None,
                repository_url: None,
            }),
        });

        let mut turn_model = Some("gpt-5.1-codex".to_string());
        let mut call_map = FunctionCallMap::default();

        let events = entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &session_meta,
            &mut turn_model,
        );

        assert_eq!(events.len(), 1);
        let event = &events[0];
        assert_eq!(event.event_type, EventType::ApiRequest);
        assert_eq!(event.tokens_input, Some(500)); // Uses last_token_usage
        assert_eq!(event.tokens_output, Some(33));
        assert_eq!(event.tokens_cache_read, Some(400));
        assert_eq!(event.model, Some("gpt-5.1-codex".to_string()));
        assert_eq!(event.cwd, Some("/test/cwd".to_string()));
        assert_eq!(event.git_branch, Some("main".to_string()));

        // Check metadata for reasoning tokens
        let metadata: serde_json::Value =
            serde_json::from_str(event.metadata.as_ref().unwrap()).unwrap();
        assert_eq!(metadata["reasoning_output_tokens"], 10);
        assert_eq!(metadata["model_context_window"], 258400);
    }

    #[test]
    fn test_parse_user_message_event() {
        let entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:07.612Z".to_string(),
            entry_type: "event_msg".to_string(),
            payload: serde_json::json!({
                "type": "user_message",
                "message": "Hello, world!",
                "images": []
            }),
        };

        let mut call_map = FunctionCallMap::default();
        let mut turn_model = None;

        let events = entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &None,
            &mut turn_model,
        );

        assert_eq!(events.len(), 1);
        let event = &events[0];
        assert_eq!(event.event_type, EventType::UserPromptSubmit);

        let payload: serde_json::Value =
            serde_json::from_str(event.payload.as_ref().unwrap()).unwrap();
        assert_eq!(payload["prompt"], "Hello, world!");
    }

    #[test]
    fn test_parse_function_call_and_output() {
        let call_entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:10.186Z".to_string(),
            entry_type: "response_item".to_string(),
            payload: serde_json::json!({
                "type": "function_call",
                "name": "shell_command",
                "arguments": "{\"command\":\"ls\",\"workdir\":\"/test\"}",
                "call_id": "call_123"
            }),
        };

        let output_entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:10.300Z".to_string(),
            entry_type: "response_item".to_string(),
            payload: serde_json::json!({
                "type": "function_call_output",
                "call_id": "call_123",
                "output": "file1.txt\nfile2.txt"
            }),
        };

        let mut call_map = FunctionCallMap::default();
        let mut turn_model = None;

        // Process function call
        let call_events = call_entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &None,
            &mut turn_model,
        );
        assert_eq!(call_events.len(), 1);
        assert_eq!(call_events[0].event_type, EventType::PreToolUse);
        assert_eq!(call_events[0].tool_use_id, Some("call_123".to_string()));
        assert_eq!(call_events[0].tool_name, Some("shell_command".to_string()));

        // Process function output
        let output_events = output_entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &None,
            &mut turn_model,
        );
        assert_eq!(output_events.len(), 1);
        assert_eq!(output_events[0].event_type, EventType::PostToolUse);
        assert_eq!(output_events[0].tool_use_id, Some("call_123".to_string()));
        assert_eq!(
            output_events[0].tool_name,
            Some("shell_command".to_string())
        );

        // Check payload has command
        let payload: serde_json::Value =
            serde_json::from_str(output_events[0].payload.as_ref().unwrap()).unwrap();
        assert_eq!(payload["command"], "ls");

        // Check metadata has truncated output
        let metadata: serde_json::Value =
            serde_json::from_str(output_events[0].metadata.as_ref().unwrap()).unwrap();
        assert_eq!(metadata["output"], "file1.txt\nfile2.txt");
    }

    #[test]
    fn test_turn_context_updates_model() {
        let ctx_entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:07.612Z".to_string(),
            entry_type: "turn_context".to_string(),
            payload: serde_json::json!({
                "cwd": "/test",
                "model": "gpt-5.2-codex",
                "approval_policy": "on-request"
            }),
        };

        let mut call_map = FunctionCallMap::default();
        let mut turn_model = None;

        let events = ctx_entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &None,
            &mut turn_model,
        );

        assert!(events.is_empty()); // No events from turn_context
        assert_eq!(turn_model, Some("gpt-5.2-codex".to_string()));
    }

    #[test]
    fn test_skip_null_token_info() {
        let entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:08.167Z".to_string(),
            entry_type: "event_msg".to_string(),
            payload: serde_json::json!({
                "type": "token_count",
                "info": null,
                "rate_limits": {}
            }),
        };

        let mut call_map = FunctionCallMap::default();
        let mut turn_model = None;

        let events = entry.into_events(
            "machine-1",
            "session-1",
            &mut call_map,
            &None,
            &mut turn_model,
        );

        assert!(events.is_empty());
    }

    #[test]
    fn test_dedup_key() {
        let entry = CodexSessionEntry {
            timestamp: "2025-11-27T01:56:10.186Z".to_string(),
            entry_type: "event_msg".to_string(),
            payload: serde_json::json!({}),
        };

        assert_eq!(
            entry.dedup_key(),
            "2025-11-27T01:56:10.186Z:event_msg".to_string()
        );
    }
}