From 1a5289a4eff274945f0ab4f2760fa367b6a65270 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Thu, 18 Dec 2025 20:32:51 -0800 Subject: [PATCH] test --- codex-rs/core/src/client.rs | 47 +++------- codex-rs/core/src/client_common.rs | 25 ++++++ codex-rs/core/src/codex.rs | 100 ++++++++------------- codex-rs/core/tests/suite/remote_models.rs | 11 +++ 4 files changed, 83 insertions(+), 100 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index c929e8e861..bd22ab2e56 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1,5 +1,4 @@ use std::sync::Arc; -use std::sync::RwLock; use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; @@ -58,8 +57,8 @@ use crate::tools::spec::create_tools_json_for_responses_api; pub struct ModelClient { config: Arc, auth_manager: Option>, - model_family: RwLock, - models_etag: RwLock>, + model_family: ModelFamily, + models_etag: Option, otel_manager: OtelManager, provider: ModelProviderInfo, conversation_id: ConversationId, @@ -85,8 +84,8 @@ impl ModelClient { Self { config, auth_manager, - model_family: RwLock::new(model_family), - models_etag: RwLock::new(models_etag), + model_family, + models_etag, otel_manager, provider, conversation_id, @@ -112,22 +111,6 @@ impl ModelClient { &self.provider } - pub fn update_models_etag(&self, models_etag: Option) { - let mut guard = self - .models_etag - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - *guard = models_etag; - } - - pub fn update_model_family(&self, model_family: ModelFamily) { - let mut guard = self - .model_family - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - *guard = model_family; - } - /// Streams a single model turn using either the Responses or Chat /// Completions wire API, depending on the configured provider. /// @@ -167,7 +150,7 @@ impl ModelClient { let auth_manager = self.auth_manager.clone(); let model_family = self.get_model_family(); - let instructions = prompt.get_full_instructions(&model_family).into_owned(); + let instructions = prompt.get_full_instructions(model_family).into_owned(); let tools_json = create_tools_json_for_chat_completions_api(&prompt.tools)?; let api_prompt = build_api_prompt(prompt, instructions, tools_json); let conversation_id = self.conversation_id.to_string(); @@ -221,7 +204,7 @@ impl ModelClient { let auth_manager = self.auth_manager.clone(); let model_family = self.get_model_family(); - let instructions = prompt.get_full_instructions(&model_family).into_owned(); + let instructions = prompt.get_full_instructions(model_family).into_owned(); let tools_json: Vec = create_tools_json_for_responses_api(&prompt.tools)?; let reasoning = if model_family.supports_reasoning_summaries { @@ -282,7 +265,7 @@ impl ModelClient { store_override: None, conversation_id: Some(conversation_id.clone()), session_source: Some(session_source.clone()), - extra_headers: beta_feature_headers(&self.config, self.get_models_etag()), + extra_headers: beta_feature_headers(&self.config, self.get_models_etag().clone()), }; let stream_result = client @@ -322,18 +305,12 @@ impl ModelClient { } /// Returns the currently configured model family. - pub fn get_model_family(&self) -> ModelFamily { - self.model_family - .read() - .map(|model_family| model_family.clone()) - .unwrap_or_else(|err| err.into_inner().clone()) + pub fn get_model_family(&self) -> &ModelFamily { + &self.model_family } - fn get_models_etag(&self) -> Option { - self.models_etag - .read() - .map(|models_etag| models_etag.clone()) - .unwrap_or_else(|err| err.into_inner().clone()) + fn get_models_etag(&self) -> &Option { + &self.models_etag } /// Returns the current reasoning effort setting. @@ -370,7 +347,7 @@ impl ModelClient { .with_telemetry(Some(request_telemetry)); let instructions = prompt - .get_full_instructions(&self.get_model_family()) + .get_full_instructions(self.get_model_family()) .into_owned(); let payload = ApiCompactionInput { model: &self.get_model(), diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 4a3bc8de23..ca8a142a3a 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,6 +1,10 @@ use crate::client_common::tools::ToolSpec; +use crate::codex::Session; +use crate::codex::TurnContext; use crate::error::Result; +use crate::features::Feature; use crate::openai_models::model_family::ModelFamily; +use crate::tools::ToolRouter; pub use codex_api::common::ResponseEvent; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; use codex_protocol::models::ResponseItem; @@ -44,6 +48,27 @@ pub struct Prompt { } impl Prompt { + pub(crate) fn new( + sess: &Session, + turn_context: &TurnContext, + router: &ToolRouter, + input: &[ResponseItem], + ) -> Prompt { + let model_supports_parallel = turn_context + .client + .get_model_family() + .supports_parallel_tool_calls; + + Prompt { + input: input.to_vec(), + tools: router.specs(), + parallel_tool_calls: model_supports_parallel + && sess.enabled(Feature::ParallelToolCalls), + base_instructions_override: turn_context.base_instructions.clone(), + output_schema: turn_context.final_output_json_schema.clone(), + } + } + pub(crate) fn get_full_instructions<'a>(&'a self, model: &'a ModelFamily) -> Cow<'a, str> { let base = self .base_instructions_override diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b0b789da4b..cbe80df26b 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2226,7 +2226,7 @@ fn errors_to_info(errors: &[SkillError]) -> Vec { /// pub(crate) async fn run_task( sess: Arc, - turn_context: Arc, + mut turn_context: Arc, input: Vec, cancellation_token: CancellationToken, ) -> Option { @@ -2309,7 +2309,7 @@ pub(crate) async fn run_task( .collect::>(); match run_turn( Arc::clone(&sess), - Arc::clone(&turn_context), + &mut turn_context, Arc::clone(&turn_diff_tracker), turn_input, cancellation_token.child_token(), @@ -2386,7 +2386,7 @@ async fn run_auto_compact(sess: &Arc, turn_context: &Arc) )] async fn run_turn( sess: Arc, - turn_context: Arc, + turn_context: &mut Arc, turn_diff_tracker: SharedTurnDiffTracker, input: Vec, cancellation_token: CancellationToken, @@ -2399,19 +2399,20 @@ async fn run_turn( .list_all_tools() .or_cancel(&cancellation_token) .await?; - let router = Arc::new(ToolRouter::from_config( - &turn_context.tools_config, - Some( - mcp_tools - .into_iter() - .map(|(name, tool)| (name, tool.tool)) - .collect(), - ), - )); let mut retries = 0; loop { - let prompt = build_prompt( + let router = Arc::new(ToolRouter::from_config( + &turn_context.tools_config, + Some( + mcp_tools + .clone() + .into_iter() + .map(|(name, tool)| (name, tool.tool)) + .collect(), + ), + )); + let prompt = Prompt::new( sess.as_ref(), turn_context.as_ref(), router.as_ref(), @@ -2421,7 +2422,7 @@ async fn run_turn( match try_run_turn( Arc::clone(&router), Arc::clone(&sess), - Arc::clone(&turn_context), + Arc::clone(turn_context), Arc::clone(&turn_diff_tracker), &prompt, cancellation_token.child_token(), @@ -2437,13 +2438,13 @@ async fn run_turn( Err(CodexErr::EnvVar(var)) => return Err(CodexErr::EnvVar(var)), Err(e @ CodexErr::Fatal(_)) => return Err(e), Err(e @ CodexErr::ContextWindowExceeded) => { - sess.set_total_tokens_full(&turn_context).await; + sess.set_total_tokens_full(turn_context).await; return Err(e); } Err(CodexErr::UsageLimitReached(e)) => { let rate_limits = e.rate_limits.clone(); if let Some(rate_limits) = rate_limits { - sess.update_rate_limits(&turn_context, rate_limits).await; + sess.update_rate_limits(turn_context, rate_limits).await; } return Err(CodexErr::UsageLimitReached(e)); } @@ -2459,7 +2460,23 @@ async fn run_turn( retries += 1; // Refresh models if we got an outdated models error if matches!(e, CodexErr::OutdatedModels) { - refresh_models_after_outdated_error(sess.as_ref(), turn_context.as_ref()) + let config = { + let state = sess.state.lock().await; + state + .session_configuration + .original_config_do_not_use + .clone() + }; + if let Err(err) = sess + .services + .models_manager + .refresh_available_models(&config) + .await + { + error!("failed to refresh models after outdated models error: {err}"); + } + *turn_context = sess + .new_default_turn_with_sub_id(turn_context.sub_id.clone()) .await; } let delay = match e { @@ -2474,7 +2491,7 @@ async fn run_turn( // user understands what is happening instead of staring // at a seemingly frozen screen. sess.notify_stream_error( - &turn_context, + turn_context, format!("Reconnecting... {retries}/{max_retries}"), e, ) @@ -2489,53 +2506,6 @@ async fn run_turn( } } -fn build_prompt( - sess: &Session, - turn_context: &TurnContext, - router: &ToolRouter, - input: &[ResponseItem], -) -> Prompt { - let model_supports_parallel = turn_context - .client - .get_model_family() - .supports_parallel_tool_calls; - - Prompt { - input: input.to_vec(), - tools: router.specs(), - parallel_tool_calls: model_supports_parallel && sess.enabled(Feature::ParallelToolCalls), - base_instructions_override: turn_context.base_instructions.clone(), - output_schema: turn_context.final_output_json_schema.clone(), - } -} - -async fn refresh_models_after_outdated_error(sess: &Session, turn_context: &TurnContext) { - let config = { - let state = sess.state.lock().await; - state - .session_configuration - .original_config_do_not_use - .clone() - }; - if let Err(err) = sess - .services - .models_manager - .refresh_available_models(&config) - .await - { - error!("failed to refresh models after outdated models error: {err}"); - } - let models_etag = sess.services.models_manager.get_models_etag().await; - let model = turn_context.client.get_model(); - let model_family = sess - .services - .models_manager - .construct_model_family(&model, &config) - .await; - turn_context.client.update_model_family(model_family); - turn_context.client.update_models_etag(models_etag); -} - #[derive(Debug)] struct TurnRunResult { needs_follow_up: bool, diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index d54405fe6d..b03f3ce09f 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -530,6 +530,17 @@ async fn remote_models_refresh_etag_after_outdated_models() -> Result<()> { }) .await?; + let stream_error = + wait_for_event(&codex, |event| matches!(event, EventMsg::StreamError(_))).await; + let EventMsg::StreamError(stream_error) = stream_error else { + unreachable!(); + }; + assert!( + stream_error.message.starts_with("Reconnecting..."), + "unexpected stream error message: {}", + stream_error.message + ); + wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; // Phase 3c: assert the refresh happened and the ETag was updated.