feat: reapply codext changes to rust v0.136.0

Author: Loongphy

codex-rs/Cargo.lock

@@ -226,6 +226,13 @@ pub struct GetAccountParams {
     /// themselves and call `account/login/start` with `chatgptAuthTokens`.
     #[serde(default, skip_serializing_if = "std::ops::Not::not")]
     pub refresh_token: bool,
+
+    /// When `true`, reloads the auth snapshot from storage before returning.
+    ///
+    /// This keeps long-lived clients in sync with `auth.json` updates without
+    /// requiring a full app-server restart.
+    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
+    pub reload_auth_from_storage: bool,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]

codex-rs/app-server-protocol/src/protocol/v2/account.rs

@@ -339,6 +339,7 @@ use codex_feedback::FeedbackUploadOptions;
 use codex_git_utils::git_diff_to_remote;
 use codex_git_utils::resolve_root_git_project_for_trust;
 use codex_login::AuthManager;
+use codex_login::AuthReloadStatus;
 use codex_login::CLIENT_ID;
 use codex_login::CodexAuth;
 use codex_login::ServerOptions as LoginServerOptions;
@@ -455,6 +456,78 @@ use uuid::Uuid;
 #[cfg(test)]
 use codex_app_server_protocol::ServerRequest;
 
+async fn reload_auth_from_storage_if_idle(
+    auth_manager: &Arc<AuthManager>,
+    thread_manager: &Arc<ThreadManager>,
+    config_manager: &ConfigManager,
+    outgoing: &OutgoingMessageSender,
+    thread_watch_manager: &ThreadWatchManager,
+    chatgpt_base_url: &str,
+    reason: &str,
+) {
+    if *thread_watch_manager.subscribe_running_turn_count().borrow() != 0 {
+        return;
+    }
+
+    let status = auth_manager.reload_with_status().await;
+    match handle_auth_reload_status(
+        status,
+        auth_manager,
+        thread_manager,
+        config_manager,
+        outgoing,
+        chatgpt_base_url,
+        reason,
+    )
+    .await
+    {
+        AuthReloadStatus::Reloaded { .. } => {}
+        AuthReloadStatus::Failed => {
+            warn!("failed to reload auth from storage before {reason}");
+        }
+    }
+}
+
+async fn handle_auth_reload_status(
+    status: AuthReloadStatus,
+    auth_manager: &Arc<AuthManager>,
+    thread_manager: &Arc<ThreadManager>,
+    config_manager: &ConfigManager,
+    outgoing: &OutgoingMessageSender,
+    chatgpt_base_url: &str,
+    reason: &str,
+) -> AuthReloadStatus {
+    match status {
+        AuthReloadStatus::Reloaded { changed } => {
+            if changed {
+                let invalidated_thread_count =
+                    thread_manager.invalidate_model_transport_caches().await;
+                info!(
+                    "auth reloaded from storage before {reason}; invalidated model transport caches for {invalidated_thread_count} tracked thread(s)"
+                );
+                config_manager.replace_cloud_requirements_loader(
+                    Arc::clone(auth_manager),
+                    chatgpt_base_url.to_string(),
+                );
+                config_manager
+                    .sync_default_client_residency_requirement()
+                    .await;
+                let auth = auth_manager.auth_cached();
+                outgoing
+                    .send_server_notification(ServerNotification::AccountUpdated(
+                        AccountUpdatedNotification {
+                            auth_mode: auth.as_ref().map(CodexAuth::api_auth_mode),
+                            plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type),
+                        },
+                    ))
+                    .await;
+            }
+            AuthReloadStatus::Reloaded { changed }
+        }
+        AuthReloadStatus::Failed => AuthReloadStatus::Failed,
+    }
+}
+
 mod account_processor;
 mod apps_processor;
 mod catalog_processor;

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

