feat: polluted memories (#13008)

Add a feature flag to disable memory creation for "polluted"

codex-rs/state/src/runtime/threads.rs

@@ -35,6 +35,14 @@ WHERE id = ?
             .transpose()
     }
 
+    pub async fn get_thread_memory_mode(&self, id: ThreadId) -> anyhow::Result<Option<String>> {
+        let row = sqlx::query("SELECT memory_mode FROM threads WHERE id = ?")
+            .bind(id.to_string())
+            .fetch_optional(self.pool.as_ref())
+            .await?;
+        Ok(row.and_then(|row| row.try_get("memory_mode").ok()))
+    }
+
     /// Get dynamic tools for a thread, if present.
     pub async fn get_dynamic_tools(
         &self,
@@ -199,6 +207,19 @@ FROM threads
             .await
     }
 
+    pub async fn set_thread_memory_mode(
+        &self,
+        thread_id: ThreadId,
+        memory_mode: &str,
+    ) -> anyhow::Result<bool> {
+        let result = sqlx::query("UPDATE threads SET memory_mode = ? WHERE id = ?")
+            .bind(memory_mode)
+            .bind(thread_id.to_string())
+            .execute(self.pool.as_ref())
+            .await?;
+        Ok(result.rows_affected() > 0)
+    }
+
     async fn upsert_thread_with_creation_memory_mode(
         &self,
         metadata: &crate::ThreadMetadata,
@@ -357,6 +378,16 @@ ON CONFLICT(thread_id, position) DO NOTHING
             }
             return Err(err);
         }
+        if let Some(memory_mode) = extract_memory_mode(items)
+            && let Err(err) = self
+                .set_thread_memory_mode(builder.id, memory_mode.as_str())
+                .await
+        {
+            if let Some(otel) = otel {
+                otel.counter(DB_ERROR_METRIC, 1, &[("stage", "set_thread_memory_mode")]);
+            }
+            return Err(err);
+        }
         let dynamic_tools = extract_dynamic_tools(items);
         if let Some(dynamic_tools) = dynamic_tools
             && let Err(err) = self
@@ -438,6 +469,16 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option<Option<Vec<
     })
 }
 
+pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option<String> {
+    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,
+    })
+}
+
 pub(super) fn push_thread_filters<'a>(
     builder: &mut QueryBuilder<'a, Sqlite>,
     archived_only: bool,
@@ -518,7 +559,11 @@ mod tests {
     use super::*;
     use crate::runtime::test_support::test_thread_metadata;
     use crate::runtime::test_support::unique_temp_dir;
+    use codex_protocol::protocol::SessionMeta;
+    use codex_protocol::protocol::SessionMetaLine;
+    use codex_protocol::protocol::SessionSource;
     use pretty_assertions::assert_eq;
+    use std::path::PathBuf;
 
     #[tokio::test]
     async fn upsert_thread_keeps_creation_memory_mode_for_existing_rows() {
@@ -557,4 +602,56 @@ mod tests {
                 .expect("memory mode should remain readable");
         assert_eq!(memory_mode, "disabled");
     }
+
+    #[tokio::test]
+    async fn apply_rollout_items_restores_memory_mode_from_session_meta() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("state db should initialize");
+        let thread_id =
+            ThreadId::from_string("00000000-0000-0000-0000-000000000456").expect("valid thread id");
+        let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone());
+
+        runtime
+            .upsert_thread(&metadata)
+            .await
+            .expect("initial upsert should succeed");
+
+        let builder = ThreadMetadataBuilder::new(
+            thread_id,
+            metadata.rollout_path.clone(),
+            metadata.created_at,
+            SessionSource::Cli,
+        );
+        let items = vec![RolloutItem::SessionMeta(SessionMetaLine {
+            meta: SessionMeta {
+                id: thread_id,
+                forked_from_id: None,
+                timestamp: metadata.created_at.to_rfc3339(),
+                cwd: PathBuf::new(),
+                originator: String::new(),
+                cli_version: String::new(),
+                source: SessionSource::Cli,
+                agent_nickname: None,
+                agent_role: None,
+                model_provider: None,
+                base_instructions: None,
+                dynamic_tools: None,
+                memory_mode: Some("polluted".to_string()),
+            },
+            git: None,
+        })];
+
+        runtime
+            .apply_rollout_items(&builder, &items, None, None)
+            .await
+            .expect("apply_rollout_items should succeed");
+
+        let memory_mode = runtime
+            .get_thread_memory_mode(thread_id)
+            .await
+            .expect("memory mode should load");
+        assert_eq!(memory_mode.as_deref(), Some("polluted"));
+    }
 }