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, session_telemetry: &SessionTelemetry, ) -> Result> { 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 { 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 { 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> { if let Ok(Value::Array(items)) = serde_json::from_str::(text) { let dict_items = items .into_iter() .filter(serde_json::Value::is_object) .collect::>(); 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::(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, path: &Path) -> Result> { 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) -> 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;