From 45d05f2b06f0141f54cc082e65850c9cf137e97f Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Mon, 8 Jun 2026 09:43:37 -0700 Subject: [PATCH] Add core integration coverage for goal startup --- codex-rs/core/src/session/initial_goal.rs | 2 +- codex-rs/core/src/session/turn_context.rs | 8 +- codex-rs/core/tests/suite/initial_goal.rs | 198 ++++++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + 4 files changed, 201 insertions(+), 8 deletions(-) create mode 100644 codex-rs/core/tests/suite/initial_goal.rs diff --git a/codex-rs/core/src/session/initial_goal.rs b/codex-rs/core/src/session/initial_goal.rs index 1a7ada8b89..b9636cb184 100644 --- a/codex-rs/core/src/session/initial_goal.rs +++ b/codex-rs/core/src/session/initial_goal.rs @@ -72,7 +72,7 @@ impl Session { .replace_for_turn(InitialGoalInput { turn_id, goal, - collaboration_mode: prepared_turn.collaboration_mode(), + collaboration_mode: &prepared_turn.session_configuration.collaboration_mode, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, }) diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 7bf666c405..d4ee339c2a 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -112,7 +112,7 @@ enum TurnMultiAgentRuntime { } pub(crate) struct PreparedTurn { - session_configuration: SessionConfiguration, + pub(super) session_configuration: SessionConfiguration, turn_environments: ResolvedTurnEnvironments, permission_profile_changed: bool, previous_cwd: AbsolutePathBuf, @@ -123,12 +123,6 @@ pub(crate) struct PreparedTurn { final_output_json_schema: Option>, } -impl PreparedTurn { - pub(crate) fn collaboration_mode(&self) -> &CollaborationMode { - &self.session_configuration.collaboration_mode - } -} - impl TurnContext { pub(crate) fn permission_profile(&self) -> PermissionProfile { self.permission_profile.clone() diff --git a/codex-rs/core/tests/suite/initial_goal.rs b/codex-rs/core/tests/suite/initial_goal.rs new file mode 100644 index 0000000000..4d9f33cd17 --- /dev/null +++ b/codex-rs/core/tests/suite/initial_goal.rs @@ -0,0 +1,198 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; + +use anyhow::Result; +use codex_core::config::Config; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::InitialGoalContributor; +use codex_extension_api::InitialGoalError; +use codex_extension_api::InitialGoalInput; +use codex_protocol::error::CodexErr; +use codex_protocol::protocol::AgentStatus; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InitialGoal; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event_match; +use pretty_assertions::assert_eq; +use serde_json::Value; +use tokio::sync::oneshot; + +#[derive(Clone, Debug, PartialEq, Eq)] +struct RecordedInitialGoal { + turn_id: String, + objective: String, +} + +#[derive(Default)] +struct RecordingInitialGoalContributor { + calls: Mutex>, +} + +impl RecordingInitialGoalContributor { + fn calls(&self) -> Vec { + self.calls + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } +} + +impl InitialGoalContributor for RecordingInitialGoalContributor { + fn replace_for_turn<'a>( + &'a self, + input: InitialGoalInput<'a>, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + self.calls + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(RecordedInitialGoal { + turn_id: input.turn_id.to_string(), + objective: input.goal.objective.clone(), + }); + Ok(()) + }) + } +} + +fn user_input(text: &str) -> Op { + Op::UserInput { + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + environments: None, + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn initial_goal_starts_one_turn_and_rejects_a_concurrent_goal() -> Result<()> { + let (completion_gate_tx, completion_gate_rx) = oneshot::channel(); + let (server, _) = start_streaming_sse_server(vec![vec![ + StreamingSseChunk { + gate: None, + body: responses::sse(vec![ev_response_created("response-1")]), + }, + StreamingSseChunk { + gate: Some(completion_gate_rx), + body: responses::sse(vec![ + ev_assistant_message("message-1", "Initial pass complete."), + ev_completed("response-1"), + ]), + }, + ]]) + .await; + + let contributor = Arc::new(RecordingInitialGoalContributor::default()); + let mut extension_builder = ExtensionRegistryBuilder::::new(); + extension_builder.initial_goal_contributor(contributor.clone()); + let mut builder = test_codex() + .with_model("gpt-5.4") + .with_extensions(Arc::new(extension_builder.build())); + let test = builder.build_with_streaming_server(&server).await?; + + let first_turn_id = test + .codex + .submit_user_input_with_client_user_message_id( + user_input("Improve benchmark coverage"), + /*trace*/ None, + /*client_user_message_id*/ None, + Some(InitialGoal { + objective: "Improve benchmark coverage".to_string(), + }), + ) + .await?; + let running_status = test.codex.agent_status().await; + let started_turn_id = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnStarted(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; + server.wait_for_request_count(1).await; + + let second_error = match test + .codex + .submit_user_input_with_client_user_message_id( + user_input("Replace the active goal"), + /*trace*/ None, + /*client_user_message_id*/ None, + Some(InitialGoal { + objective: "Replace the active goal".to_string(), + }), + ) + .await + { + Err(CodexErr::InvalidRequest(message)) => message, + Err(err) => anyhow::bail!("expected invalid request, got {err}"), + Ok(turn_id) => anyhow::bail!("concurrent goal unexpectedly started turn {turn_id}"), + }; + + completion_gate_tx + .send(()) + .map_err(|()| anyhow::anyhow!("response completion gate closed"))?; + let completed_turn_id = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnComplete(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; + let requests = server.requests().await; + let request_body: Value = serde_json::from_slice( + requests + .first() + .ok_or_else(|| anyhow::anyhow!("expected initial model request"))?, + )?; + let original_user_texts = request_body + .get("input") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter(|item| item.get("type").and_then(Value::as_str) == Some("message")) + .filter(|item| item.get("role").and_then(Value::as_str) == Some("user")) + .filter_map(|item| item.get("content").and_then(Value::as_array)) + .flatten() + .filter_map(|content| content.get("text").and_then(Value::as_str)) + .filter(|text| !text.starts_with("")) + .map(str::to_string) + .collect::>(); + + assert_eq!( + ( + contributor.calls(), + running_status, + started_turn_id, + completed_turn_id, + second_error, + requests.len(), + original_user_texts, + ), + ( + vec![RecordedInitialGoal { + turn_id: first_turn_id.clone(), + objective: "Improve benchmark coverage".to_string(), + }], + AgentStatus::Running, + first_turn_id.clone(), + first_turn_id, + "cannot start a goal while another turn is active".to_string(), + 1, + vec!["Improve benchmark coverage".to_string()], + ) + ); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 234fb8a2f0..52d5412336 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -60,6 +60,7 @@ mod hooks; #[cfg(not(target_os = "windows"))] mod hooks_mcp; mod image_rollout; +mod initial_goal; mod items; mod json_result; mod live_cli;