From d0221bc9471109ae027082d6fcacbc73b8b8e3aa Mon Sep 17 00:00:00 2001 From: neerajrj-oai Date: Thu, 12 Nov 2026 20:26:45 +0000 Subject: [PATCH] Keep plan item IDs unique after completion (#61186) ## Why Plan items used a single ID derived from the turn ID, so multiple completed plans within one turn shared an ID across retries or tool continuations. ## What changed Generate UUIDv7 plan item IDs and carry an unfinished plan's ID across model requests within a turn. Once a plan completes, assign a fresh ID to the next plan. ## Testing Add unit and integration coverage for unfinished retries, retries after plan completion, and tool continuations, including matching streamed events to saved completions. Update the app-server test to check plan IDs against the start notification. GitOrigin-RevId: 87b4b3e6bd74224ea96f4308f739bcd2efb62299 --- .../app-server/tests/suite/v2/plan_item.rs | 14 +- codex-rs/core/src/session/turn.rs | 29 ++- codex-rs/core/src/session/turn_tests.rs | 66 ++++++- codex-rs/core/tests/suite/items.rs | 182 ++++++++++++++++++ 4 files changed, 277 insertions(+), 14 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/plan_item.rs b/codex-rs/app-server/tests/suite/v2/plan_item.rs index bb2cb65780..b037ce1766 100644 --- a/codex-rs/app-server/tests/suite/v2/plan_item.rs +++ b/codex-rs/app-server/tests/suite/v2/plan_item.rs @@ -56,18 +56,24 @@ async fn plan_mode_uses_proposed_plan_block_for_plan_item() -> Result<()> { .await?; let turn = start_plan_mode_turn(&mut mcp).await?; - let (_, completed_items, plan_deltas, turn_completed) = + let (started_items, completed_items, plan_deltas, turn_completed) = collect_turn_notifications(&mut mcp).await?; wait_for_responses_request_count(&server, /*expected_count*/ 1).await?; assert_eq!(turn_completed.turn.id, turn.id); assert_eq!(turn_completed.turn.status, TurnStatus::Completed); + let expected_plan_id = started_items + .iter() + .find_map(|item| match item { + ThreadItemPayload::Plan { id, .. } => Some(id), + _ => None, + }) + .ok_or_else(|| anyhow!("plan start notification missing"))?; let expected_plan = ThreadItemPayload::Plan { - id: format!("{}-plan", turn.id), + id: expected_plan_id.clone(), text: "# Final plan\n- first\n- second\n".to_string(), }; - let expected_plan_id = format!("{}-plan", turn.id); let streamed_plan = plan_deltas .iter() .map(|delta| delta.delta.as_str()) @@ -76,7 +82,7 @@ async fn plan_mode_uses_proposed_plan_block_for_plan_item() -> Result<()> { assert!( plan_deltas .iter() - .all(|delta| delta.item_id == expected_plan_id) + .all(|delta| &delta.item_id == expected_plan_id) ); let plan_items = completed_items .iter() diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 0d3da231c9..5ea72dcab7 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -156,6 +156,7 @@ use tracing::instrument; use tracing::trace; use tracing::trace_span; use tracing::warn; +use uuid::Uuid; const POST_SAMPLING_TOKEN_ESTIMATE_TARGET: &str = "codex_core::post_sampling_token_estimate"; @@ -540,6 +541,7 @@ pub(crate) async fn run_turn( let mut next_step_context = Some(first_step_context); let mut guardian_budget_compacted = false; + let mut unfinished_plan_id = None; loop { // Note that pending_input would be something like a message the user // submitted through the UI while the model was running. Though the UI @@ -651,6 +653,7 @@ pub(crate) async fn run_turn( Arc::clone(&turn_diff_tracker), &mut client_session, sampling_request_input, + &mut unfinished_plan_id, cancellation_token.child_token(), ) .await @@ -1895,6 +1898,7 @@ async fn run_sampling_request( turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, input: crate::context_manager::ContextManager, + unfinished_plan_id: &mut Option, cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { let turn_context = Arc::clone(&step_context.turn); @@ -2016,6 +2020,7 @@ async fn run_sampling_request( prompt, pending_audio, input_item_count, + unfinished_plan_id, cancellation_token.child_token(), &mut draining_for_mailbox, ) @@ -2259,12 +2264,16 @@ struct PlanModeStreamState { } impl PlanModeStreamState { - fn new(turn_id: &str) -> Self { + fn new(unfinished_plan_id: Option) -> Self { Self { pending_agent_message_items: HashMap::new(), started_agent_message_items: HashSet::new(), leading_whitespace_by_item: HashMap::new(), - plan_item_state: ProposedPlanItemState::new(turn_id), + plan_item_state: ProposedPlanItemState { + item_id: unfinished_plan_id.unwrap_or_else(|| Uuid::now_v7().to_string()), + started: false, + completed: false, + }, } } } @@ -2320,12 +2329,9 @@ impl AssistantMessageStreamParsers { } impl ProposedPlanItemState { - fn new(turn_id: &str) -> Self { - Self { - item_id: format!("{turn_id}-plan"), - started: false, - completed: false, - } + /// Keep the ID across model requests until a completion claims it. + fn unfinished_id(self) -> Option { + (!self.completed).then_some(self.item_id) } async fn start(&mut self, sess: &Session, turn_context: &TurnContext) { @@ -2898,6 +2904,7 @@ async fn try_run_sampling_request( prompt: Prompt, mut pending_audio: Option, input_item_count: usize, + unfinished_plan_id: &mut Option, cancellation_token: CancellationToken, draining_for_mailbox: &mut bool, ) -> CodexResult { @@ -3058,7 +3065,8 @@ async fn try_run_sampling_request( .unwrap_or_else(|| "default".to_string()); let plan_mode = turn_context.mode() == ModeKind::Plan; let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); - let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id)); + let mut plan_mode_state = + plan_mode.then(|| PlanModeStreamState::new(unfinished_plan_id.clone())); let defer_streamed_turn_items_for_contributors = !sess.services.extensions.turn_item_contributors().is_empty(); let mut active_item_is_streaming_to_client = false; @@ -3791,6 +3799,9 @@ async fn try_run_sampling_request( &mut assistant_message_stream_parsers, ) .await; + if let Some(state) = plan_mode_state { + *unfinished_plan_id = state.plan_item_state.unfinished_id(); + } if !in_flight.is_empty() { let _tool_blocking_timing_guard = turn_context diff --git a/codex-rs/core/src/session/turn_tests.rs b/codex-rs/core/src/session/turn_tests.rs index d201607c56..5ba9b18821 100644 --- a/codex-rs/core/src/session/turn_tests.rs +++ b/codex-rs/core/src/session/turn_tests.rs @@ -3,6 +3,8 @@ use codex_extension_api::ExtensionData; use codex_extension_api::TurnItemContributor; use codex_protocol::ResponseItemId; use codex_protocol::items::AgentMessageContent; +use codex_protocol::protocol::ItemCompletedEvent; +use codex_protocol::protocol::ItemStartedEvent; use pretty_assertions::assert_eq; use std::sync::Arc; use tracing_subscriber::prelude::*; @@ -71,7 +73,7 @@ async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() { builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor)); session.services.extensions = Arc::new(builder.build()); let turn_store = ExtensionData::new(turn_context.sub_id.clone()); - let mut state = PlanModeStreamState::new(&turn_context.sub_id); + let mut state = PlanModeStreamState::new(/*unfinished_plan_id*/ None); let mut last_agent_message = None; let item = assistant_output_text("original assistant text"); @@ -95,6 +97,68 @@ async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() { ); } +/// A model retry can reuse an unfinished plan. A later completion must use its own ID. +#[tokio::test] +async fn plan_attempts_only_reuse_unfinished_item_ids() { + let (session, turn, events) = crate::session::tests::make_session_and_context_with_rx().await; + let mut first = PlanModeStreamState::new(/*unfinished_plan_id*/ None).plan_item_state; + let first_id = first.item_id.clone(); + first.start(&session, &turn).await; + first.push_delta(&session, &turn, "Partial").await; + + let mut retry = PlanModeStreamState::new(first.unfinished_id()).plan_item_state; + retry.start(&session, &turn).await; + retry.push_delta(&session, &turn, "Plan A").await; + retry + .complete_with_text(&session, &turn, "Plan A".into()) + .await; + retry + .complete_with_text(&session, &turn, "Repeated A".into()) + .await; + + let mut next = PlanModeStreamState::new(retry.unfinished_id()).plan_item_state; + next.start(&session, &turn).await; + next.push_delta(&session, &turn, "Plan B").await; + next.complete_with_text(&session, &turn, "Plan B".into()) + .await; + + let mut observed = Vec::new(); + while let Ok(event) = events.try_recv() { + match event.msg { + EventMsg::ItemStarted(ItemStartedEvent { + item: TurnItem::Plan(plan), + .. + }) => { + observed.push(("start", plan.id, plan.text)); + } + EventMsg::PlanDelta(event) => observed.push(("delta", event.item_id, event.delta)), + EventMsg::ItemCompleted(ItemCompletedEvent { + item: TurnItem::Plan(plan), + .. + }) => { + observed.push(("complete", plan.id, plan.text)); + } + _ => {} + } + } + let next_id = next.item_id; + assert_eq!(Uuid::parse_str(&first_id).unwrap().get_version_num(), 7); + assert_ne!(first_id, next_id); + assert_eq!( + observed, + vec![ + ("start", first_id.clone(), String::new()), + ("delta", first_id.clone(), "Partial".into()), + ("start", first_id.clone(), String::new()), + ("delta", first_id.clone(), "Plan A".into()), + ("complete", first_id, "Plan A".into()), + ("start", next_id.clone(), String::new()), + ("delta", next_id.clone(), "Plan B".into()), + ("complete", next_id, "Plan B".into()), + ], + ); +} + #[test] fn realtime_user_verification_notice_excludes_request_payload() { let event = EventMsg::ElicitationRequest(codex_protocol::approvals::ElicitationRequestEvent { diff --git a/codex-rs/core/tests/suite/items.rs b/codex-rs/core/tests/suite/items.rs index 034f443be4..980cd325c7 100644 --- a/codex-rs/core/tests/suite/items.rs +++ b/codex-rs/core/tests/suite/items.rs @@ -25,6 +25,7 @@ use codex_protocol::user_input::UserInput; use core_test_support::PathBufExt; 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_message_item_added; use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_reasoning_item; @@ -35,6 +36,7 @@ use core_test_support::responses::ev_response_created; use core_test_support::responses::ev_web_search_call_added_partial; use core_test_support::responses::ev_web_search_call_done; use core_test_support::responses::mount_sse_once; +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; @@ -521,6 +523,186 @@ async fn plan_mode_emits_plan_item_from_proposed_plan_block() -> anyhow::Result< Ok(()) } +#[derive(Clone, Copy)] +enum FirstPlanRequest { + Interrupted, + CompletedBeforeDisconnect, + ToolContinuation, +} + +#[test_case::test_case(FirstPlanRequest::Interrupted; "unfinished_retry")] +#[test_case::test_case(FirstPlanRequest::CompletedBeforeDisconnect; "completed_retry")] +#[test_case::test_case(FirstPlanRequest::ToolContinuation; "tool_continuation")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn plan_mode_keeps_only_unfinished_ids_across_model_requests( + first_request: FirstPlanRequest, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + // The fixture owns the directories used by paginated history. + let test = test_codex() + .with_history_mode(ThreadHistoryMode::Paginated) + .with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(1); + config.model_provider.supports_websockets = false; + config.update_plan_enabled = true; + }) + .build_with_auto_env(&server) + .await?; + let codex = &test.codex; + let session_configured = &test.session_configured; + + let first_text = "\nPlan A\n"; + let first_complete = format!("{first_text}\n"); + let mut first_events = vec![ + ev_response_created("resp-1"), + ev_message_item_added("msg-1", ""), + ev_output_text_delta(first_text), + ]; + if !matches!(first_request, FirstPlanRequest::Interrupted) { + first_events.push(ev_assistant_message("msg-1", &first_complete)); + } + if matches!(first_request, FirstPlanRequest::ToolContinuation) { + first_events.push(ev_function_call( + "plan-tool", + "update_plan", + r#"{"plan":[]}"#, + )); + first_events.push(ev_completed("resp-1")); + } + // The other cases close the stream before response.completed and enter the real retry loop. + let second_text = "\nPlan B\n\n"; + let second_message_id = if matches!(first_request, FirstPlanRequest::Interrupted) { + "msg-1" + } else { + "msg-2" + }; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(first_events), + sse(vec![ + ev_response_created("resp-2"), + ev_message_item_added(second_message_id, ""), + ev_output_text_delta(second_text), + ev_assistant_message(second_message_id, second_text), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + codex + .start_or_steer_turn(disabled_plan_turn( + "please plan", + session_configured.model.clone(), + CollaborationMode { + mode: ModeKind::Plan, + settings: Settings { + model: session_configured.model.clone(), + reasoning_effort: None, + developer_instructions: None, + }, + }, + )?) + .await?; + + let mut started = Vec::new(); + let mut deltas = Vec::new(); + let mut completed = Vec::new(); + let mut retries = 0; + loop { + match wait_for_event(codex, |_| true).await { + EventMsg::ItemStarted(ItemStartedEvent { + item: TurnItem::Plan(item), + turn_id, + .. + }) => started.push((turn_id, item)), + EventMsg::PlanDelta(event) => deltas.push(event), + EventMsg::ItemCompleted(event) if matches!(event.item, TurnItem::Plan(_)) => { + completed.push(event); + } + EventMsg::StreamError(_) => retries += 1, + EventMsg::TurnComplete(event) => { + assert!( + event.error.is_none(), + "unexpected turn error: {:?}", + event.error + ); + break; + } + _ => {} + } + } + + assert_eq!(response_mock.requests().len(), 2); + assert_eq!( + retries, + usize::from(!matches!(first_request, FirstPlanRequest::ToolContinuation)) + ); + assert_eq!(started.len(), 2); + let (turn_id, first_plan) = &started[0]; + let (second_turn_id, second_plan) = &started[1]; + assert_eq!(turn_id, second_turn_id); + assert_eq!( + first_plan.id == second_plan.id, + matches!(first_request, FirstPlanRequest::Interrupted), + ); + assert_eq!( + deltas + .iter() + .map(|event| (&event.turn_id, &event.item_id, event.delta.as_str())) + .collect::>(), + vec![ + (turn_id, &first_plan.id, "Plan A\n"), + (turn_id, &second_plan.id, "Plan B\n") + ], + ); + let mut expected_completions = Vec::new(); + if !matches!(first_request, FirstPlanRequest::Interrupted) { + expected_completions.push((&first_plan.id, "Plan A\n")); + } + expected_completions.push((&second_plan.id, "Plan B\n")); + assert_eq!( + completed + .iter() + .map(|event| { + assert_eq!(&event.turn_id, turn_id); + let TurnItem::Plan(item) = &event.item else { + unreachable!() + }; + (&item.id, item.text.as_str()) + }) + .collect::>(), + expected_completions, + ); + + codex.flush_rollout().await?; + let rollout = std::fs::read_to_string(codex.rollout_path().expect("paginated rollout path"))?; + let saved_completions = rollout + .lines() + .map(codex_rollout::parse_rollout_line) + .collect::, _>>()? + .into_iter() + .filter_map(|line| match line.item { + RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) + if matches!(event.item, TurnItem::Plan(_)) => + { + Some(event) + } + _ => None, + }) + .collect::>(); + assert_eq!( + serde_json::to_value(saved_completions)?, + serde_json::to_value(completed)?, + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn plan_mode_strips_plan_from_agent_messages() -> anyhow::Result<()> { skip_if_no_network!(Ok(()));