@@ -808,6 +808,26 @@ impl AccountRequestProcessor {
     ) -> Result<GetAccountResponse, JSONRPCErrorError> {
         let do_refresh = params.refresh_token;
 
+        if params.reload_auth_from_storage {
+            let status = self.auth_manager.reload_with_status().await;
+            match handle_auth_reload_status(
+                status,
+                &self.auth_manager,
+                &self.thread_manager,
+                &self.config_manager,
+                &self.outgoing,
+                &self.config.chatgpt_base_url,
+                "account/get",
+            )
+            .await
+            {
+                AuthReloadStatus::Reloaded { .. } => {}
+                AuthReloadStatus::Failed => {
+                    return Err(internal_error("failed to reload auth from storage"));
+                }
+            }
+        }
+
         self.refresh_token_if_requested(do_refresh).await;
 
         let provider = create_model_provider(

codex-rs/app-server/src/request_processors/account_processor.rs

@@ -827,6 +827,17 @@ impl ThreadRequestProcessor {
         app_server_client_version: Option<String>,
         request_context: RequestContext,
     ) -> Result<(), JSONRPCErrorError> {
+        reload_auth_from_storage_if_idle(
+            &self.auth_manager,
+            &self.thread_manager,
+            &self.config_manager,
+            &self.outgoing,
+            &self.thread_watch_manager,
+            &self.config.chatgpt_base_url,
+            "thread/start",
+        )
+        .await;
+
         let ThreadStartParams {
             model,
             model_provider,
@@ -2471,6 +2482,17 @@ impl ThreadRequestProcessor {
             }
         }
 
+        reload_auth_from_storage_if_idle(
+            &self.auth_manager,
+            &self.thread_manager,
+            &self.config_manager,
+            &self.outgoing,
+            &self.thread_watch_manager,
+            &self.config.chatgpt_base_url,
+            "thread/resume",
+        )
+        .await;
+
         let ThreadResumeParams {
             thread_id,
             history,

codex-rs/app-server/src/request_processors/thread_processor.rs

@@ -384,6 +384,17 @@ impl TurnRequestProcessor {
         app_server_client_name: Option<String>,
         app_server_client_version: Option<String>,
     ) -> Result<TurnStartResponse, JSONRPCErrorError> {
+        reload_auth_from_storage_if_idle(
+            &self.auth_manager,
+            &self.thread_manager,
+            &self.config_manager,
+            &self.outgoing,
+            &self.thread_watch_manager,
+            &self.config.chatgpt_base_url,
+            "turn/start",
+        )
+        .await;
+
         if let Err(error) = Self::validate_v2_input_limit(&params.input) {
             self.track_error_response(
                 &request_id,

codex-rs/app-server/src/request_processors/turn_processor.rs

@@ -731,6 +731,14 @@ pub struct Tui {
     #[serde(default)]
     pub keymap: TuiKeymap,
 
+    /// Optional synthetic user-turn prompt injected after a turn fails with
+    /// `UsageLimitExceeded`.
+    ///
+    /// When unset, Codex uses the built-in default recovery prompt.
+    /// When set to an empty string, Codex disables this automatic recovery turn.
+    #[serde(default)]
+    pub usage_limit_resume_prompt: Option<String>,
+
     /// Startup tooltip availability NUX state persisted by the TUI.
     #[serde(default)]
     pub model_availability_nux: ModelAvailabilityNuxConfig,

codex-rs/config/src/types.rs

@@ -399,6 +399,10 @@ impl ModelClient {
         self.store_cached_websocket_session(WebsocketSession::default());
     }
 
+    pub(crate) fn invalidate_cached_transport_state(&self) {
+        self.store_cached_websocket_session(WebsocketSession::default());
+    }
+
     pub(crate) fn current_window_id(&self) -> String {
         let thread_id = self.state.thread_id;
         let window_generation = self.state.window_generation.load(Ordering::Relaxed);

codex-rs/core/src/client.rs

@@ -142,6 +142,14 @@ impl CodexThread {
         self.codex.shutdown_and_wait().await
     }
 
+    pub(crate) fn invalidate_model_transport_cache(&self) {
+        self.codex
+            .session
+            .services
+            .model_client
+            .invalidate_cached_transport_state();
+    }
+
     /// Wait until the underlying session loop has terminated.
     pub async fn wait_until_terminated(&self) {
         self.codex.session_loop_termination.clone().await;

codex-rs/core/src/codex_thread.rs

@@ -729,6 +729,13 @@ pub struct Config {
     /// 3. built-in defaults
     pub tui_keymap: TuiKeymap,
 
+    /// Synthetic user-turn prompt injected after a `UsageLimitExceeded` turn
+    /// failure.
+    ///
+    /// `None` uses the built-in default prompt. `Some("")` disables the
+    /// automatic recovery turn.
+    pub tui_usage_limit_resume_prompt: Option<String>,
+
     /// The absolute directory that should be treated as the current working
     /// directory for the session. All relative paths inside the business-logic
     /// layer are resolved against this path.
@@ -3557,6 +3564,10 @@ impl Config {
                 .as_ref()
                 .map(|t| t.keymap.clone())
                 .unwrap_or_default(),
+            tui_usage_limit_resume_prompt: cfg
+                .tui
+                .as_ref()
+                .and_then(|t| t.usage_limit_resume_prompt.clone()),
             otel,
         };
         Ok(config)