diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ce091297a6..b31ab6e054 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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::(); 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()) + }); } } } diff --git a/codex-rs/core/src/session/retained_context.rs b/codex-rs/core/src/session/retained_context.rs index 8d1eb3e615..72f22c7367 100644 --- a/codex-rs/core/src/session/retained_context.rs +++ b/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>); + 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 { if self.guardian_context_mode == GuardianContextMode::Legacy { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8f96220f62..cdb61e7597 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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, diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index a89923dbea..1e21a4cec1 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/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 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}\n"), "\n") + } 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![ + (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,