openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-16 12:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

split memories out of codex-core

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Ahmed Ibrahim
2026-03-18 04:56:44 +00:00
parent 1cf68f940c
commit f5dc4dc801
20 changed files with 902 additions and 650 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
codex-rs/Cargo.lock generated
View File

@@ -1849,6 +1849,7 @@ dependencies = [
"codex-git",
"codex-hooks",
"codex-keyring-store",
"codex-memories",
"codex-network-proxy",
"codex-otel",
"codex-protocol",
@@ -2194,6 +2195,23 @@ dependencies = [
"wiremock",
]
[[package]]
name = "codex-memories"
version = "0.0.0"
dependencies = [
"anyhow",
"chrono",
"codex-api",
"codex-protocol",
"codex-state",
"pretty_assertions",
"serde_json",
"tempfile",
"tokio",
"tracing",
"uuid",
]
[[package]]
name = "codex-network-proxy"
version = "0.0.0"

2
codex-rs/Cargo.toml
View File

@@ -32,6 +32,7 @@ members = [
"linux-sandbox",
"lmstudio",
"login",
"memories",
"mcp-server",
"network-proxy",
"ollama",
@@ -114,6 +115,7 @@ codex-keyring-store = { path = "keyring-store" }
codex-linux-sandbox = { path = "linux-sandbox" }
codex-lmstudio = { path = "lmstudio" }
codex-login = { path = "login" }
codex-memories = { path = "memories" }
codex-mcp-server = { path = "mcp-server" }
codex-network-proxy = { path = "network-proxy" }
codex-ollama = { path = "ollama" }

1
codex-rs/core/Cargo.toml
View File

@@ -41,6 +41,7 @@ codex-file-search = { workspace = true }
codex-git = { workspace = true }
codex-hooks = { workspace = true }
codex-keyring-store = { workspace = true }
codex-memories = { workspace = true }
codex-network-proxy = { workspace = true }
codex-otel = { workspace = true }
codex-artifacts = { workspace = true }

6
codex-rs/core/src/codex.rs
View File

@@ -4798,8 +4798,10 @@ mod handlers {
errors.push("state db unavailable; memory rows were not cleared".to_string());
}
let memory_root = crate::memories::memory_root(&config.codex_home);
if let Err(err) = crate::memories::clear_memory_root_contents(&memory_root).await {
let memory_root = codex_memories::memories::memory_root(&config.codex_home);
if let Err(err) =
codex_memories::memories::control::clear_memory_root_contents(&memory_root).await
{
errors.push(format!(
"failed clearing memory directory {}: {err}",
memory_root.display()

2
codex-rs/core/src/config/mod.rs
View File

@@ -41,7 +41,6 @@ use crate::features::FeatureOverrides;
use crate::features::Features;
use crate::features::FeaturesToml;
use crate::git_info::resolve_root_git_project_for_trust;
use crate::memories::memory_root;
use crate::model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID;
use crate::model_provider_info::LMSTUDIO_OSS_PROVIDER_ID;
use crate::model_provider_info::ModelProviderInfo;
@@ -62,6 +61,7 @@ use crate::windows_sandbox::resolve_windows_sandbox_mode;
use crate::windows_sandbox::resolve_windows_sandbox_private_desktop;
use codex_app_server_protocol::Tools;
use codex_app_server_protocol::UserSavedConfig;
use codex_memories::memories::memory_root;
use codex_protocol::config_types::AltScreenMode;
use codex_protocol::config_types::ForcedLoginMethod;
use codex_protocol::config_types::Personality;

26
codex-rs/core/src/memories/citations_tests.rs
View File

@@ -1,26 +0,0 @@
use super::get_thread_id_from_citations;
use codex_protocol::ThreadId;
use pretty_assertions::assert_eq;
#[test]
fn get_thread_id_from_citations_extracts_thread_ids() {
let first = ThreadId::new();
let second = ThreadId::new();
let citations = vec![format!(
"<memory_citation>\n<citation_entries>\nMEMORY.md:1-2|note=[x]\n</citation_entries>\n<thread_ids>\n{first}\nnot-a-uuid\n{second}\n</thread_ids>\n</memory_citation>"
)];
assert_eq!(get_thread_id_from_citations(citations), vec![first, second]);
}
#[test]
fn get_thread_id_from_citations_supports_legacy_rollout_ids() {
let thread_id = ThreadId::new();
let citations = vec![format!(
"<memory_citation>\n<rollout_ids>\n{thread_id}\n</rollout_ids>\n</memory_citation>"
)];
assert_eq!(get_thread_id_from_citations(citations), vec![thread_id]);
}

30
codex-rs/core/src/memories/mod.rs
View File

@@ -4,29 +4,29 @@
//! - Phase 1: select rollouts, extract stage-1 raw memories, persist stage-1 outputs, and enqueue consolidation.
//! - Phase 2: claim a global consolidation lock, materialize consolidation inputs, and dispatch one consolidation agent.
pub(crate) mod citations;
mod control;
mod phase1;
mod phase2;
pub(crate) mod prompts;
mod start;
mod storage;
#[cfg(test)]
mod tests;
pub(crate) mod usage;
use codex_protocol::openai_models::ReasoningEffort;
pub(crate) use control::clear_memory_root_contents;
pub(crate) use codex_memories::memories::memory_root;
/// Starts the memory startup pipeline for eligible root sessions.
/// This is the single entrypoint that `codex` uses to trigger memory startup.
///
/// This is the entry point to read and understand this module.
pub(crate) use start::start_memories_startup_task;
mod artifacts {
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
pub(super) const RAW_MEMORIES_FILENAME: &str = "raw_memories.md";
pub(crate) mod citations {
pub(crate) use codex_memories::memories::citations::*;
}
pub(crate) mod storage {
pub(crate) use codex_memories::memories::storage::*;
}
/// Phase 1 (startup extraction).
@@ -98,18 +98,18 @@ mod metrics {
use std::path::Path;
use std::path::PathBuf;
pub fn memory_root(codex_home: &Path) -> PathBuf {
codex_home.join("memories")
pub(crate) async fn clear_memory_root_contents(memory_root: &Path) -> std::io::Result<()> {
codex_memories::memories::control::clear_memory_root_contents(memory_root).await
}
fn rollout_summaries_dir(root: &Path) -> PathBuf {
root.join(artifacts::ROLLOUT_SUMMARIES_SUBDIR)
pub(crate) fn rollout_summaries_dir(root: &Path) -> PathBuf {
codex_memories::memories::rollout_summaries_dir(root)
}
fn raw_memories_file(root: &Path) -> PathBuf {
root.join(artifacts::RAW_MEMORIES_FILENAME)
pub(crate) fn raw_memories_file(root: &Path) -> PathBuf {
codex_memories::memories::raw_memories_file(root)
}
async fn ensure_layout(root: &Path) -> std::io::Result<()> {
tokio::fs::create_dir_all(rollout_summaries_dir(root)).await
pub(crate) async fn ensure_layout(root: &Path) -> std::io::Result<()> {
codex_memories::memories::ensure_layout(root).await
}

260
codex-rs/core/src/memories/storage.rs
View File

@@ -1,260 +0,0 @@
use codex_state::Stage1Output;
use std::collections::HashSet;
use std::fmt::Write as _;
use std::path::Path;
use tracing::warn;
use uuid::Uuid;
use crate::memories::ensure_layout;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
/// Rebuild `raw_memories.md` from DB-backed stage-1 outputs.
pub(super) async fn rebuild_raw_memories_file_from_memories(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
ensure_layout(root).await?;
rebuild_raw_memories_file(root, memories, max_raw_memories_for_consolidation).await
}
/// Syncs canonical rollout summary files from DB-backed stage-1 output rows.
pub(super) async fn sync_rollout_summaries_from_memories(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
ensure_layout(root).await?;
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
let keep = retained
.iter()
.map(rollout_summary_file_stem)
.collect::<HashSet<_>>();
prune_rollout_summaries(root, &keep).await?;
for memory in retained {
write_rollout_summary_for_thread(root, memory).await?;
}
if retained.is_empty() {
for file_name in ["MEMORY.md", "memory_summary.md"] {
let path = root.join(file_name);
if let Err(err) = tokio::fs::remove_file(path).await
&& err.kind() != std::io::ErrorKind::NotFound
{
return Err(err);
}
}
let skills_dir = root.join("skills");
if let Err(err) = tokio::fs::remove_dir_all(skills_dir).await
&& err.kind() != std::io::ErrorKind::NotFound
{
return Err(err);
}
}
Ok(())
}
async fn rebuild_raw_memories_file(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
let mut body = String::from("# Raw Memories\n\n");
if retained.is_empty() {
body.push_str("No raw memories yet.\n");
return tokio::fs::write(raw_memories_file(root), body).await;
}
body.push_str("Merged stage-1 raw memories (latest first):\n\n");
for memory in retained {
writeln!(body, "## Thread `{}`", memory.thread_id).map_err(raw_memories_format_error)?;
writeln!(
body,
"updated_at: {}",
memory.source_updated_at.to_rfc3339()
)
.map_err(raw_memories_format_error)?;
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(raw_memories_format_error)?;
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
.map_err(raw_memories_format_error)?;
let rollout_summary_file = format!("{}.md", rollout_summary_file_stem(memory));
writeln!(body, "rollout_summary_file: {rollout_summary_file}")
.map_err(raw_memories_format_error)?;
writeln!(body).map_err(raw_memories_format_error)?;
body.push_str(memory.raw_memory.trim());
body.push_str("\n\n");
}
tokio::fs::write(raw_memories_file(root), body).await
}
async fn prune_rollout_summaries(root: &Path, keep: &HashSet<String>) -> std::io::Result<()> {
let dir_path = rollout_summaries_dir(root);
let mut dir = match tokio::fs::read_dir(&dir_path).await {
Ok(dir) => dir,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(err) => return Err(err),
};
while let Some(entry) = dir.next_entry().await? {
let path = entry.path();
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
let Some(stem) = file_name.strip_suffix(".md") else {
continue;
};
if !keep.contains(stem)
&& let Err(err) = tokio::fs::remove_file(&path).await
&& err.kind() != std::io::ErrorKind::NotFound
{
warn!(
"failed pruning outdated rollout summary {}: {err}",
path.display()
);
}
}
Ok(())
}
async fn write_rollout_summary_for_thread(
root: &Path,
memory: &Stage1Output,
) -> std::io::Result<()> {
let file_stem = rollout_summary_file_stem(memory);
let path = rollout_summaries_dir(root).join(format!("{file_stem}.md"));
let mut body = String::new();
writeln!(body, "thread_id: {}", memory.thread_id).map_err(rollout_summary_format_error)?;
writeln!(
body,
"updated_at: {}",
memory.source_updated_at.to_rfc3339()
)
.map_err(rollout_summary_format_error)?;
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
.map_err(rollout_summary_format_error)?;
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(rollout_summary_format_error)?;
if let Some(git_branch) = memory.git_branch.as_deref() {
writeln!(body, "git_branch: {git_branch}").map_err(rollout_summary_format_error)?;
}
writeln!(body).map_err(rollout_summary_format_error)?;
body.push_str(&memory.rollout_summary);
body.push('\n');
tokio::fs::write(path, body).await
}
fn retained_memories(
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> &[Stage1Output] {
&memories[..memories.len().min(max_raw_memories_for_consolidation)]
}
fn raw_memories_format_error(err: std::fmt::Error) -> std::io::Error {
std::io::Error::other(format!("format raw memories: {err}"))
}
fn rollout_summary_format_error(err: std::fmt::Error) -> std::io::Error {
std::io::Error::other(format!("format rollout summary: {err}"))
}
pub(crate) fn rollout_summary_file_stem(memory: &Stage1Output) -> String {
rollout_summary_file_stem_from_parts(
memory.thread_id,
memory.source_updated_at,
memory.rollout_slug.as_deref(),
)
}
pub(super) fn rollout_summary_file_stem_from_parts(
thread_id: codex_protocol::ThreadId,
source_updated_at: chrono::DateTime<chrono::Utc>,
rollout_slug: Option<&str>,
) -> String {
const ROLLOUT_SLUG_MAX_LEN: usize = 60;
const SHORT_HASH_ALPHABET: &[u8; 62] =
b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
const SHORT_HASH_SPACE: u32 = 14_776_336;
let thread_id = thread_id.to_string();
let (timestamp_fragment, short_hash_seed) = match Uuid::parse_str(&thread_id) {
Ok(thread_uuid) => {
let timestamp = thread_uuid
.get_timestamp()
.and_then(|uuid_timestamp| {
let (seconds, nanos) = uuid_timestamp.to_unix();
i64::try_from(seconds).ok().and_then(|secs| {
chrono::DateTime::<chrono::Utc>::from_timestamp(secs, nanos)
})
})
.unwrap_or(source_updated_at);
let short_hash_seed = (thread_uuid.as_u128() & 0xFFFF_FFFF) as u32;
(
timestamp.format("%Y-%m-%dT%H-%M-%S").to_string(),
short_hash_seed,
)
}
Err(_) => {
let mut short_hash_seed = 0u32;
for byte in thread_id.bytes() {
short_hash_seed = short_hash_seed
.wrapping_mul(31)
.wrapping_add(u32::from(byte));
}
(
source_updated_at.format("%Y-%m-%dT%H-%M-%S").to_string(),
short_hash_seed,
)
}
};
let mut short_hash_value = short_hash_seed % SHORT_HASH_SPACE;
let mut short_hash_chars = ['0'; 4];
for idx in (0..short_hash_chars.len()).rev() {
let alphabet_idx = (short_hash_value % SHORT_HASH_ALPHABET.len() as u32) as usize;
short_hash_chars[idx] = SHORT_HASH_ALPHABET[alphabet_idx] as char;
short_hash_value /= SHORT_HASH_ALPHABET.len() as u32;
}
let short_hash: String = short_hash_chars.iter().collect();
let file_prefix = format!("{timestamp_fragment}-{short_hash}");
let Some(raw_slug) = rollout_slug else {
return file_prefix;
};
let mut slug = String::with_capacity(ROLLOUT_SLUG_MAX_LEN);
for ch in raw_slug.chars() {
if slug.len() >= ROLLOUT_SLUG_MAX_LEN {
break;
}
if ch.is_ascii_alphanumeric() {
slug.push(ch.to_ascii_lowercase());
} else {
slug.push('_');
}
}
while slug.ends_with('_') {
slug.pop();
}
if slug.is_empty() {
file_prefix
} else {
format!("{file_prefix}-{slug}")
}
}
#[cfg(test)]
#[path = "storage_tests.rs"]
mod tests;

70
codex-rs/core/src/memories/storage_tests.rs
View File

@@ -1,70 +0,0 @@
use super::rollout_summary_file_stem;
use super::rollout_summary_file_stem_from_parts;
use chrono::TimeZone;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_state::Stage1Output;
use pretty_assertions::assert_eq;
use std::path::PathBuf;
const FIXED_PREFIX: &str = "2025-02-11T15-35-19-jqmb";
fn stage1_output_with_slug(thread_id: ThreadId, rollout_slug: Option<&str>) -> Stage1Output {
Stage1Output {
thread_id,
source_updated_at: Utc.timestamp_opt(123, 0).single().expect("timestamp"),
raw_memory: "raw memory".to_string(),
rollout_summary: "summary".to_string(),
rollout_slug: rollout_slug.map(ToString::to_string),
rollout_path: PathBuf::from("/tmp/rollout.jsonl"),
cwd: PathBuf::from("/tmp/workspace"),
git_branch: None,
generated_at: Utc.timestamp_opt(124, 0).single().expect("timestamp"),
}
}
fn fixed_thread_id() -> ThreadId {
ThreadId::try_from("0194f5a6-89ab-7cde-8123-456789abcdef").expect("valid thread id")
}
#[test]
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_missing() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(thread_id, None);
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
assert_eq!(
rollout_summary_file_stem_from_parts(
memory.thread_id,
memory.source_updated_at,
memory.rollout_slug.as_deref(),
),
FIXED_PREFIX
);
}
#[test]
fn rollout_summary_file_stem_sanitizes_and_truncates_slug() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(
thread_id,
Some("Unsafe Slug/With Spaces & Symbols + EXTRA_LONG_12345_67890_ABCDE_fghij_klmno"),
);
let stem = rollout_summary_file_stem(&memory);
let slug = stem
.strip_prefix(&format!("{FIXED_PREFIX}-"))
.expect("slug suffix should be present");
assert_eq!(slug.len(), 60);
assert_eq!(
slug,
"unsafe_slug_with_spaces___symbols___extra_long_12345_67890_a"
);
}
#[test]
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_is_empty() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(thread_id, Some(""));
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
}

205
codex-rs/core/src/memory_trace.rs
View File

@@ -1,30 +1,11 @@
use std::path::Path;
use std::path::PathBuf;
use crate::ModelClient;
use crate::error::CodexErr;
use crate::error::Result;
use codex_api::RawMemory as ApiRawMemory;
use codex_api::RawMemoryMetadata as ApiRawMemoryMetadata;
pub use codex_memories::memory_trace::BuiltMemory;
use codex_otel::SessionTelemetry;
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.
///
@@ -44,187 +25,13 @@ pub async fn build_memories_from_trace_files(
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 prepared = codex_memories::memory_trace::load_trace_requests(trace_paths)
.await
.map_err(|err| crate::error::CodexErr::InvalidRequest(err.to_string()))?;
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())
codex_memories::memory_trace::build_memories_from_output(prepared, output)
.map_err(|err| crate::error::CodexErr::InvalidRequest(err.to_string()))
}
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;

73
codex-rs/core/src/memory_trace_tests.rs
View File

@@ -1,73 +0,0 @@
use super::*;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
#[test]
fn normalize_trace_items_handles_payload_wrapper_and_message_role_filtering() {
let items = vec![
serde_json::json!({
"type": "response_item",
"payload": {"type": "message", "role": "assistant", "content": []}
}),
serde_json::json!({
"type": "response_item",
"payload": [
{"type": "message", "role": "user", "content": []},
{"type": "message", "role": "tool", "content": []},
{"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"}
]
}),
serde_json::json!({
"type": "not_response_item",
"payload": {"type": "message", "role": "assistant", "content": []}
}),
serde_json::json!({
"type": "message",
"role": "developer",
"content": []
}),
];
let normalized = normalize_trace_items(items, Path::new("trace.json")).expect("normalize");
let expected = vec![
serde_json::json!({"type": "message", "role": "assistant", "content": []}),
serde_json::json!({"type": "message", "role": "user", "content": []}),
serde_json::json!({"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"}),
serde_json::json!({"type": "message", "role": "developer", "content": []}),
];
assert_eq!(normalized, expected);
}
#[test]
fn load_trace_items_supports_jsonl_arrays_and_objects() {
let text = r#"
{"type":"response_item","payload":{"type":"message","role":"assistant","content":[]}}
[{"type":"message","role":"user","content":[]},{"type":"message","role":"tool","content":[]}]
"#;
let loaded = load_trace_items(Path::new("trace.jsonl"), text).expect("load");
let expected = vec![
serde_json::json!({"type":"message","role":"assistant","content":[]}),
serde_json::json!({"type":"message","role":"user","content":[]}),
];
assert_eq!(loaded, expected);
}
#[tokio::test]
async fn load_trace_text_decodes_utf8_sig() {
let dir = tempdir().expect("tempdir");
let path = dir.path().join("trace.json");
tokio::fs::write(
&path,
[
0xEF, 0xBB, 0xBF, b'[', b'{', b'"', b't', b'y', b'p', b'e', b'"', b':', b'"', b'm',
b'e', b's', b's', b'a', b'g', b'e', b'"', b',', b'"', b'r', b'o', b'l', b'e', b'"',
b':', b'"', b'u', b's', b'e', b'r', b'"', b',', b'"', b'c', b'o', b'n', b't', b'e',
b'n', b't', b'"', b':', b'[', b']', b'}', b']',
],
)
.await
.expect("write");
let text = load_trace_text(&path).await.expect("decode");
assert!(text.starts_with('['));
}

2
codex-rs/core/src/stream_events_utils.rs
View File

@@ -4,6 +4,7 @@ use std::sync::Arc;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_memories::memories::citations::get_thread_id_from_citations;
use codex_protocol::config_types::ModeKind;
use codex_protocol::items::TurnItem;
use codex_utils_stream_parser::strip_citations;
@@ -14,7 +15,6 @@ use crate::codex::TurnContext;
use crate::error::CodexErr;
use crate::error::Result;
use crate::function_tool::FunctionCallError;
use crate::memories::citations::get_thread_id_from_citations;
use crate::parse_turn_item;
use crate::state_db;
use crate::tools::parallel::ToolCallRuntime;

6
codex-rs/memories/BUILD.bazel Normal file
View File

@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "memories",
crate_name = "codex_memories",
)

27
codex-rs/memories/Cargo.toml Normal file
View File

@@ -0,0 +1,27 @@
[package]
name = "codex-memories"
version.workspace = true
edition.workspace = true
license.workspace = true
[lib]
name = "codex_memories"
path = "src/lib.rs"
[lints]
workspace = true
[dependencies]
anyhow = { workspace = true }
chrono = { workspace = true }
codex-api = { workspace = true }
codex-protocol = { workspace = true }
codex-state = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["fs", "macros", "rt"] }
tracing = { workspace = true }
uuid = { workspace = true, features = ["serde", "v4", "v5"] }
[dev-dependencies]
pretty_assertions = { workspace = true }
tempfile = { workspace = true }

2
codex-rs/memories/src/lib.rs Normal file
View File

@@ -0,0 +1,2 @@
pub mod memories;
pub mod memory_trace;

30
codex-rs/core/src/memories/citations.rs → codex-rs/memories/src/memories/citations.rs
View File

@@ -32,5 +32,31 @@ pub fn get_thread_id_from_citations(citations: Vec<String>) -> Vec<ThreadId> {
}
#[cfg(test)]
#[path = "citations_tests.rs"]
mod tests;
mod tests {
use super::get_thread_id_from_citations;
use codex_protocol::ThreadId;
use pretty_assertions::assert_eq;
#[test]
fn get_thread_id_from_citations_extracts_thread_ids() {
let first = ThreadId::new();
let second = ThreadId::new();
let citations = vec![format!(
"<memory_citation>\n<citation_entries>\nMEMORY.md:1-2|note=[x]\n</citation_entries>\n<thread_ids>\n{first}\nnot-a-uuid\n{second}\n</thread_ids>\n</memory_citation>"
)];
assert_eq!(get_thread_id_from_citations(citations), vec![first, second]);
}
#[test]
fn get_thread_id_from_citations_supports_legacy_rollout_ids() {
let thread_id = ThreadId::new();
let citations = vec![format!(
"<memory_citation>\n<rollout_ids>\n{thread_id}\n</rollout_ids>\n</memory_citation>"
)];
assert_eq!(get_thread_id_from_citations(citations), vec![thread_id]);
}
}

2
codex-rs/core/src/memories/control.rs → codex-rs/memories/src/memories/control.rs
View File

@@ -1,6 +1,6 @@
use std::path::Path;
pub(crate) async fn clear_memory_root_contents(memory_root: &Path) -> std::io::Result<()> {
pub async fn clear_memory_root_contents(memory_root: &Path) -> std::io::Result<()> {
match tokio::fs::symlink_metadata(memory_root).await {
Ok(metadata) if metadata.file_type().is_symlink() => {
return Err(std::io::Error::new(

29
codex-rs/memories/src/memories/mod.rs Normal file
View File

@@ -0,0 +1,29 @@
//! Shared memory filesystem and artifact utilities.
pub mod citations;
pub mod control;
pub mod storage;
use std::path::Path;
use std::path::PathBuf;
mod artifacts {
pub(super) const ROLLOUT_SUMMARIES_SUBDIR: &str = "rollout_summaries";
pub(super) const RAW_MEMORIES_FILENAME: &str = "raw_memories.md";
}
pub fn memory_root(codex_home: &Path) -> PathBuf {
codex_home.join("memories")
}
pub fn rollout_summaries_dir(root: &Path) -> PathBuf {
root.join(artifacts::ROLLOUT_SUMMARIES_SUBDIR)
}
pub fn raw_memories_file(root: &Path) -> PathBuf {
root.join(artifacts::RAW_MEMORIES_FILENAME)
}
pub async fn ensure_layout(root: &Path) -> std::io::Result<()> {
tokio::fs::create_dir_all(rollout_summaries_dir(root)).await
}

480
codex-rs/memories/src/memories/storage.rs Normal file
View File

@@ -0,0 +1,480 @@
use crate::memories::ensure_layout;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
use codex_state::Stage1Output;
use std::collections::HashSet;
use std::fmt::Write as _;
use std::path::Path;
use tracing::warn;
use uuid::Uuid;
/// Rebuild `raw_memories.md` from DB-backed stage-1 outputs.
pub async fn rebuild_raw_memories_file_from_memories(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
ensure_layout(root).await?;
rebuild_raw_memories_file(root, memories, max_raw_memories_for_consolidation).await
}
/// Syncs canonical rollout summary files from DB-backed stage-1 output rows.
pub async fn sync_rollout_summaries_from_memories(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
ensure_layout(root).await?;
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
let keep = retained
.iter()
.map(rollout_summary_file_stem)
.collect::<HashSet<_>>();
prune_rollout_summaries(root, &keep).await?;
for memory in retained {
write_rollout_summary_for_thread(root, memory).await?;
}
if retained.is_empty() {
for file_name in ["MEMORY.md", "memory_summary.md"] {
let path = root.join(file_name);
if let Err(err) = tokio::fs::remove_file(path).await
&& err.kind() != std::io::ErrorKind::NotFound
{
return Err(err);
}
}
let skills_dir = root.join("skills");
if let Err(err) = tokio::fs::remove_dir_all(skills_dir).await
&& err.kind() != std::io::ErrorKind::NotFound
{
return Err(err);
}
}
Ok(())
}
async fn rebuild_raw_memories_file(
root: &Path,
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> std::io::Result<()> {
let retained = retained_memories(memories, max_raw_memories_for_consolidation);
let mut body = String::from("# Raw Memories\n\n");
if retained.is_empty() {
body.push_str("No raw memories yet.\n");
return tokio::fs::write(raw_memories_file(root), body).await;
}
body.push_str("Merged stage-1 raw memories (latest first):\n\n");
for memory in retained {
writeln!(body, "## Thread `{}`", memory.thread_id).map_err(raw_memories_format_error)?;
writeln!(
body,
"updated_at: {}",
memory.source_updated_at.to_rfc3339()
)
.map_err(raw_memories_format_error)?;
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(raw_memories_format_error)?;
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
.map_err(raw_memories_format_error)?;
let rollout_summary_file = format!("{}.md", rollout_summary_file_stem(memory));
writeln!(body, "rollout_summary_file: {rollout_summary_file}")
.map_err(raw_memories_format_error)?;
writeln!(body).map_err(raw_memories_format_error)?;
body.push_str(memory.raw_memory.trim());
body.push_str("\n\n");
}
tokio::fs::write(raw_memories_file(root), body).await
}
async fn prune_rollout_summaries(root: &Path, keep: &HashSet<String>) -> std::io::Result<()> {
let dir_path = rollout_summaries_dir(root);
let mut dir = match tokio::fs::read_dir(&dir_path).await {
Ok(dir) => dir,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(err) => return Err(err),
};
while let Some(entry) = dir.next_entry().await? {
let path = entry.path();
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
let Some(stem) = file_name.strip_suffix(".md") else {
continue;
};
if !keep.contains(stem)
&& let Err(err) = tokio::fs::remove_file(&path).await
&& err.kind() != std::io::ErrorKind::NotFound
{
warn!(
"failed pruning outdated rollout summary {}: {err}",
path.display()
);
}
}
Ok(())
}
async fn write_rollout_summary_for_thread(
root: &Path,
memory: &Stage1Output,
) -> std::io::Result<()> {
let file_stem = rollout_summary_file_stem(memory);
let path = rollout_summaries_dir(root).join(format!("{file_stem}.md"));
let mut body = String::new();
writeln!(body, "thread_id: {}", memory.thread_id).map_err(rollout_summary_format_error)?;
writeln!(
body,
"updated_at: {}",
memory.source_updated_at.to_rfc3339()
)
.map_err(rollout_summary_format_error)?;
writeln!(body, "rollout_path: {}", memory.rollout_path.display())
.map_err(rollout_summary_format_error)?;
writeln!(body, "cwd: {}", memory.cwd.display()).map_err(rollout_summary_format_error)?;
if let Some(git_branch) = memory.git_branch.as_deref() {
writeln!(body, "git_branch: {git_branch}").map_err(rollout_summary_format_error)?;
}
writeln!(body).map_err(rollout_summary_format_error)?;
body.push_str(&memory.rollout_summary);
body.push('\n');
tokio::fs::write(path, body).await
}
fn retained_memories(
memories: &[Stage1Output],
max_raw_memories_for_consolidation: usize,
) -> &[Stage1Output] {
&memories[..memories.len().min(max_raw_memories_for_consolidation)]
}
fn raw_memories_format_error(err: std::fmt::Error) -> std::io::Error {
std::io::Error::other(format!("format raw memories: {err}"))
}
fn rollout_summary_format_error(err: std::fmt::Error) -> std::io::Error {
std::io::Error::other(format!("format rollout summary: {err}"))
}
pub fn rollout_summary_file_stem(memory: &Stage1Output) -> String {
rollout_summary_file_stem_from_parts(
memory.thread_id,
memory.source_updated_at,
memory.rollout_slug.as_deref(),
)
}
pub fn rollout_summary_file_stem_from_parts(
thread_id: codex_protocol::ThreadId,
source_updated_at: chrono::DateTime<chrono::Utc>,
rollout_slug: Option<&str>,
) -> String {
const ROLLOUT_SLUG_MAX_LEN: usize = 60;
const SHORT_HASH_ALPHABET: &[u8; 62] =
b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
const SHORT_HASH_SPACE: u32 = 14_776_336;
let thread_id = thread_id.to_string();
let (timestamp_fragment, short_hash_seed) = match Uuid::parse_str(&thread_id) {
Ok(thread_uuid) => {
let timestamp = thread_uuid
.get_timestamp()
.and_then(|uuid_timestamp| {
let (seconds, nanos) = uuid_timestamp.to_unix();
i64::try_from(seconds).ok().and_then(|secs| {
chrono::DateTime::<chrono::Utc>::from_timestamp(secs, nanos)
})
})
.unwrap_or(source_updated_at);
let short_hash_seed = (thread_uuid.as_u128() & 0xFFFF_FFFF) as u32;
(
timestamp.format("%Y-%m-%dT%H-%M-%S").to_string(),
short_hash_seed,
)
}
Err(_) => {
let mut short_hash_seed = 0u32;
for byte in thread_id.bytes() {
short_hash_seed = short_hash_seed
.wrapping_mul(31)
.wrapping_add(u32::from(byte));
}
(
source_updated_at.format("%Y-%m-%dT%H-%M-%S").to_string(),
short_hash_seed,
)
}
};
let mut short_hash_value = short_hash_seed % SHORT_HASH_SPACE;
let mut short_hash = ['0'; 4];
for digit in short_hash.iter_mut().rev() {
let idx = usize::try_from(short_hash_value % 62).unwrap_or(0);
*digit = char::from(SHORT_HASH_ALPHABET[idx]);
short_hash_value /= 62;
}
let mut stem = format!(
"{timestamp_fragment}-{}",
short_hash.iter().collect::<String>()
);
let sanitized_slug = rollout_slug
.map(str::trim)
.filter(|slug| !slug.is_empty())
.map(|slug| {
slug.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() {
ch.to_ascii_lowercase()
} else {
'_'
}
})
.collect::<String>()
})
.map(|slug| slug.chars().take(ROLLOUT_SLUG_MAX_LEN).collect::<String>())
.filter(|slug| !slug.is_empty());
if let Some(slug) = sanitized_slug {
stem.push('-');
stem.push_str(&slug);
}
stem
}
#[cfg(test)]
mod tests {
use super::rebuild_raw_memories_file_from_memories;
use super::rollout_summary_file_stem;
use super::rollout_summary_file_stem_from_parts;
use super::sync_rollout_summaries_from_memories;
use crate::memories::control::clear_memory_root_contents;
use crate::memories::ensure_layout;
use crate::memories::memory_root;
use crate::memories::raw_memories_file;
use crate::memories::rollout_summaries_dir;
use chrono::TimeZone;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_state::Stage1Output;
use pretty_assertions::assert_eq;
use std::path::PathBuf;
use tempfile::tempdir;
const FIXED_PREFIX: &str = "2025-02-11T15-35-19-jqmb";
fn stage1_output_with_slug(thread_id: ThreadId, rollout_slug: Option<&str>) -> Stage1Output {
Stage1Output {
thread_id,
source_updated_at: Utc.timestamp_opt(123, 0).single().expect("timestamp"),
raw_memory: "raw memory".to_string(),
rollout_summary: "summary".to_string(),
rollout_slug: rollout_slug.map(ToString::to_string),
rollout_path: PathBuf::from("/tmp/rollout.jsonl"),
cwd: PathBuf::from("/tmp/workspace"),
git_branch: None,
generated_at: Utc.timestamp_opt(124, 0).single().expect("timestamp"),
}
}
fn fixed_thread_id() -> ThreadId {
ThreadId::try_from("0194f5a6-89ab-7cde-8123-456789abcdef").expect("valid thread id")
}
#[test]
fn memory_root_uses_shared_global_path() {
let dir = tempdir().expect("tempdir");
let codex_home = dir.path().join("codex");
assert_eq!(memory_root(&codex_home), codex_home.join("memories"));
}
#[tokio::test]
async fn clear_memory_root_contents_preserves_root_directory() {
let dir = tempdir().expect("tempdir");
let root = dir.path().join("memory");
let nested_dir = root.join("rollout_summaries");
tokio::fs::create_dir_all(&nested_dir)
.await
.expect("create rollout summaries dir");
tokio::fs::write(root.join("MEMORY.md"), "stale memory index\n")
.await
.expect("write memory index");
tokio::fs::write(nested_dir.join("rollout.md"), "stale rollout\n")
.await
.expect("write rollout summary");
clear_memory_root_contents(&root)
.await
.expect("clear memory root contents");
assert!(
tokio::fs::try_exists(&root)
.await
.expect("check memory root existence"),
"memory root should still exist after clearing contents"
);
let mut entries = tokio::fs::read_dir(&root)
.await
.expect("read memory root after clear");
assert!(
entries
.next_entry()
.await
.expect("read next entry")
.is_none(),
"memory root should be empty after clearing contents"
);
}
#[cfg(unix)]
#[tokio::test]
async fn clear_memory_root_contents_rejects_symlinked_root() {
let dir = tempdir().expect("tempdir");
let target = dir.path().join("outside");
tokio::fs::create_dir_all(&target)
.await
.expect("create symlink target dir");
let target_file = target.join("keep.txt");
tokio::fs::write(&target_file, "keep\n")
.await
.expect("write target file");
let root = dir.path().join("memory");
std::os::unix::fs::symlink(&target, &root).expect("create memory root symlink");
let err = clear_memory_root_contents(&root)
.await
.expect_err("symlinked memory root should be rejected");
assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
assert!(
tokio::fs::try_exists(&target_file)
.await
.expect("check target file existence"),
"rejecting a symlinked memory root should not delete the symlink target"
);
}
#[tokio::test]
async fn sync_rollout_summaries_and_raw_memories_file_keeps_latest_memories_only() {
let dir = tempdir().expect("tempdir");
let root = dir.path().join("memory");
ensure_layout(&root).await.expect("ensure layout");
let keep_id = ThreadId::default().to_string();
let drop_id = ThreadId::default().to_string();
let keep_path = rollout_summaries_dir(&root).join(format!("{keep_id}.md"));
let drop_path = rollout_summaries_dir(&root).join(format!("{drop_id}.md"));
tokio::fs::write(&keep_path, "keep")
.await
.expect("write keep");
tokio::fs::write(&drop_path, "drop")
.await
.expect("write drop");
let memories = vec![Stage1Output {
thread_id: ThreadId::try_from(keep_id.clone()).expect("thread id"),
source_updated_at: Utc.timestamp_opt(100, 0).single().expect("timestamp"),
raw_memory: "raw memory".to_string(),
rollout_summary: "short summary".to_string(),
rollout_slug: None,
rollout_path: PathBuf::from("/tmp/rollout-100.jsonl"),
cwd: PathBuf::from("/tmp/workspace"),
git_branch: None,
generated_at: Utc.timestamp_opt(101, 0).single().expect("timestamp"),
}];
sync_rollout_summaries_from_memories(&root, &memories, 100)
.await
.expect("sync rollout summaries");
rebuild_raw_memories_file_from_memories(&root, &memories, 100)
.await
.expect("rebuild raw memories");
assert!(
!tokio::fs::try_exists(&keep_path)
.await
.expect("check stale keep path"),
"sync should prune stale filename that used thread id only"
);
assert!(
!tokio::fs::try_exists(&drop_path)
.await
.expect("check stale drop path"),
"sync should prune stale filename for dropped thread"
);
let mut dir = tokio::fs::read_dir(rollout_summaries_dir(&root))
.await
.expect("open rollout summaries dir");
let mut files = Vec::new();
while let Some(entry) = dir.next_entry().await.expect("read dir entry") {
files.push(entry.file_name().to_string_lossy().to_string());
}
files.sort_unstable();
assert_eq!(files.len(), 1);
let canonical_rollout_summary_file = &files[0];
let raw_memories = tokio::fs::read_to_string(raw_memories_file(&root))
.await
.expect("read raw memories");
assert!(raw_memories.contains("raw memory"));
assert!(raw_memories.contains(&keep_id));
assert!(raw_memories.contains("cwd: /tmp/workspace"));
assert!(raw_memories.contains("rollout_path: /tmp/rollout-100.jsonl"));
assert!(raw_memories.contains(&format!(
"rollout_summary_file: {canonical_rollout_summary_file}"
)));
}
#[test]
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_missing() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(thread_id, None);
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
assert_eq!(
rollout_summary_file_stem_from_parts(
memory.thread_id,
memory.source_updated_at,
memory.rollout_slug.as_deref(),
),
FIXED_PREFIX
);
}
#[test]
fn rollout_summary_file_stem_sanitizes_and_truncates_slug() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(
thread_id,
Some("Unsafe Slug/With Spaces & Symbols + EXTRA_LONG_12345_67890_ABCDE_fghij_klmno"),
);
let stem = rollout_summary_file_stem(&memory);
let slug = stem
.strip_prefix(&format!("{FIXED_PREFIX}-"))
.expect("slug suffix should be present");
assert_eq!(slug.len(), 60);
assert_eq!(
slug,
"unsafe_slug_with_spaces___symbols___extra_long_12345_67890_a"
);
}
#[test]
fn rollout_summary_file_stem_uses_uuid_timestamp_and_hash_when_slug_is_empty() {
let thread_id = fixed_thread_id();
let memory = stage1_output_with_slug(thread_id, Some(""));
assert_eq!(rollout_summary_file_stem(&memory), FIXED_PREFIX);
}
}

281
codex-rs/memories/src/memory_trace.rs Normal file
View File

@@ -0,0 +1,281 @@
use anyhow::Result;
use codex_api::MemorySummarizeOutput;
use codex_api::RawMemory as ApiRawMemory;
use codex_api::RawMemoryMetadata as ApiRawMemoryMetadata;
use serde_json::Map;
use serde_json::Value;
use std::path::Path;
use std::path::PathBuf;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BuiltMemory {
pub memory_id: String,
pub source_path: PathBuf,
pub raw_memory: String,
pub memory_summary: String,
}
#[derive(Debug, Clone)]
pub struct PreparedTrace {
pub memory_id: String,
pub source_path: PathBuf,
pub payload: ApiRawMemory,
}
pub async fn load_trace_requests(trace_paths: &[PathBuf]) -> Result<Vec<PreparedTrace>> {
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?);
}
Ok(prepared)
}
pub fn build_memories_from_output(
prepared: Vec<PreparedTrace>,
output: Vec<MemorySummarizeOutput>,
) -> Result<Vec<BuiltMemory>> {
if output.len() != prepared.len() {
anyhow::bail!(
"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() {
anyhow::bail!("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() {
anyhow::bail!("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() {
anyhow::bail!(
"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)]
mod tests {
use super::load_trace_items;
use super::load_trace_text;
use super::normalize_trace_items;
use pretty_assertions::assert_eq;
use std::path::Path;
use tempfile::tempdir;
#[test]
fn normalize_trace_items_handles_payload_wrapper_and_message_role_filtering() {
let items = vec![
serde_json::json!({
"type": "response_item",
"payload": {"type": "message", "role": "assistant", "content": []}
}),
serde_json::json!({
"type": "response_item",
"payload": [
{"type": "message", "role": "user", "content": []},
{"type": "message", "role": "tool", "content": []},
{"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"}
]
}),
serde_json::json!({
"type": "not_response_item",
"payload": {"type": "message", "role": "assistant", "content": []}
}),
serde_json::json!({
"type": "message",
"role": "developer",
"content": []
}),
];
let normalized = normalize_trace_items(items, Path::new("trace.json")).expect("normalize");
let expected = vec![
serde_json::json!({"type": "message", "role": "assistant", "content": []}),
serde_json::json!({"type": "message", "role": "user", "content": []}),
serde_json::json!({"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "c1"}),
serde_json::json!({"type": "message", "role": "developer", "content": []}),
];
assert_eq!(normalized, expected);
}
#[test]
fn load_trace_items_supports_jsonl_arrays_and_objects() {
let text = r#"
{"type":"response_item","payload":{"type":"message","role":"assistant","content":[]}}
[{"type":"message","role":"user","content":[]},{"type":"message","role":"tool","content":[]}]
"#;
let loaded = load_trace_items(Path::new("trace.jsonl"), text).expect("load");
let expected = vec![
serde_json::json!({"type":"message","role":"assistant","content":[]}),
serde_json::json!({"type":"message","role":"user","content":[]}),
];
assert_eq!(loaded, expected);
}
#[tokio::test]
async fn load_trace_text_decodes_utf8_sig() {
let dir = tempdir().expect("tempdir");
let path = dir.path().join("trace.json");
tokio::fs::write(
&path,
[
0xEF, 0xBB, 0xBF, b'[', b'{', b'"', b't', b'y', b'p', b'e', b'"', b':', b'"', b'm',
b'e', b's', b's', b'a', b'g', b'e', b'"', b',', b'"', b'r', b'o', b'l', b'e', b'"',
b':', b'"', b'u', b's', b'e', b'r', b'"', b',', b'"', b'c', b'o', b'n', b't', b'e',
b'n', b't', b'"', b':', b'[', b']', b'}', b']',
],
)
.await
.expect("write");
let text = load_trace_text(&path).await.expect("decode");
assert!(text.starts_with('['));
}
}