Track collaboration mode instructions in world state (#33876)

## What changed

- Model collaboration mode as a world-state section whose persisted snapshot is
  the active mode.
- Emit collaboration instructions when the mode changes, while ignoring
  instruction text changes within the same mode.
- Restore persisted collaboration instructions when they are missing from
  retained history, including after a fork.

## Testing

- Add snapshot coverage for initial, unchanged, changed, and unknown mode state.
- Cover retained-history restoration and same-mode instruction changes.

GitOrigin-RevId: ffd1683499f09acef6516b0bfcb5bcb4d3f28df4

codex-rs/core/src/context/collaboration_mode_instructions.rs

@@ -1,40 +0,0 @@
-use super::ContextualUserFragment;
-use codex_protocol::config_types::CollaborationMode;
-use codex_protocol::protocol::COLLABORATION_MODE_CLOSE_TAG;
-use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG;
-
-#[derive(Debug, Clone, PartialEq)]
-pub(crate) struct CollaborationModeInstructions {
-    instructions: String,
-}
-
-impl CollaborationModeInstructions {
-    pub(crate) fn from_collaboration_mode(collaboration_mode: &CollaborationMode) -> Option<Self> {
-        collaboration_mode
-            .settings
-            .developer_instructions
-            .as_ref()
-            .filter(|instructions| !instructions.is_empty())
-            .map(|instructions| Self {
-                instructions: instructions.clone(),
-            })
-    }
-}
-
-impl ContextualUserFragment for CollaborationModeInstructions {
-    fn role(&self) -> &'static str {
-        "developer"
-    }
-
-    fn markers(&self) -> (&'static str, &'static str) {
-        Self::type_markers()
-    }
-
-    fn type_markers() -> (&'static str, &'static str) {
-        (COLLABORATION_MODE_OPEN_TAG, COLLABORATION_MODE_CLOSE_TAG)
-    }
-
-    fn body(&self) -> String {
-        self.instructions.clone()
-    }
-}

codex-rs/core/src/context/mod.rs

@@ -4,7 +4,6 @@ mod approved_command_prefix_saved;
 mod apps_instructions;
 mod available_plugins_instructions;
 mod available_skills_instructions;
-mod collaboration_mode_instructions;
 mod contextual_user_message;
 mod current_time_reminder;
 mod environment_context;
@@ -46,7 +45,6 @@ pub use codex_context_fragments::ContextualUserFragment;
 pub(crate) use codex_context_fragments::FragmentRegistration;
 pub(crate) use codex_context_fragments::FragmentRegistrationProxy;
 pub(crate) use codex_core_skills::SkillInstructions;
-pub(crate) use collaboration_mode_instructions::CollaborationModeInstructions;
 pub(crate) use contextual_user_message::is_contextual_user_fragment;
 pub(crate) use contextual_user_message::parse_visible_hook_prompt_message;
 pub(crate) use current_time_reminder::CurrentTimeReminder;

codex-rs/core/src/context/world_state/collaboration_mode.rs

