Preserve streamed assistant message order in retained context (#47584)

## Why

A user can reply to a streamed question before its completed message is recorded. Assigning the assistant message's order at completion can place the reply before the question in retained context.

## What changed

Reserve an assistant message's order when its source stream starts in thread-owned Guardian context, then reuse that reservation when recording the message. Reserve completion-only messages before publishing their text, and give text and plan output from the same source message a single order.

## Testing

Add streaming regression coverage that accepts a reply before the question completes and verifies question-before-reply ordering both live and after resume. Cover server-provided and generated message IDs, plan output, and text combined with a plan.

GitOrigin-RevId: cb67e7459c804df7beb0841908244202448c655a

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

@@ -3623,6 +3623,9 @@ impl Session {
                 .current_time_reminder
                 .note_recorded_items(&response_items);
             if self.guardian_context_mode == GuardianContextMode::ThreadOwned {
+                let pending_orders = turn_context
+                    .extension_data
+                    .get::<retained_context::PendingAssistantMessageOrders>();
                 for envelope in &mut items {
                     if envelope
                         .metadata
@@ -3635,12 +3638,22 @@ impl Session {
                         || matches!(&envelope.item, ResponseItem::FunctionCall { .. })
                         || crate::context::is_user_authorization_message(&envelope.item)
                     {
-                        // Preserve accepted input order; assign delivery order to other originals.
+                        let message_order = pending_orders.as_ref().and_then(|orders| {
+                            orders
+                                .0
+                                .lock()
+                                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                                .remove(envelope.item.id()?.as_str())
+                        });
+                        // Preserve input acceptance and source-message start order.
+                        // Synthetic messages still receive their order here.
                         envelope
                             .metadata
                             .get_or_insert_default()
                             .user_input_order
-                            .get_or_insert_with(|| state.history.reserve_input_order());
+                            .get_or_insert_with(|| {
+                                message_order.unwrap_or_else(|| state.history.reserve_input_order())
+                            });
                     }
                 }
             }

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

@@ -1,13 +1,54 @@
-//! Orders retained inputs and records host-verified facts at the context checkpoint boundary.
+//! Orders retained inputs at acceptance and assistant messages at source-stream start.
+
+use std::collections::HashMap;
+use std::sync::Mutex;
+use std::sync::PoisonError;
 
 use crate::context::GuardianContextMode;
 use codex_history::RetainedContextEvent;
 use codex_history::RolloutItem;
+use codex_protocol::models::ResponseItem;
 
 use super::Session;
+use super::TurnContext;
 use super::thread_settings;
 
+/// Only unrecorded starts need a reservation; abandoned entries expire with the turn.
+#[derive(Default)]
+pub(super) struct PendingAssistantMessageOrders(pub(super) Mutex<HashMap<String, u64>>);
+
 impl Session {
+    /// Reserve before deriving display items, including plans, from the source message.
+    /// Completion-only responses use the same path before publishing their text.
+    pub(super) async fn reserve_assistant_message_order(
+        &self,
+        turn_context: &TurnContext,
+        item: &ResponseItem,
+    ) {
+        if self.guardian_context_mode == GuardianContextMode::ThreadOwned
+            && let ResponseItem::Message {
+                id: Some(id), role, ..
+            } = item
+            && role == "assistant"
+        {
+            let mut state = self.state.lock().await;
+            if !state
+                .history
+                .raw_items()
+                .any(|item| item.id().is_some_and(|recorded_id| recorded_id == id))
+            {
+                turn_context
+                    .extension_data
+                    .get_or_init(PendingAssistantMessageOrders::default)
+                    .0
+                    .lock()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .entry(id.to_string())
+                    .or_insert_with(|| state.history.reserve_input_order());
+            }
+        }
+    }
+
     /// Legacy mode neither reserves a sequence nor takes the session-state lock.
     pub(crate) async fn reserve_user_input_order(&self) -> Option<u64> {
         if self.guardian_context_mode == GuardianContextMode::Legacy {

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

@@ -2603,6 +2603,8 @@ async fn try_run_sampling_request(
             }
             ResponseEvent::OutputItemDone(mut item) => {
                 assign_missing_streamed_response_item_id(&mut item, active_item.as_ref());
+                sess.reserve_assistant_message_order(&turn_context, &item)
+                    .await;
                 if analytics_tool_call_ids.len() < MAX_ANALYTICS_TOOL_CALL_IDS_PER_RESPONSE {
                     let call_id = match &item {
                         ResponseItem::FunctionCall { call_id, .. }
@@ -2715,6 +2717,8 @@ async fn try_run_sampling_request(
             }
             ResponseEvent::OutputItemAdded(mut item) => {
                 assign_missing_streamed_response_item_id(&mut item, /*active_item*/ None);
+                sess.reserve_assistant_message_order(&turn_context, &item)
+                    .await;
                 if let ResponseItem::CustomToolCall {
                     call_id,
                     name,

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

@@ -31,12 +31,16 @@ use codex_history::RolloutItem;
 use codex_history::VerifiedAnswer;
 use codex_history::VerifiedQuestionAnswer;
 use codex_protocol::ResponseItemId;
+use codex_protocol::config_types::CollaborationMode;
+use codex_protocol::config_types::ModeKind;
+use codex_protocol::config_types::Settings;
 use codex_protocol::mcp::ClientMcpExtensions;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::ThreadHistoryMode;
+use codex_protocol::protocol::ThreadSettingsOverrides;
 use codex_protocol::request_user_input::RequestUserInputAnswer;
 use codex_protocol::request_user_input::RequestUserInputResponse;
 use codex_protocol::user_input::UserInput;
@@ -47,11 +51,16 @@ use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_function_call;
 use core_test_support::responses::ev_function_call_with_namespace;
+use core_test_support::responses::ev_message_item_added;
+use core_test_support::responses::ev_output_text_delta;
+use core_test_support::responses::ev_response_created;
 use core_test_support::responses::mount_sse_once_match;
 use core_test_support::responses::mount_sse_sequence;
 use core_test_support::responses::sse;
 use core_test_support::responses::start_mock_server;
 use core_test_support::skip_if_no_network;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
 use core_test_support::test_codex::TestCodex;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
@@ -60,6 +69,7 @@ use pretty_assertions::assert_eq;
 use serde_json::json;
 use test_case::test_case;
 use tokio::sync::Notify;
+use tokio::sync::oneshot;
 use wiremock::MockServer;
 use wiremock::matchers::header;
 
@@ -228,6 +238,143 @@ async fn resume(test: &TestCodex, thread: &CodexThread) -> Result<Arc<CodexThrea
         .thread)
 }
 
+#[test_case("question", ModeKind::Default, ""; "server message id")]
+#[test_case("", ModeKind::Default, ""; "generated message id")]
+#[test_case("question", ModeKind::Plan, ""; "plan with server message id")]
+#[test_case("", ModeKind::Plan, ""; "plan with generated message id")]
+#[test_case("question", ModeKind::Plan, "Here is the plan:\n"; "text and plan share one order")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn streamed_question_precedes_reply_across_resume(
+    message_id: &str,
+    mode: ModeKind,
+    preamble: &str,
+) -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    const INITIAL: &str = "Prepare the deployment. Staging only.";
+    const QUESTION: &str = "Deploy abc123?";
+    const ANSWER: &str = "Yes.";
+    let (prefix, suffix) = if mode == ModeKind::Plan {
+        (format!("{preamble}<proposed_plan>\n"), "\n</proposed_plan>")
+    } else {
+        (String::new(), "")
+    };
+    let question = format!("{prefix}{QUESTION}{suffix}");
+    let (release, gate) = oneshot::channel();
+    let (server, _completions) = start_streaming_sse_server(vec![
+        vec![
+            StreamingSseChunk {
+                gate: None,
+                body: sse(vec![
+                    ev_response_created("question-response"),
+                    ev_message_item_added(message_id, ""),
+                    ev_output_text_delta(&format!("{prefix}Deploy ")),
+                    ev_output_text_delta("abc123?"),
+                    ev_output_text_delta(suffix),
+                ]),
+            },
+            StreamingSseChunk {
+                gate: Some(gate),
+                body: sse(vec![
+                    ev_assistant_message(message_id, &question),
+                    ev_completed("question-response"),
+                ]),
+            },
+        ],
+        vec![StreamingSseChunk {
+            gate: None,
+            body: sse(vec![ev_completed("answer-response")]),
+        }],
+    ])
+    .await;
+    let test = test_codex()
+        .with_model("gpt-5.4")
+        .with_history_mode(ThreadHistoryMode::Paginated)
+        .with_config(|config| {
+            config.experimental_thread_store = ThreadStoreConfig::Local;
+            config
+                .features
+                .enable(Feature::GuardianThreadContext)
+                .expect("enable retained context");
+        })
+        .build_with_streaming_server(&server)
+        .await?;
+    test.codex
+        .start_or_steer_turn(
+            TurnInputRequest::user_input(vec![UserInput::Text {
+                text: INITIAL.to_owned(),
+                text_elements: Vec::new(),
+            }])
+            .with_thread_settings(ThreadSettingsOverrides {
+                collaboration_mode: Some(CollaborationMode {
+                    mode,
+                    settings: Settings {
+                        model: test.session_configured.model.clone(),
+                        reasoning_effort: None,
+                        developer_instructions: None,
+                    },
+                }),
+                ..Default::default()
+            }),
+        )
+        .await?;
+    wait_for_event(&test.codex, |event| match event {
+        EventMsg::AgentMessageContentDelta(delta) => delta.delta.contains("abc123?"),
+        EventMsg::PlanDelta(delta) => delta.delta.contains("abc123?"),
+        _ => false,
+    })
+    .await;
+    // The reply is accepted while the question's completed item is still withheld.
+    test.codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: ANSWER.to_owned(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    release.send(()).expect("release question completion");
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+
+    let history = test.codex.conversation_history_snapshot().await;
+    let retained = history
+        .retained_context()
+        .context("live retained context")?;
+    assert_eq!(
+        retained
+            .ordered_entries()
+            .map(|(order, entry)| match entry {
+                RetainedContextEntry::UserMessage(message) =>
+                    (order, "user", message.text.as_str()),
+                RetainedContextEntry::AssistantMessage(message) => {
+                    (order, "assistant", message.text.as_str())
+                }
+                RetainedContextEntry::VerifiedAnswer(_) => panic!("expected plain chat messages"),
+            })
+            .collect::<Vec<_>>(),
+        vec![
+            (RetainedContextOrder::Local(0), "user", INITIAL),
+            (
+                RetainedContextOrder::Local(1),
+                "assistant",
+                question.as_str()
+            ),
+            (RetainedContextOrder::Local(2), "user", ANSWER),
+        ],
+    );
+    let resumed = resume(&test, &test.codex).await?;
+    assert_eq!(
+        resumed
+            .conversation_history_snapshot()
+            .await
+            .retained_context(),
+        Some(retained),
+    );
+    resumed.shutdown_and_wait().await?;
+    server.shutdown().await;
+    Ok(())
+}
+
 async fn compact_and_assert_answers(
     test: &TestCodex,
     thread: &CodexThread,