codex/codex-rs/core/src/memory_trace.rs

use std::path::Path;
use std::path::PathBuf;

use crate::ModelClient;
use codex_api::RawMemory as ApiRawMemory;
use codex_api::RawMemoryMetadata as ApiRawMemoryMetadata;
use codex_otel::SessionTelemetry;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use serde_json::Map;
use serde_json::Value;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BuiltMemory {
    pub memory_id: String,
    pub source_path: PathBuf,
    pub raw_memory: String,
    pub memory_summary: String,
}

struct PreparedTrace {
    memory_id: String,
    source_path: PathBuf,
    payload: ApiRawMemory,
}

/// Loads raw trace files, normalizes items, and builds memory summaries.
///
/// The request/response wiring mirrors the memory summarize E2E flow:
/// `/v1/memories/trace_summarize` with one output object per input raw memory.
///
/// The caller provides the model selection, reasoning effort, and telemetry context explicitly so
/// the session-scoped [`ModelClient`] can be reused across turns.
pub async fn build_memories_from_trace_files(
    client: &ModelClient,
    trace_paths: &[PathBuf],
    model_info: &ModelInfo,
    effort: Option<ReasoningEffortConfig>,
    session_telemetry: &SessionTelemetry,
) -> Result<Vec<BuiltMemory>> {
    if trace_paths.is_empty() {
        return Ok(Vec::new());
    }

    let mut prepared = Vec::with_capacity(trace_paths.len());
    for (index, path) in trace_paths.iter().enumerate() {
        prepared.push(prepare_trace(index + 1, path).await?);
    }

    let raw_memories = prepared.iter().map(|trace| trace.payload.clone()).collect();
    let output = client
        .summarize_memories(raw_memories, model_info, effort, session_telemetry)
        .await?;
    if output.len() != prepared.len() {
        return Err(CodexErr::InvalidRequest(format!(
            "unexpected memory summarize output length: expected {}, got {}",
            prepared.len(),
            output.len()
        )));
    }

    Ok(prepared
        .into_iter()
        .zip(output)
        .map(|(trace, summary)| BuiltMemory {
            memory_id: trace.memory_id,
            source_path: trace.source_path,
            raw_memory: summary.raw_memory,
            memory_summary: summary.memory_summary,
        })
        .collect())
}

async fn prepare_trace(index: usize, path: &Path) -> Result<PreparedTrace> {
    let text = load_trace_text(path).await?;
    let items = load_trace_items(path, &text)?;
    let memory_id = build_memory_id(index, path);
    let source_path = path.to_path_buf();

    Ok(PreparedTrace {
        memory_id: memory_id.clone(),
        source_path: source_path.clone(),
        payload: ApiRawMemory {
            id: memory_id,
            metadata: ApiRawMemoryMetadata {
                source_path: source_path.display().to_string(),
            },
            items,
        },
    })
}

async fn load_trace_text(path: &Path) -> Result<String> {
    let raw = tokio::fs::read(path).await?;
    Ok(decode_trace_bytes(&raw))
}

fn decode_trace_bytes(raw: &[u8]) -> String {
    if let Some(without_bom) = raw.strip_prefix(&[0xEF, 0xBB, 0xBF])
        && let Ok(text) = String::from_utf8(without_bom.to_vec())
    {
        return text;
    }
    if let Ok(text) = String::from_utf8(raw.to_vec()) {
        return text;
    }
    raw.iter().map(|b| char::from(*b)).collect()
}

fn load_trace_items(path: &Path, text: &str) -> Result<Vec<Value>> {
    if let Ok(Value::Array(items)) = serde_json::from_str::<Value>(text) {
        let dict_items = items
            .into_iter()
            .filter(serde_json::Value::is_object)
            .collect::<Vec<_>>();
        if dict_items.is_empty() {
            return Err(CodexErr::InvalidRequest(format!(
                "no object items found in trace file: {}",
                path.display()
            )));
        }
        return normalize_trace_items(dict_items, path);
    }

    let mut parsed_items = Vec::new();
    for line in text.lines() {
        let line = line.trim();
        if line.is_empty() || (!line.starts_with('{') && !line.starts_with('[')) {
            continue;
        }

        let Ok(obj) = serde_json::from_str::<Value>(line) else {
            continue;
        };

        match obj {
            Value::Object(_) => parsed_items.push(obj),
            Value::Array(inner) => {
                parsed_items.extend(inner.into_iter().filter(serde_json::Value::is_object))
            }
            _ => {}
        }
    }

    if parsed_items.is_empty() {
        return Err(CodexErr::InvalidRequest(format!(
            "no JSON items parsed from trace file: {}",
            path.display()
        )));
    }

    normalize_trace_items(parsed_items, path)
}

fn normalize_trace_items(items: Vec<Value>, path: &Path) -> Result<Vec<Value>> {
    let mut normalized = Vec::new();

    for item in items {
        let Value::Object(obj) = item else {
            continue;
        };

        if let Some(payload) = obj.get("payload") {
            if obj.get("type").and_then(Value::as_str) != Some("response_item") {
                continue;
            }

            match payload {
                Value::Object(payload_item) => {
                    if is_allowed_trace_item(payload_item) {
                        normalized.push(Value::Object(payload_item.clone()));
                    }
                }
                Value::Array(payload_items) => {
                    for payload_item in payload_items {
                        if let Value::Object(payload_item) = payload_item
                            && is_allowed_trace_item(payload_item)
                        {
                            normalized.push(Value::Object(payload_item.clone()));
                        }
                    }
                }
                _ => {}
            }
            continue;
        }

        if is_allowed_trace_item(&obj) {
            normalized.push(Value::Object(obj));
        }
    }

    if normalized.is_empty() {
        return Err(CodexErr::InvalidRequest(format!(
            "no valid trace items after normalization: {}",
            path.display()
        )));
    }
    Ok(normalized)
}

fn is_allowed_trace_item(item: &Map<String, Value>) -> bool {
    let Some(item_type) = item.get("type").and_then(Value::as_str) else {
        return false;
    };

    if item_type == "message" {
        return matches!(
            item.get("role").and_then(Value::as_str),
            Some("assistant" | "system" | "developer" | "user")
        );
    }

    true
}

fn build_memory_id(index: usize, path: &Path) -> String {
    let stem = path
        .file_stem()
        .map(|stem| stem.to_string_lossy().into_owned())
        .filter(|stem| !stem.is_empty())
        .unwrap_or_else(|| "memory".to_string());
    format!("memory_{index}_{stem}")
}

#[cfg(test)]
#[path = "memory_trace_tests.rs"]
mod tests;