Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions codex-rs/app-server/tests/common/rollout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ pub fn create_fake_rollout_with_source(
model_provider: model_provider.map(str::to_string),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
};
let payload = serde_json::to_value(SessionMetaLine {
meta,
Expand Down Expand Up @@ -165,6 +166,7 @@ pub fn create_fake_rollout_with_text_elements(
model_provider: model_provider.map(str::to_string),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
};
let payload = serde_json::to_value(SessionMetaLine {
meta,
Expand Down
4 changes: 4 additions & 0 deletions codex-rs/core/config.schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -644,6 +644,10 @@
"format": "int64",
"type": "integer"
},
"no_memories_if_mcp_or_web_search": {
"description": "When `true`, web searches and MCP tool calls mark the thread `memory_mode` as `\"polluted\"`.",
"type": "boolean"
},
"phase_1_model": {
"description": "Model used for thread summarisation.",
"type": "string"
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/core/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2490,6 +2490,7 @@ persistence = "none"

let memories = r#"
[memories]
no_memories_if_mcp_or_web_search = true
generate_memories = false
use_memories = false
max_raw_memories_for_global = 512
Expand All @@ -2504,6 +2505,7 @@ phase_2_model = "gpt-5"
toml::from_str::<ConfigToml>(memories).expect("TOML deserialization should succeed");
assert_eq!(
Some(MemoriesToml {
no_memories_if_mcp_or_web_search: Some(true),
generate_memories: Some(false),
use_memories: Some(false),
max_raw_memories_for_global: Some(512),
Expand All @@ -2526,6 +2528,7 @@ phase_2_model = "gpt-5"
assert_eq!(
config.memories,
MemoriesConfig {
no_memories_if_mcp_or_web_search: true,
generate_memories: false,
use_memories: false,
max_raw_memories_for_global: 512,
Expand Down
7 changes: 7 additions & 0 deletions codex-rs/core/src/config/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,8 @@ pub struct FeedbackConfigToml {
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[schemars(deny_unknown_fields)]
pub struct MemoriesToml {
/// When `true`, web searches and MCP tool calls mark the thread `memory_mode` as `"polluted"`.
pub no_memories_if_mcp_or_web_search: Option<bool>,
/// When `false`, newly created threads are stored with `memory_mode = "disabled"` in the state DB.
pub generate_memories: Option<bool>,
/// When `false`, skip injecting memory usage instructions into developer prompts.
Expand All @@ -394,6 +396,7 @@ pub struct MemoriesToml {
/// Effective memories settings after defaults are applied.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoriesConfig {
pub no_memories_if_mcp_or_web_search: bool,
pub generate_memories: bool,
pub use_memories: bool,
pub max_raw_memories_for_global: usize,
Expand All @@ -408,6 +411,7 @@ pub struct MemoriesConfig {
impl Default for MemoriesConfig {
fn default() -> Self {
Self {
no_memories_if_mcp_or_web_search: false,
generate_memories: true,
use_memories: true,
max_raw_memories_for_global: DEFAULT_MEMORIES_MAX_RAW_MEMORIES_FOR_GLOBAL,
Expand All @@ -425,6 +429,9 @@ impl From<MemoriesToml> for MemoriesConfig {
fn from(toml: MemoriesToml) -> Self {
let defaults = Self::default();
Self {
no_memories_if_mcp_or_web_search: toml
.no_memories_if_mcp_or_web_search
.unwrap_or(defaults.no_memories_if_mcp_or_web_search),
generate_memories: toml.generate_memories.unwrap_or(defaults.generate_memories),
use_memories: toml.use_memories.unwrap_or(defaults.use_memories),
max_raw_memories_for_global: toml
Expand Down
19 changes: 19 additions & 0 deletions codex-rs/core/src/mcp_tool_call.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ use crate::protocol::EventMsg;
use crate::protocol::McpInvocation;
use crate::protocol::McpToolCallBeginEvent;
use crate::protocol::McpToolCallEndEvent;
use crate::state_db;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputPayload;
Expand Down Expand Up @@ -121,6 +122,7 @@ pub(crate) async fn handle_mcp_tool_call(
});
notify_mcp_tool_call_event(sess.as_ref(), turn_context, tool_call_begin_event)
.await;
maybe_mark_thread_memory_mode_polluted(sess.as_ref(), turn_context).await;

let start = Instant::now();
let result = sess
Expand Down Expand Up @@ -189,6 +191,7 @@ pub(crate) async fn handle_mcp_tool_call(
invocation: invocation.clone(),
});
notify_mcp_tool_call_event(sess.as_ref(), turn_context, tool_call_begin_event).await;
maybe_mark_thread_memory_mode_polluted(sess.as_ref(), turn_context).await;

let start = Instant::now();
// Perform the tool call.
Expand Down Expand Up @@ -224,6 +227,22 @@ pub(crate) async fn handle_mcp_tool_call(
ResponseInputItem::McpToolCallOutput { call_id, result }
}

async fn maybe_mark_thread_memory_mode_polluted(sess: &Session, turn_context: &TurnContext) {
if !turn_context
.config
.memories
.no_memories_if_mcp_or_web_search
{
return;
}
state_db::mark_thread_memory_mode_polluted(
sess.services.state_db.as_deref(),
sess.conversation_id,
"mcp_tool_call",
)
.await;
}

fn sanitize_mcp_tool_result_for_model(
supports_image_input: bool,
result: Result<CallToolResult, String>,
Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/src/personality_migration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ mod tests {
model_provider: None,
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
},
git: None,
};
Expand Down
83 changes: 83 additions & 0 deletions codex-rs/core/src/rollout/metadata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,13 @@ pub(crate) async fn extract_metadata_from_rollout(
}
Ok(ExtractionOutcome {
metadata,
memory_mode: items.iter().rev().find_map(|item| match item {
RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(),
RolloutItem::ResponseItem(_)
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::EventMsg(_) => None,
}),
parse_errors,
})
}
Expand Down Expand Up @@ -272,6 +279,7 @@ pub(crate) async fn backfill_sessions(
);
}
let mut metadata = outcome.metadata;
let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Avoid forcing enabled mode during backfill replay

Backfill now defaults missing rollout memory_mode to "enabled" and writes it back unconditionally. On any backfill pass over legacy rollouts, existing threads.memory_mode values can be clobbered, causing polluted/disabled threads to become eligible for stage1/phase2 processing again.

Useful? React with 👍 / 👎.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't happen

if rollout.archived && metadata.archived_at.is_none() {
let fallback_archived_at = metadata.updated_at;
metadata.archived_at = file_modified_time_utc(&rollout.path)
Expand All @@ -282,6 +290,17 @@ pub(crate) async fn backfill_sessions(
stats.failed = stats.failed.saturating_add(1);
warn!("failed to upsert rollout {}: {err}", rollout.path.display());
} else {
if let Err(err) = runtime
.set_thread_memory_mode(metadata.id, memory_mode.as_str())
.await
{
stats.failed = stats.failed.saturating_add(1);
warn!(
"failed to restore memory mode for {}: {err}",
rollout.path.display()
);
continue;
}
stats.upserted = stats.upserted.saturating_add(1);
if let Ok(meta_line) =
rollout::list::read_session_meta_line(&rollout.path).await
Expand Down Expand Up @@ -519,6 +538,7 @@ mod tests {
model_provider: Some("openai".to_string()),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
};
let session_meta_line = SessionMetaLine {
meta: session_meta,
Expand All @@ -543,9 +563,71 @@ mod tests {
expected.updated_at = file_modified_time_utc(&path).await.expect("mtime");

assert_eq!(outcome.metadata, expected);
assert_eq!(outcome.memory_mode, None);
assert_eq!(outcome.parse_errors, 0);
}

#[tokio::test]
async fn extract_metadata_from_rollout_returns_latest_memory_mode() {
let dir = tempdir().expect("tempdir");
let uuid = Uuid::new_v4();
let id = ThreadId::from_string(&uuid.to_string()).expect("thread id");
let path = dir
.path()
.join(format!("rollout-2026-01-27T12-34-56-{uuid}.jsonl"));

let session_meta = SessionMeta {
id,
forked_from_id: None,
timestamp: "2026-01-27T12:34:56Z".to_string(),
cwd: dir.path().to_path_buf(),
originator: "cli".to_string(),
cli_version: "0.0.0".to_string(),
source: SessionSource::default(),
agent_nickname: None,
agent_role: None,
model_provider: Some("openai".to_string()),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
};
let polluted_meta = SessionMeta {
memory_mode: Some("polluted".to_string()),
..session_meta.clone()
};
let lines = vec![
RolloutLine {
timestamp: "2026-01-27T12:34:56Z".to_string(),
item: RolloutItem::SessionMeta(SessionMetaLine {
meta: session_meta,
git: None,
}),
},
RolloutLine {
timestamp: "2026-01-27T12:35:00Z".to_string(),
item: RolloutItem::SessionMeta(SessionMetaLine {
meta: polluted_meta,
git: None,
}),
},
];
let mut file = File::create(&path).expect("create rollout");
for line in lines {
writeln!(
file,
"{}",
serde_json::to_string(&line).expect("serialize rollout line")
)
.expect("write rollout line");
}

let outcome = extract_metadata_from_rollout(&path, "openai", None)
.await
.expect("extract");

assert_eq!(outcome.memory_mode.as_deref(), Some("polluted"));
}

#[test]
fn builder_from_items_falls_back_to_filename() {
let dir = tempdir().expect("tempdir");
Expand Down Expand Up @@ -669,6 +751,7 @@ mod tests {
model_provider: Some("test-provider".to_string()),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
};
let session_meta_line = SessionMetaLine {
meta: session_meta,
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/rollout/recorder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -412,6 +412,8 @@ impl RolloutRecorder {
} else {
Some(dynamic_tools)
},
memory_mode: (!config.memories.generate_memories)
.then_some("disabled".to_string()),
};

(
Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/src/rollout/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1109,6 +1109,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> {
model_provider: Some("test-provider".into()),
base_instructions: None,
dynamic_tools: None,
memory_mode: None,
},
git: None,
}),
Expand Down
24 changes: 24 additions & 0 deletions codex-rs/core/src/state_db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,19 @@ pub async fn persist_dynamic_tools(
}
}

pub async fn mark_thread_memory_mode_polluted(
context: Option<&codex_state::StateRuntime>,
thread_id: ThreadId,
stage: &str,
) {
let Some(ctx) = context else {
return;
};
if let Err(err) = ctx.mark_thread_memory_mode_polluted(thread_id).await {
warn!("state db mark_thread_memory_mode_polluted failed during {stage}: {err}");
}
}

/// Reconcile rollout items into SQLite, falling back to scanning the rollout file.
pub async fn reconcile_rollout(
context: Option<&codex_state::StateRuntime>,
Expand Down Expand Up @@ -375,6 +388,7 @@ pub async fn reconcile_rollout(
}
};
let mut metadata = outcome.metadata;
let memory_mode = outcome.memory_mode.unwrap_or_else(|| "enabled".to_string());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve memory_mode when reconcile lacks rollout metadata

Do not default to "enabled" here. For rollouts without session_meta.memory_mode (legacy files, or threads later marked polluted only in DB), this path overwrites the current DB value and re-enables memory generation. A read-repair/reconcile can therefore silently undo disabled/polluted state.

Useful? React with 👍 / 👎.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't do

metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd);
match archived_only {
Some(true) if metadata.archived_at.is_none() => {
Expand All @@ -392,6 +406,16 @@ pub async fn reconcile_rollout(
);
return;
}
if let Err(err) = ctx
.set_thread_memory_mode(metadata.id, memory_mode.as_str())
.await
{
warn!(
"state db reconcile_rollout memory_mode update failed {}: {err}",
rollout_path.display()
);
return;
}
if let Ok(meta_line) = crate::rollout::list::read_session_meta_line(rollout_path).await {
persist_dynamic_tools(
Some(ctx),
Expand Down
22 changes: 22 additions & 0 deletions codex-rs/core/src/stream_events_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,31 @@ pub(crate) async fn record_completed_response_item(
) {
sess.record_conversation_items(turn_context, std::slice::from_ref(item))
.await;
maybe_mark_thread_memory_mode_polluted_from_web_search(sess, turn_context, item).await;
record_stage1_output_usage_for_completed_item(turn_context, item).await;
}

async fn maybe_mark_thread_memory_mode_polluted_from_web_search(
sess: &Session,
turn_context: &TurnContext,
item: &ResponseItem,
) {
if !turn_context
.config
.memories
.no_memories_if_mcp_or_web_search
|| !matches!(item, ResponseItem::WebSearchCall { .. })
{
return;
}
state_db::mark_thread_memory_mode_polluted(
sess.services.state_db.as_deref(),
sess.conversation_id,
"record_completed_response_item",
)
.await;
}

async fn record_stage1_output_usage_for_completed_item(
turn_context: &TurnContext,
item: &ResponseItem,
Expand Down
Loading
Loading