@@ -0,0 +1,91 @@
+use super::PreviousSectionState;
+use super::WorldStateSection;
+use crate::context::ContextualUserFragment;
+use codex_protocol::config_types::CollaborationMode;
+use codex_protocol::config_types::ModeKind;
+use codex_protocol::protocol::COLLABORATION_MODE_CLOSE_TAG;
+use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG;
+
+/// Collaboration-mode instructions currently visible to the model.
+#[derive(Clone, Debug)]
+pub(crate) struct CollaborationModeState {
+    mode: ModeKind,
+    instructions: String,
+}
+
+impl CollaborationModeState {
+    pub(crate) fn from_collaboration_mode(collaboration_mode: &CollaborationMode) -> Option<Self> {
+        collaboration_mode
+            .settings
+            .developer_instructions
+            .clone()
+            .filter(|instructions| !instructions.is_empty())
+            .map(|instructions| Self {
+                mode: collaboration_mode.mode,
+                instructions,
+            })
+    }
+}
+
+impl WorldStateSection for CollaborationModeState {
+    const ID: &'static str = "collaboration_mode";
+    type Snapshot = ModeKind;
+
+    fn snapshot(&self) -> Self::Snapshot {
+        self.mode
+    }
+
+    fn matches_legacy_fragment(role: &str, text: &str) -> bool {
+        role == "developer" && CollaborationModeInstructions::matches_text(text)
+    }
+
+    fn has_retained_fragment_matcher() -> bool {
+        true
+    }
+
+    fn matches_retained_fragment(role: &str, text: &str) -> bool {
+        Self::matches_legacy_fragment(role, text)
+    }
+
+    fn render_diff(
+        &self,
+        previous: PreviousSectionState<'_, Self::Snapshot>,
+    ) -> Option<Box<dyn ContextualUserFragment>> {
+        if matches!(previous, PreviousSectionState::Known(previous) if previous == &self.mode)
+            || matches!(previous, PreviousSectionState::Unknown)
+        {
+            return None;
+        }
+
+        Some(Box::new(CollaborationModeInstructions {
+            instructions: self.instructions.clone(),
+        }))
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+struct CollaborationModeInstructions {
+    instructions: String,
+}
+
+impl ContextualUserFragment for CollaborationModeInstructions {
+    fn role(&self) -> &'static str {
+        "developer"
+    }
+
+    fn markers(&self) -> (&'static str, &'static str) {
+        Self::type_markers()
+    }
+
+    fn type_markers() -> (&'static str, &'static str) {
+        (COLLABORATION_MODE_OPEN_TAG, COLLABORATION_MODE_CLOSE_TAG)
+    }
+
+    fn body(&self) -> String {
+        self.instructions.clone()
+    }
+}
+
+#[cfg(test)]
+#[path = "collaboration_mode_tests.rs"]
+mod tests;

codex-rs/core/src/context/world_state/collaboration_mode_tests.rs

@@ -0,0 +1,66 @@
+use super::super::PreviousSectionState;
+use super::super::test_support::render_section_cases;
+use super::*;
+use crate::context::world_state::WorldState;
+use codex_protocol::config_types::ModeKind;
+use codex_protocol::config_types::Settings;
+use codex_protocol::models::ResponseItem;
+
+#[test]
+fn snapshots() {
+    use PreviousSectionState::Absent;
+    use PreviousSectionState::Known;
+    use PreviousSectionState::Unknown;
+
+    let default = collaboration_mode_state(ModeKind::Default, "pair with the user");
+    let old_default = collaboration_mode_state(ModeKind::Default, "old instructions");
+    let new_default = collaboration_mode_state(ModeKind::Default, "new instructions");
+    let plan = collaboration_mode_state(ModeKind::Plan, "make a plan");
+
+    insta::assert_snapshot!(render_section_cases(&[
+        (Absent, Absent),
+        (Absent, Known(&default)),
+        (Known(&default), Known(&default)),
+        (Known(&old_default), Known(&new_default)),
+        (Known(&default), Known(&plan)),
+        (Unknown, Known(&default)),
+    ]));
+}
+
+#[test]
+fn persisted_instructions_are_restored_only_when_missing_from_history() {
+    let state = collaboration_mode_state(ModeKind::Default, "pair with the user");
+    let retained: ResponseItem = ContextualUserFragment::into(CollaborationModeInstructions {
+        instructions: state.instructions.clone(),
+    });
+    let mut world_state = WorldState::default();
+    world_state.add_section(state);
+    let snapshot = world_state.snapshot();
+
+    assert!(
+        world_state
+            .render_history_diff(/*previous*/ None, std::slice::from_ref(&retained))
+            .is_empty()
+    );
+    assert_eq!(
+        world_state.render_history_diff(Some(&snapshot), &[]).len(),
+        1,
+    );
+    assert!(
+        world_state
+            .render_history_diff(Some(&snapshot), &[retained])
+            .is_empty()
+    );
+}
+
+fn collaboration_mode_state(mode: ModeKind, instructions: &str) -> CollaborationModeState {
+    CollaborationModeState::from_collaboration_mode(&CollaborationMode {
+        mode,
+        settings: Settings {
+            model: "test-model".to_string(),
+            reasoning_effort: None,
+            developer_instructions: Some(instructions.to_string()),
+        },
+    })
+    .expect("test collaboration mode should have instructions")
+}

codex-rs/core/src/context/world_state/mod.rs

@@ -1,5 +1,6 @@
 mod agents_md;
 mod apps_instructions;
+mod collaboration_mode;
 mod environment;
 mod environments_instructions;
 mod plugins_instructions;
@@ -22,6 +23,7 @@ use std::fmt;
 
 pub(crate) use agents_md::AgentsMdState;
 pub(crate) use apps_instructions::AppsInstructionsState;
+pub(crate) use collaboration_mode::CollaborationModeState;
 pub(crate) use environment::EnvironmentsState;
 pub(crate) use environments_instructions::EnvironmentsInstructionsState;
 pub(crate) use plugins_instructions::PluginsInstructionsState;

codex-rs/core/src/context/world_state/snapshots/codex_core__context__world_state__collaboration_mode__tests__snapshots.snap

@@ -0,0 +1,21 @@
+---
+source: core/src/context/world_state/collaboration_mode_tests.rs
+expression: "render_section_cases(&[(Absent, Absent), (Absent, Known(&default)),\n(Known(&default), Known(&default)),\n(Known(&old_default), Known(&new_default)), (Known(&default), Known(&plan)),\n(Unknown, Known(&default)),])"
+---
+Absent -> Absent
+None
+
+Absent -> "default" (role - developer)
+<collaboration_mode>pair with the user</collaboration_mode>
+
+"default" -> "default"
+None
+
+"default" -> "default"
+None
+
+"default" -> "plan" (role - developer)
+<collaboration_mode>make a plan</collaboration_mode>
+
+Unknown -> "default"
+None

codex-rs/core/src/context_manager/updates.rs

@@ -1,5 +1,4 @@
 use crate::context::ApprovalPromptContext;
-use crate::context::CollaborationModeInstructions;
 use crate::context::ContextualUserFragment;
 use crate::context::ModelSwitchInstructions;
 use crate::context::MultiAgentModeInstructions;
@@ -65,25 +64,6 @@ fn build_permissions_update_item(
     )
 }
 
