Skip to content

Commit 1e1146c

Browse files
authored
Reload cloud requirements after user login (openai#10725)
Reload cloud requirements after user login so it could take effect immediately.
1 parent dc7007b commit 1e1146c

2 files changed

Lines changed: 114 additions & 14 deletions

File tree

codex-rs/app-server/src/codex_message_processor.rs

Lines changed: 97 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ use codex_app_server_protocol::UserSavedConfig;
140140
use codex_app_server_protocol::build_turns_from_event_msgs;
141141
use codex_backend_client::Client as BackendClient;
142142
use codex_chatgpt::connectors;
143+
use codex_cloud_requirements::cloud_requirements_loader;
143144
use codex_core::AuthManager;
144145
use codex_core::CodexAuth;
145146
use codex_core::CodexThread;
@@ -162,6 +163,7 @@ use codex_core::config::edit::ConfigEditsBuilder;
162163
use codex_core::config::types::McpServerTransportConfig;
163164
use codex_core::config_loader::CloudRequirementsLoader;
164165
use codex_core::default_client::get_codex_user_agent;
166+
use codex_core::default_client::set_default_client_residency_requirement;
165167
use codex_core::error::CodexErr;
166168
use codex_core::exec::ExecParams;
167169
use codex_core::exec_env::create_env;
@@ -218,6 +220,7 @@ use std::io::Error as IoError;
218220
use std::path::Path;
219221
use std::path::PathBuf;
220222
use std::sync::Arc;
223+
use std::sync::RwLock;
221224
use std::sync::atomic::AtomicBool;
222225
use std::sync::atomic::Ordering;
223226
use std::time::Duration;
@@ -277,7 +280,7 @@ pub(crate) struct CodexMessageProcessor {
277280
codex_linux_sandbox_exe: Option<PathBuf>,
278281
config: Arc<Config>,
279282
cli_overrides: Vec<(String, TomlValue)>,
280-
cloud_requirements: CloudRequirementsLoader,
283+
cloud_requirements: Arc<RwLock<CloudRequirementsLoader>>,
281284
conversation_listeners: HashMap<Uuid, oneshot::Sender<()>>,
282285
listener_thread_ids_by_subscription: HashMap<Uuid, ThreadId>,
283286
active_login: Arc<Mutex<Option<ActiveLogin>>>,
@@ -349,7 +352,7 @@ impl CodexMessageProcessor {
349352
codex_linux_sandbox_exe,
350353
config,
351354
cli_overrides,
352-
cloud_requirements,
355+
cloud_requirements: Arc::new(RwLock::new(cloud_requirements)),
353356
conversation_listeners: HashMap::new(),
354357
listener_thread_ids_by_subscription: HashMap::new(),
355358
active_login: Arc::new(Mutex::new(None)),
@@ -362,9 +365,10 @@ impl CodexMessageProcessor {
362365
}
363366

364367
async fn load_latest_config(&self) -> Result<Config, JSONRPCErrorError> {
368+
let cloud_requirements = self.current_cloud_requirements();
365369
codex_core::config::ConfigBuilder::default()
366370
.cli_overrides(self.cli_overrides.clone())
367-
.cloud_requirements(self.cloud_requirements.clone())
371+
.cloud_requirements(cloud_requirements)
368372
.build()
369373
.await
370374
.map_err(|err| JSONRPCErrorError {
@@ -374,6 +378,13 @@ impl CodexMessageProcessor {
374378
})
375379
}
376380

381+
fn current_cloud_requirements(&self) -> CloudRequirementsLoader {
382+
self.cloud_requirements
383+
.read()
384+
.map(|guard| guard.clone())
385+
.unwrap_or_default()
386+
}
387+
377388
fn review_request_from_target(
378389
target: ApiReviewTarget,
379390
) -> Result<(ReviewRequest, String), JSONRPCErrorError> {
@@ -834,6 +845,9 @@ impl CodexMessageProcessor {
834845
let outgoing_clone = self.outgoing.clone();
835846
let active_login = self.active_login.clone();
836847
let auth_manager = self.auth_manager.clone();
848+
let cloud_requirements = self.cloud_requirements.clone();
849+
let chatgpt_base_url = self.config.chatgpt_base_url.clone();
850+
let cli_overrides = self.cli_overrides.clone();
837851
let auth_url = server.auth_url.clone();
838852
tokio::spawn(async move {
839853
let (success, error_msg) = match tokio::time::timeout(
@@ -863,6 +877,16 @@ impl CodexMessageProcessor {
863877

864878
if success {
865879
auth_manager.reload();
880+
replace_cloud_requirements_loader(
881+
cloud_requirements.as_ref(),
882+
auth_manager.clone(),
883+
chatgpt_base_url,
884+
);
885+
sync_default_client_residency_requirement(
886+
&cli_overrides,
887+
cloud_requirements.as_ref(),
888+
)
889+
.await;
866890

867891
// Notify clients with the actual current auth mode.
868892
let current_auth_method = auth_manager
@@ -927,6 +951,9 @@ impl CodexMessageProcessor {
927951
let outgoing_clone = self.outgoing.clone();
928952
let active_login = self.active_login.clone();
929953
let auth_manager = self.auth_manager.clone();
954+
let cloud_requirements = self.cloud_requirements.clone();
955+
let chatgpt_base_url = self.config.chatgpt_base_url.clone();
956+
let cli_overrides = self.cli_overrides.clone();
930957
let auth_url = server.auth_url.clone();
931958
tokio::spawn(async move {
932959
let (success, error_msg) = match tokio::time::timeout(
@@ -956,6 +983,16 @@ impl CodexMessageProcessor {
956983

957984
if success {
958985
auth_manager.reload();
986+
replace_cloud_requirements_loader(
987+
cloud_requirements.as_ref(),
988+
auth_manager.clone(),
989+
chatgpt_base_url,
990+
);
991+
sync_default_client_residency_requirement(
992+
&cli_overrides,
993+
cloud_requirements.as_ref(),
994+
)
995+
.await;
959996

960997
// Notify clients with the actual current auth mode.
961998
let current_auth_method = auth_manager
@@ -1123,6 +1160,16 @@ impl CodexMessageProcessor {
11231160
return;
11241161
}
11251162
self.auth_manager.reload();
1163+
replace_cloud_requirements_loader(
1164+
self.cloud_requirements.as_ref(),
1165+
self.auth_manager.clone(),
1166+
self.config.chatgpt_base_url.clone(),
1167+
);
1168+
sync_default_client_residency_requirement(
1169+
&self.cli_overrides,
1170+
self.cloud_requirements.as_ref(),
1171+
)
1172+
.await;
11261173

11271174
self.outgoing
11281175
.send_response(request_id, LoginAccountResponse::ChatgptAuthTokens {})
@@ -1561,11 +1608,12 @@ impl CodexMessageProcessor {
15611608
);
15621609
}
15631610

1611+
let cloud_requirements = self.current_cloud_requirements();
15641612
let config = match derive_config_from_params(
15651613
&self.cli_overrides,
15661614
Some(request_overrides),
15671615
typesafe_overrides,
1568-
&self.cloud_requirements,
1616+
&cloud_requirements,
15691617
)
15701618
.await
15711619
{
@@ -1647,11 +1695,12 @@ impl CodexMessageProcessor {
16471695
);
16481696
typesafe_overrides.ephemeral = ephemeral;
16491697

1698+
let cloud_requirements = self.current_cloud_requirements();
16501699
let config = match derive_config_from_params(
16511700
&self.cli_overrides,
16521701
config,
16531702
typesafe_overrides,
1654-
&self.cloud_requirements,
1703+
&cloud_requirements,
16551704
)
16561705
.await
16571706
{
@@ -2440,12 +2489,13 @@ impl CodexMessageProcessor {
24402489
);
24412490

24422491
// Derive a Config using the same logic as new conversation, honoring overrides if provided.
2492+
let cloud_requirements = self.current_cloud_requirements();
24432493
let config = match derive_config_for_cwd(
24442494
&self.cli_overrides,
24452495
request_overrides,
24462496
typesafe_overrides,
24472497
history_cwd,
2448-
&self.cloud_requirements,
2498+
&cloud_requirements,
24492499
)
24502500
.await
24512501
{
@@ -2633,12 +2683,13 @@ impl CodexMessageProcessor {
26332683
None,
26342684
);
26352685
// Derive a Config using the same logic as new conversation, honoring overrides if provided.
2686+
let cloud_requirements = self.current_cloud_requirements();
26362687
let config = match derive_config_for_cwd(
26372688
&self.cli_overrides,
26382689
request_overrides,
26392690
typesafe_overrides,
26402691
history_cwd,
2641-
&self.cloud_requirements,
2692+
&cloud_requirements,
26422693
)
26432694
.await
26442695
{
@@ -3438,12 +3489,13 @@ impl CodexMessageProcessor {
34383489
),
34393490
};
34403491

3492+
let cloud_requirements = self.current_cloud_requirements();
34413493
let config = match derive_config_for_cwd(
34423494
&self.cli_overrides,
34433495
request_overrides,
34443496
typesafe_overrides,
34453497
history_cwd,
3446-
&self.cloud_requirements,
3498+
&cloud_requirements,
34473499
)
34483500
.await
34493501
{
@@ -3627,12 +3679,13 @@ impl CodexMessageProcessor {
36273679
),
36283680
};
36293681

3682+
let cloud_requirements = self.current_cloud_requirements();
36303683
let config = match derive_config_for_cwd(
36313684
&self.cli_overrides,
36323685
request_overrides,
36333686
typesafe_overrides,
36343687
history_cwd,
3635-
&self.cloud_requirements,
3688+
&cloud_requirements,
36363689
)
36373690
.await
36383691
{
@@ -4944,6 +4997,41 @@ fn validate_dynamic_tools(
49444997
Ok(())
49454998
}
49464999

5000+
fn replace_cloud_requirements_loader(
5001+
cloud_requirements: &RwLock<CloudRequirementsLoader>,
5002+
auth_manager: Arc<AuthManager>,
5003+
chatgpt_base_url: String,
5004+
) {
5005+
let loader = cloud_requirements_loader(auth_manager, chatgpt_base_url);
5006+
if let Ok(mut guard) = cloud_requirements.write() {
5007+
*guard = loader;
5008+
} else {
5009+
warn!("failed to update cloud requirements loader");
5010+
}
5011+
}
5012+
5013+
async fn sync_default_client_residency_requirement(
5014+
cli_overrides: &[(String, TomlValue)],
5015+
cloud_requirements: &RwLock<CloudRequirementsLoader>,
5016+
) {
5017+
let loader = cloud_requirements
5018+
.read()
5019+
.map(|guard| guard.clone())
5020+
.unwrap_or_default();
5021+
match codex_core::config::ConfigBuilder::default()
5022+
.cli_overrides(cli_overrides.to_vec())
5023+
.cloud_requirements(loader)
5024+
.build()
5025+
.await
5026+
{
5027+
Ok(config) => set_default_client_residency_requirement(config.enforce_residency.value()),
5028+
Err(err) => warn!(
5029+
error = %err,
5030+
"failed to sync default client residency requirement after auth refresh"
5031+
),
5032+
}
5033+
}
5034+
49475035
/// Derive the effective [`Config`] by layering three override sources.
49485036
///
49495037
/// Precedence (lowest to highest):

codex-rs/tui/src/lib.rs

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,7 @@ async fn run_ratatui_app(
413413
initial_config: Config,
414414
overrides: ConfigOverrides,
415415
cli_kv_overrides: Vec<(String, toml::Value)>,
416-
cloud_requirements: CloudRequirementsLoader,
416+
mut cloud_requirements: CloudRequirementsLoader,
417417
feedback: codex_feedback::CodexFeedback,
418418
) -> color_eyre::Result<AppExitInfo> {
419419
color_eyre::install()?;
@@ -470,9 +470,10 @@ async fn run_ratatui_app(
470470
should_show_onboarding(login_status, &initial_config, should_show_trust_screen_flag);
471471

472472
let config = if should_show_onboarding {
473+
let show_login_screen = should_show_login_screen(login_status, &initial_config);
473474
let onboarding_result = run_onboarding_app(
474475
OnboardingScreenArgs {
475-
show_login_screen: should_show_login_screen(login_status, &initial_config),
476+
show_login_screen,
476477
show_trust_screen: should_show_trust_screen_flag,
477478
login_status,
478479
auth_manager: auth_manager.clone(),
@@ -493,9 +494,19 @@ async fn run_ratatui_app(
493494
exit_reason: ExitReason::UserRequested,
494495
});
495496
}
496-
// If the user made an explicit trust decision, reload config so current
497-
// process state reflects what was persisted to config.toml.
498-
if onboarding_result.directory_trust_decision.is_some() {
497+
// If this onboarding run included the login step, always refresh cloud requirements and
498+
// rebuild config. This avoids missing newly available cloud requirements due to login
499+
// status detection edge cases.
500+
if show_login_screen {
501+
cloud_requirements = cloud_requirements_loader(
502+
auth_manager.clone(),
503+
initial_config.chatgpt_base_url.clone(),
504+
);
505+
}
506+
507+
// If the user made an explicit trust decision, or we showed the login flow, reload config
508+
// so current process state reflects persisted trust/auth changes.
509+
if onboarding_result.directory_trust_decision.is_some() || show_login_screen {
499510
load_config_or_exit(
500511
cli_kv_overrides.clone(),
501512
overrides.clone(),
@@ -668,6 +679,7 @@ async fn run_ratatui_app(
668679
}
669680
_ => config,
670681
};
682+
set_default_client_residency_requirement(config.enforce_residency.value());
671683
let active_profile = config.active_profile.clone();
672684
let should_show_trust_screen = should_show_trust_screen(&config);
673685

0 commit comments

Comments
 (0)