diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 807705667f..c7ba315ac8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2985,6 +2985,214 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> Ok(()) } +#[tokio::test] +async fn thread_goal_keeps_original_root_until_external_objective_edit() -> Result<()> { + let (release_original_turn, original_turn_gate) = oneshot::channel(); + let (release_edited_turn, edited_turn_gate) = oneshot::channel(); + let (server, _response_completions) = start_streaming_sse_server(vec![ + ungated_goal_response(responses::sse(vec![ + responses::ev_response_created("create-original-goal"), + responses::ev_function_call( + "create-original-goal-call", + "create_goal", + r#"{"objective":"keep its original owner","token_budget":100}"#, + ), + responses::ev_completed_with_tokens("create-original-goal", /*total_tokens*/ 5), + ])), + vec![StreamingSseChunk { + gate: Some(original_turn_gate), + body: responses::sse_completed("finish-original-user-turn"), + }], + ungated_goal_response(responses::sse_completed("reopen-original-user-turn")), + ungated_goal_response(responses::sse_completed("finish-intervening-user-turn")), + ungated_goal_response(responses::sse(vec![ + responses::ev_response_created("goal-continuation-after-intervening-turn"), + responses::ev_completed_with_tokens( + "goal-continuation-after-intervening-turn", + /*total_tokens*/ 40, + ), + ])), + vec![StreamingSseChunk { + gate: Some(edited_turn_gate), + body: responses::sse_completed("second-goal-continuation"), + }], + ungated_goal_response(responses::sse_completed("reopened-goal-turn")), + ungated_goal_response(responses::sse(vec![ + responses::ev_response_created("rootless-goal-continuation"), + responses::ev_completed_with_tokens( + "rootless-goal-continuation", + /*total_tokens*/ 100, + ), + ])), + ]) + .await; + let codex_home = TempDir::new()?; + mock_responses_config(server.uri()) + .enable_feature(Feature::Goals) + .write(codex_home.path())?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_managed_config() + .build_initialized() + .await?; + let thread = mcp.start_thread(ThreadStartParams::default()).await?.thread; + + let start_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "create the original goal".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let original_turn: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; + timeout( + DEFAULT_READ_TIMEOUT, + server.wait_for_request_count(/*count*/ 2), + ) + .await?; + + let injection_id = mcp + .send_raw_request( + "thread/inject_items", + Some(json!({ + "threadId": thread.id, + "items": [{ + "type": "message", + "role": "developer", + "content": [{ + "type": "input_text", + "text": "externally injected context", + }], + }], + })), + ) + .await?; + let _: serde_json::Value = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(injection_id)).await??; + + let queue_id = mcp + .send_raw_request( + "thread/queue/add", + Some(json!({ + "threadId": thread.id, + "input": [{ + "type": "text", + "text": "an intervening user message", + "textElements": [], + }], + "clientUserMessageId": "intervening-goal-message", + })), + ) + .await?; + let _: serde_json::Value = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(queue_id)).await??; + release_original_turn + .send(()) + .expect("original turn should remain open until the user message is queued"); + + for _ in 0..3 { + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + } + timeout( + DEFAULT_READ_TIMEOUT, + server.wait_for_request_count(/*count*/ 6), + ) + .await?; + + let edit_id = mcp + .send_raw_request( + "thread/goal/set", + Some(json!({ + "threadId": thread.id, + "objective": "externally updated goal", + "status": "active", + })), + ) + .await?; + let edited_goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(edit_id)).await??; + assert_eq!(edited_goal.goal.objective, "externally updated goal"); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/goal/updated"), + ) + .await??; + + let get_id = mcp + .send_raw_request("thread/goal/get", Some(json!({ "threadId": thread.id }))) + .await?; + let _: codex_app_server_protocol::ThreadGoalGetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; + release_edited_turn + .send(()) + .expect("goal turn should remain open until its external edit"); + + for _ in 0..2 { + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + } + + let requests = server.requests().await; + assert_eq!(requests.len(), 8); + let reopened_original_request = serde_json::from_slice::(&requests[2])?; + assert_eq!( + reopened_original_request["client_metadata"]["turn_id"].as_str(), + Some(original_turn.turn.id.as_str()) + ); + responses::assert_root_turn(&reopened_original_request, /*expected*/ None)?; + let intervening_request = serde_json::from_slice::(&requests[3])?; + let intervening_turn_id = intervening_request["client_metadata"]["turn_id"] + .as_str() + .expect("intervening user turn ID"); + responses::assert_root_turn(&intervening_request, Some(intervening_turn_id))?; + let first_continuation = serde_json::from_slice::(&requests[4])?; + let first_continuation_turn_id = first_continuation["client_metadata"]["turn_id"] + .as_str() + .expect("first continuation turn ID"); + let second_continuation = serde_json::from_slice::(&requests[5])?; + for (request, parent_turn_id) in [ + (&first_continuation, intervening_turn_id), + (&second_continuation, first_continuation_turn_id), + ] { + responses::assert_root_turn(request, Some(original_turn.turn.id.as_str()))?; + responses::assert_parent_turn(request, Some(parent_turn_id))?; + } + let edited_turn_id = second_continuation["client_metadata"]["turn_id"] + .as_str() + .expect("second continuation turn ID"); + + let reopened_request = serde_json::from_slice::(&requests[6])?; + assert_eq!( + reopened_request["client_metadata"]["turn_id"].as_str(), + Some(edited_turn_id) + ); + responses::assert_root_turn(&reopened_request, /*expected*/ None)?; + let continuation_request = serde_json::from_slice::(&requests[7])?; + assert_ne!( + continuation_request["client_metadata"]["turn_id"].as_str(), + Some(edited_turn_id) + ); + responses::assert_root_turn(&continuation_request, /*expected*/ None)?; + responses::assert_parent_turn(&continuation_request, /*expected*/ None)?; + + server.shutdown().await; + Ok(()) +} + +fn ungated_goal_response(body: String) -> Vec { + vec![StreamingSseChunk { gate: None, body }] +} + #[tokio::test] async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(vec![ @@ -3131,6 +3339,20 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul serde_json::Value::Null ); + let requests = server.received_requests().await.expect("wiremock requests"); + let goal_request = requests + .iter() + .filter(|request| request.url.path().ends_with("/responses")) + .nth(1) + .expect("externally created goal continuation request"); + let goal_request_body = goal_request.body_json::()?; + assert_eq!( + goal_request_body["client_metadata"]["turn_id"], + causal_turn_id + ); + responses::assert_root_turn(&goal_request_body, /*expected*/ None)?; + responses::assert_parent_turn(&goal_request_body, /*expected*/ None)?; + let clear_id = mcp .send_raw_request( "thread/goal/clear", diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 3527df18d5..533d1351c2 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -510,6 +510,28 @@ impl CodexThread { self.session.inject_if_running(items).await } + /// Returns the trusted root when the expected turn is currently active. + pub async fn active_turn_root(&self, expected_turn_id: &str) -> Option { + let active = self.session.active_turn.lock().await; + let task = active.as_ref()?.task.as_ref()?; + if task.turn_context.sub_id != expected_turn_id { + return None; + } + task.turn_context.turn_metadata_state.root_turn_id() + } + + /// Invalidates the trusted root when the expected turn is currently active. + pub async fn invalidate_turn_lineage(&self, expected_turn_id: &str) { + let active = self.session.active_turn.lock().await; + if let Some(task) = active.as_ref().and_then(|turn| turn.task.as_ref()) + && task.turn_context.sub_id == expected_turn_id + { + task.turn_context + .turn_metadata_state + .mark_root_turn_ambiguous(); + } + } + pub async fn set_app_server_client_info( &self, app_server_client_name: Option, diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index e8202d38f5..f3a0ec11c9 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -769,10 +769,21 @@ pub(crate) async fn drain_async_hook_results( .collect::>(); if before_user_prompt { + // A fresh user turn owns its root; automatic turns only inherit one. + let current_turn_id = Some(turn_context.sub_id.as_str()); + if !additional_contexts.is_empty() + && result.turn_id.as_deref() != current_turn_id + && turn_context.turn_metadata_state.root_turn_id().as_deref() != current_turn_id + { + turn_context.turn_metadata_state.mark_root_turn_ambiguous(); + } record_additional_contexts(sess, turn_context, additional_contexts).await; } else if !additional_contexts.is_empty() { let _ = sess - .inject_if_running(additional_context_messages(additional_contexts)) + .inject_hook_context_if_running( + additional_context_messages(additional_contexts), + result.turn_id.as_deref(), + ) .await; } diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 8c17a8ae5e..3c57e7033f 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -35,6 +35,41 @@ impl Session { } } + /// Injects hook context while classifying its actual receiving turn atomically. + #[expect( + clippy::await_holding_invalid_type, + reason = "active turn provenance and turn state updates must remain atomic" + )] + pub(crate) async fn inject_hook_context_if_running( + &self, + input: Vec, + source_turn_id: Option<&str>, + ) -> Result<(), Vec> { + let mut active = self.active_turn.lock().await; + let Some(active_turn) = active.as_mut() else { + return Err(input); + }; + let Some(task) = active_turn.task.as_ref() else { + return Err(input); + }; + if source_turn_id != Some(task.turn_context.sub_id.as_str()) { + task.turn_context + .turn_metadata_state + .mark_root_turn_ambiguous(); + } + self.input_queue + .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + active_turn.turn_state.as_ref(), + input + .into_iter() + .map(ResponseItemEnvelope::new) + .map(PendingTurnInput::ResponseItem) + .collect(), + ) + .await; + Ok(()) + } + /// Preserves trusted client provenance while items wait for an active turn. #[expect( clippy::await_holding_invalid_type, @@ -51,6 +86,11 @@ impl Session { .collect::>(); let mut active = self.active_turn.lock().await; if let Some(active_turn) = active.as_mut() { + if let Some(task) = active_turn.task.as_ref() { + task.turn_context + .turn_metadata_state + .mark_root_turn_ambiguous(); + } self.input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( active_turn.turn_state.as_ref(), diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 1e88b17147..e4874b489c 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -5,7 +5,9 @@ use anyhow::Context; use anyhow::Result; use codex_config::test_support::CloudConfigBundleFixture; use codex_core::StartThreadOptions; +use codex_core::TurnInput; use codex_core::TurnInputRequest; +use codex_core::TurnStartOptions; use codex_core::config::Config; use codex_core::config::Constrained; use codex_core::config::ThreadStoreConfig; @@ -39,6 +41,7 @@ use core_test_support::fs_wait; use core_test_support::hooks::trust_discovered_hooks; use core_test_support::hooks::trust_hooks; use core_test_support::managed_network_requirements_loader; +use core_test_support::responses; use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -1704,6 +1707,11 @@ async fn async_hook_context_is_injected_into_the_active_turn() -> Result<()> { let requests = responses.requests(); assert_eq!(requests.len(), 2); + let first_request = requests[0].body_json(); + let turn_id = first_request["client_metadata"]["turn_id"] + .as_str() + .context("first model request should include its turn ID")?; + responses::assert_root_turn(&requests[1].body_json(), Some(turn_id))?; assert!( requests[1] .message_input_texts("developer") @@ -1714,8 +1722,12 @@ async fn async_hook_context_is_injected_into_the_active_turn() -> Result<()> { Ok(()) } +#[test_case::test_case(/*automatic_continuation*/ false; "user_turn")] +#[test_case::test_case(/*automatic_continuation*/ true; "automatic_continuation")] #[tokio::test] -async fn async_hook_finishing_while_idle_waits_for_the_next_turn() -> Result<()> { +async fn async_hook_finishing_while_idle_waits_for_the_next_turn( + automatic_continuation: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -1785,12 +1797,22 @@ async fn async_hook_finishing_while_idle_waits_for_the_next_turn() -> Result<()> ); let next_prompt = "observe the buffered async context"; - test.codex - .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + let next_turn = if automatic_continuation { + TurnInputRequest::new(TurnInput::ResponseItem(responses::user_message_item( + next_prompt, + ))) + .on_start(TurnStartOptions { + root_turn_id: Some(first_turn_id.clone()), + parent_turn_id: Some(first_turn_id.clone()), + ..Default::default() + }) + } else { + TurnInputRequest::user_input(vec![UserInput::Text { text: next_prompt.to_string(), text_elements: Vec::new(), - }])) - .await?; + }]) + }; + test.codex.start_turn_if_idle(next_turn).await?; let mut warning_event = None; timeout(Duration::from_secs(5), async { @@ -1819,6 +1841,10 @@ async fn async_hook_finishing_while_idle_waits_for_the_next_turn() -> Result<()> .context("second model request should include its turn ID")? .to_string(); assert_ne!(first_turn_id, second_turn_id); + responses::assert_root_turn( + &requests[1].body_json(), + (!automatic_continuation).then_some(second_turn_id.as_str()), + )?; assert_eq!( warning_event .context("buffered async hook warning should be delivered during the next turn")? diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index aa9b2f7918..440ba1d42d 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -1043,6 +1043,7 @@ async fn queued_inter_agent_mail_does_not_restart_after_final_answer() { async fn injected_response_item_reopens_turn_after_final_answer() { const INITIAL_PROMPT: &str = "first prompt"; const INJECTED_CONTEXT: &str = "late injected context"; + const EXTERNAL_CONTEXT: &str = "external injected context"; let (gate_completed_tx, gate_completed_rx) = oneshot::channel(); let first_chunks = vec![ @@ -1083,18 +1084,35 @@ async fn injected_response_item_reopens_turn_after_final_answer() { .await .is_ok() ); + codex + .inject_response_items(vec![responses::user_message_item(EXTERNAL_CONTEXT)]) + .await + .expect("external context should be injected"); let _ = gate_completed_tx.send(()); wait_for_turn_complete(&codex).await; let requests = server.requests().await; assert_eq!(requests.len(), 2); + let first: Value = from_slice(&requests[0]).expect("parse first request"); + let first_turn_id = first["client_metadata"]["turn_id"] + .as_str() + .expect("first request should include its turn ID"); + responses::assert_root_turn(&first, Some(first_turn_id)) + .expect("initial root should be trusted"); let second: Value = from_slice(&requests[1]).expect("parse second request"); + responses::assert_root_turn(&second, /*expected*/ None) + .expect("external injection should invalidate the active turn root"); let relevant_user_input = message_input_texts(&second, "user") .into_iter() - .filter(|text| text == INITIAL_PROMPT || text == INJECTED_CONTEXT) + .filter(|text| { + text == INITIAL_PROMPT || text == INJECTED_CONTEXT || text == EXTERNAL_CONTEXT + }) .collect::>(); - assert_eq!(relevant_user_input, vec![INITIAL_PROMPT, INJECTED_CONTEXT]); + assert_eq!( + relevant_user_input, + vec![INITIAL_PROMPT, INJECTED_CONTEXT, EXTERNAL_CONTEXT] + ); server.shutdown().await; } diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 94397d55b4..a643b90ff2 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/codex-rs/ext/goal/src/api.rs @@ -275,6 +275,10 @@ impl GoalService { .map(|goal| (goal, Some(previous_goal)))? }; + if let Some(runtime) = runtime.as_ref() { + runtime.invalidate_turn_lineage().await; + } + if objective.is_some() { fill_empty_thread_preview_if_possible(state_db, thread_id, &goal).await; } @@ -316,6 +320,9 @@ impl GoalService { GoalServiceError::Internal(format!("failed to clear thread goal: {err}")) })?; let cleared = cleared_goal.is_some(); + if cleared && let Some(runtime) = runtime.as_ref() { + runtime.invalidate_turn_lineage().await; + } drop(goal_state_permit); drop(runtime); diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 0e8d9e7522..89848996fa 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -3,6 +3,7 @@ use std::sync::Weak; use codex_analytics::AnalyticsEventsClient; use codex_core::ThreadManager; +use codex_core::TurnStartOptions; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; @@ -41,6 +42,7 @@ use crate::metrics::GoalMetrics; use crate::runtime::ActiveGoalStopReason; use crate::runtime::GoalRuntimeConfig; use crate::runtime::GoalRuntimeHandle; +use crate::spec::CREATE_GOAL_TOOL_NAME; use crate::spec::UPDATE_GOAL_TOOL_NAME; use crate::steering::budget_limit_steering_item; use crate::tool::GoalToolExecutor; @@ -288,6 +290,7 @@ where ) .await { + input.thread_store.remove::(); tracing::warn!( "failed to stop active goal after repeated execution failures for {turn_id}: {err}" ); @@ -302,12 +305,27 @@ where ) .await { + input.thread_store.remove::(); tracing::warn!( "failed to account active goal progress at turn stop for {turn_id}: {err}" ); return; } - runtime.accounting_state().finish_turn(turn_id); + let accounting = runtime.accounting_state(); + if accounting + .current_active_goal_id_for_turn(turn_id) + .is_some() + && let Some(options) = input.thread_store.get::() + { + input.thread_store.insert_if( + TurnStartOptions { + parent_turn_id: Some(turn_id.to_string()), + ..options.as_ref().clone() + }, + |current| current.is_some(), + ); + } + accounting.finish_turn(turn_id); }) } @@ -321,6 +339,7 @@ where } let turn_id = input.turn_store.level_id(); + input.thread_store.remove::(); if let Err(err) = runtime .account_active_goal_progress( turn_id, @@ -412,6 +431,23 @@ where input.tool_name, input.outcome, ); + if input.tool_name.is_default_namespace() + && input.tool_name.name == CREATE_GOAL_TOOL_NAME + && matches!(input.outcome, ToolCallOutcome::Completed { success: true }) + { + input.thread_store.remove::(); + if let Ok(_goal_state_permit) = runtime.goal_state_permit().await + && let Some(thread_manager) = self.thread_manager.upgrade() + && let Ok(thread) = thread_manager.get_thread(runtime.thread_id()).await + && let Some(root_turn_id) = thread.active_turn_root(input.turn_id).await + { + input.thread_store.insert(TurnStartOptions { + root_turn_id: Some(root_turn_id), + parent_turn_id: Some(input.turn_id.to_string()), + ..Default::default() + }); + } + } let should_count_for_goal_progress = tool_attempt_counts_for_goal_progress(input.outcome) && !(input.tool_name.is_default_namespace() diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 027378f837..dca0736c91 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -135,6 +135,19 @@ impl GoalRuntimeHandle { self.inner.root_accounting_state.clone() } + pub(crate) async fn invalidate_turn_lineage(&self) { + let Some(thread_manager) = self.inner.thread_manager.upgrade() else { + return; + }; + let Ok(thread) = thread_manager.get_thread(self.inner.thread_id).await else { + return; + }; + thread.thread_extension_data().remove::(); + if let Some(turn_id) = self.inner.accounting_state.current_turn_id() { + thread.invalidate_turn_lineage(turn_id.as_str()).await; + } + } + pub(crate) async fn goal_state_permit(&self) -> Result, String> { self.inner .goal_state_lock @@ -430,13 +443,18 @@ impl GoalRuntimeHandle { self.inner.accounting_state.clear_active_goal(); return Ok(()); } + let start_options = thread + .thread_extension_data() + .get::() + .map(|options| options.as_ref().clone()) + .unwrap_or_default(); let item = continuation_steering_item(&protocol_goal_from_state(goal)); match thread .start_turn_if_idle( TurnInputRequest::new(TurnInput::ResponseItem(item)).on_start(TurnStartOptions { turn_trigger: Some("goal".to_string()), - ..Default::default() + ..start_options }), ) .await