-fn build_collaboration_mode_update_item(
-    previous: Option<&TurnContextItem>,
-    next: &TurnContext,
-) -> Option<String> {
-    if !next.config.include_collaboration_mode_instructions {
-        return None;
-    }
-
-    let prev = previous?;
-    let collaboration_mode = next.collaboration_mode();
-    if prev.collaboration_mode.as_ref() != Some(&collaboration_mode) {
-        // If the next mode has empty developer instructions, this returns None and we emit no
-        // update, so prior collaboration instructions remain in the prompt history.
-        Some(CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode)?.render())
-    } else {
-        None
-    }
-}
-
 fn build_multi_agent_mode_update_item(
     previous: Option<&TurnContextItem>,
     next: &TurnContext,
@@ -257,7 +237,6 @@ pub(crate) fn build_settings_update_items(
         // any other context diffs on this turn.
         build_model_instructions_update_item(previous_turn_settings, next),
         build_permissions_update_item(previous, next, exec_policy),
-        build_collaboration_mode_update_item(previous, next),
         build_multi_agent_mode_update_item(previous, next),
         build_realtime_update_item(previous, previous_turn_settings, next),
         build_personality_update_item(previous, next, personality_feature_enabled),

codex-rs/core/src/session/mod.rs

@@ -24,7 +24,6 @@ use crate::config::resolve_tool_suggest_config_from_layer_stack;
 use crate::context::ApprovalPromptContext;
 use crate::context::ApprovedCommandPrefixSaved;
 use crate::context::AvailableSkillsInstructions;
-use crate::context::CollaborationModeInstructions;
 use crate::context::ContextualUserFragment;
 use crate::context::MultiAgentModeInstructions;
 use crate::context::NetworkRuleSaved;
@@ -3196,7 +3195,6 @@ impl Session {
         let (
             reference_context_item,
             previous_turn_settings,
-            collaboration_mode,
             base_instructions,
             session_source,
             auto_compact_window_ids,
@@ -3205,7 +3203,6 @@ impl Session {
             (
                 state.reference_context_item(),
                 state.previous_turn_settings(),
-                state.session_configuration.collaboration_mode.clone(),
                 state.session_configuration.base_instructions.clone(),
                 state.session_configuration.session_source.clone(),
                 state.auto_compact_window_ids(),
@@ -3262,13 +3259,6 @@ impl Session {
         {
             developer_sections.push(developer_instructions.to_string());
         }
-        // Add developer instructions from collaboration_mode if they exist and are non-empty
-        if turn_context.config.include_collaboration_mode_instructions
-            && let Some(collab_instructions) =
-                CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode)
-        {
-            developer_sections.push(collab_instructions.render());
-        }
         if let Some(realtime_update) = crate::context_manager::updates::build_initial_realtime_item(
             reference_context_item.as_ref(),
             previous_turn_settings.as_ref(),

codex-rs/core/src/session/snapshots/codex_core__codex_tests__fork_startup_context_then_first_turn_diff.snap

@@ -8,7 +8,6 @@ Scenario: First request after fork when startup preserves the parent baseline, t
 00:message/developer:<PERMISSIONS_INSTRUCTIONS>
 01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
 02:message/user:fork seed
-03:message/developer[2]:
+03:message/developer:<PERMISSIONS_INSTRUCTIONS>
-    [01] <PERMISSIONS_INSTRUCTIONS>
+04:message/developer:<collaboration_mode>Fork turn collaboration instructions.</collaboration_mode>
-    [02] <collaboration_mode>Fork turn collaboration instructions.</collaboration_mode>
+05:message/user:after fork
-04:message/user:after fork

codex-rs/core/src/session/world_state.rs

@@ -3,6 +3,7 @@ use super::step_context::StepContext;
 use crate::connectors;
 use crate::context::world_state::AgentsMdState;
 use crate::context::world_state::AppsInstructionsState;
+use crate::context::world_state::CollaborationModeState;
 use crate::context::world_state::EnvironmentsInstructionsState;
 use crate::context::world_state::EnvironmentsState;
 use crate::context::world_state::PluginsInstructionsState;
@@ -32,6 +33,12 @@ impl Session {
 
         let mut world_state = WorldState::default();
         world_state.add_section(AgentsMdState::new(step_context.loaded_agents_md.as_deref()));
+        if turn_context.config.include_collaboration_mode_instructions
+            && let Some(collaboration_mode) =
+                CollaborationModeState::from_collaboration_mode(&turn_context.collaboration_mode())
+        {
+            world_state.add_section(collaboration_mode);
+        }
         if turn_context.config.include_environment_context {
             world_state.add_section(
                 EnvironmentsState::from_turn_context_with_environments(

codex-rs/core/tests/suite/collaboration_instructions.rs

@@ -355,7 +355,7 @@ async fn user_turn_overrides_collaboration_instructions_after_override() -> Resu
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn collaboration_mode_update_emits_new_instruction_message() -> Result<()> {
+async fn collaboration_mode_update_ignores_instruction_changes_within_same_mode() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
     let server = start_mock_server().await;
@@ -425,7 +425,7 @@ async fn collaboration_mode_update_emits_new_instruction_message() -> Result<()>
     let first_text = collab_xml(first_text);
     let second_text = collab_xml(second_text);
     assert_eq!(count_messages_containing(&dev_texts, &first_text), 1);
-    assert_eq!(count_messages_containing(&dev_texts, &second_text), 1);
+    assert_eq!(count_messages_containing(&dev_texts, &second_text), 0);
 
     Ok(())
 }