Add device-code onboarding and ChatGPT token refresh to app-server TUI

codex-rs/tui_app_server/src/app/app_server_adapter.rs

@@ -16,9 +16,12 @@ use crate::app_event::AppEvent;
 use crate::app_server_session::AppServerSession;
 use crate::app_server_session::app_server_rate_limit_snapshot_to_core;
 use crate::app_server_session::status_account_display_from_auth_mode;
+use crate::local_chatgpt_auth::load_local_chatgpt_auth;
 use codex_app_server_client::AppServerEvent;
+use codex_app_server_protocol::ChatgptAuthTokensRefreshParams;
 use codex_app_server_protocol::JSONRPCErrorError;
 use codex_app_server_protocol::ServerNotification;
+use codex_app_server_protocol::ServerRequest;
 use codex_app_server_protocol::Thread;
 use codex_app_server_protocol::ThreadItem;
 use codex_app_server_protocol::Turn;
@@ -90,6 +93,7 @@ impl App {
                         matches!(
                             notification.auth_mode,
                             Some(codex_app_server_protocol::AuthMode::Chatgpt)
+                                | Some(codex_app_server_protocol::AuthMode::ChatgptAuthTokens)
                         ),
                     );
                 }
@@ -150,6 +154,15 @@ impl App {
                 }
             }
             AppServerEvent::ServerRequest(request) => {
+                if let ServerRequest::ChatgptAuthTokensRefresh { request_id, params } = request {
+                    self.handle_chatgpt_auth_tokens_refresh_request(
+                        app_server_client,
+                        request_id,
+                        params,
+                    )
+                    .await;
+                    return;
+                }
                 if let Some(unsupported) = self
                     .pending_app_server_requests
                     .note_server_request(&request)
@@ -181,6 +194,70 @@ impl App {
         }
     }
 
+    async fn handle_chatgpt_auth_tokens_refresh_request(
+        &mut self,
+        app_server_client: &AppServerSession,
+        request_id: codex_app_server_protocol::RequestId,
+        params: ChatgptAuthTokensRefreshParams,
+    ) {
+        let config = self.config.clone();
+        let result = tokio::task::spawn_blocking(move || {
+            resolve_chatgpt_auth_tokens_refresh_response(
+                &config.codex_home,
+                config.cli_auth_credentials_store_mode,
+                config.forced_chatgpt_workspace_id.as_deref(),
+                &params,
+            )
+        })
+        .await;
+
+        match result {
+            Ok(Ok(response)) => {
+                let response = serde_json::to_value(response).map_err(|err| {
+                    format!("failed to serialize chatgpt auth refresh response: {err}")
+                });
+                match response {
+                    Ok(response) => {
+                        if let Err(err) = app_server_client
+                            .resolve_server_request(request_id, response)
+                            .await
+                        {
+                            tracing::warn!("failed to resolve chatgpt auth refresh request: {err}");
+                        }
+                    }
+                    Err(err) => {
+                        self.chat_widget.add_error_message(err.clone());
+                        if let Err(reject_err) = self
+                            .reject_app_server_request(app_server_client, request_id, err)
+                            .await
+                        {
+                            tracing::warn!("{reject_err}");
+                        }
+                    }
+                }
+            }
+            Ok(Err(err)) => {
+                self.chat_widget.add_error_message(err.clone());
+                if let Err(reject_err) = self
+                    .reject_app_server_request(app_server_client, request_id, err)
+                    .await
+                {
+                    tracing::warn!("{reject_err}");
+                }
+            }
+            Err(err) => {
+                let message = format!("chatgpt auth refresh task failed: {err}");
+                self.chat_widget.add_error_message(message.clone());
+                if let Err(reject_err) = self
+                    .reject_app_server_request(app_server_client, request_id, message)
+                    .await
+                {
+                    tracing::warn!("{reject_err}");
+                }
+            }
+        }
+    }
+
     async fn reject_app_server_request(
         &self,
         app_server_client: &AppServerSession,
@@ -201,6 +278,28 @@ impl App {
     }
 }
 
+fn resolve_chatgpt_auth_tokens_refresh_response(
+    codex_home: &std::path::Path,
+    auth_credentials_store_mode: codex_core::auth::AuthCredentialsStoreMode,
+    forced_chatgpt_workspace_id: Option<&str>,
+    params: &ChatgptAuthTokensRefreshParams,
+) -> Result<codex_app_server_protocol::ChatgptAuthTokensRefreshResponse, String> {
+    let auth = load_local_chatgpt_auth(
+        codex_home,
+        auth_credentials_store_mode,
+        forced_chatgpt_workspace_id,
+    )?;
+    if let Some(previous_account_id) = params.previous_account_id.as_deref()
+        && previous_account_id != auth.chatgpt_account_id
+    {
+        return Err(format!(
+            "local ChatGPT auth refresh account mismatch: expected `{previous_account_id}`, got `{}`",
+            auth.chatgpt_account_id
+        ));
+    }
+    Ok(auth.to_refresh_response())
+}
+
 /// Convert a `Thread` snapshot into a flat sequence of protocol `Event`s
 /// suitable for replaying into the TUI event store.
 ///
@@ -624,6 +723,113 @@ fn thread_item_to_core(item: &ThreadItem) -> Option<TurnItem> {
     }
 }
 
+#[cfg(test)]
+mod refresh_tests {
+    use super::*;
+
+    use base64::Engine;
+    use chrono::Utc;
+    use codex_app_server_protocol::AuthMode;
+    use codex_core::auth::AuthCredentialsStoreMode;
+    use codex_core::auth::AuthDotJson;
+    use codex_core::auth::save_auth;
+    use codex_core::token_data::TokenData;
+    use pretty_assertions::assert_eq;
+    use serde::Serialize;
+    use serde_json::json;
+    use tempfile::TempDir;
+
+    fn fake_jwt(account_id: &str, plan_type: &str) -> String {
+        #[derive(Serialize)]
+        struct Header {
+            alg: &'static str,
+            typ: &'static str,
+        }
+
+        let header = Header {
+            alg: "none",
+            typ: "JWT",
+        };
+        let payload = json!({
+            "email": "user@example.com",
+            "https://api.openai.com/auth": {
+                "chatgpt_account_id": account_id,
+                "chatgpt_plan_type": plan_type,
+            },
+        });
+        let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
+        let header_b64 = encode(&serde_json::to_vec(&header).expect("serialize header"));
+        let payload_b64 = encode(&serde_json::to_vec(&payload).expect("serialize payload"));
+        let signature_b64 = encode(b"sig");
+        format!("{header_b64}.{payload_b64}.{signature_b64}")
+    }
+
+    fn write_chatgpt_auth(codex_home: &std::path::Path) {
+        let id_token = fake_jwt("workspace-1", "business");
+        let access_token = fake_jwt("workspace-1", "business");
+        save_auth(
+            codex_home,
+            &AuthDotJson {
+                auth_mode: Some(AuthMode::Chatgpt),
+                openai_api_key: None,
+                tokens: Some(TokenData {
+                    id_token: codex_core::token_data::parse_chatgpt_jwt_claims(&id_token)
+                        .expect("id token should parse"),
+                    access_token,
+                    refresh_token: "refresh-token".to_string(),
+                    account_id: Some("workspace-1".to_string()),
+                }),
+                last_refresh: Some(Utc::now()),
+            },
+            AuthCredentialsStoreMode::File,
+        )
+        .expect("chatgpt auth should save");
+    }
+
+    #[test]
+    fn refresh_request_uses_local_chatgpt_auth() {
+        let codex_home = TempDir::new().expect("tempdir");
+        write_chatgpt_auth(codex_home.path());
+
+        let response = resolve_chatgpt_auth_tokens_refresh_response(
+            codex_home.path(),
+            AuthCredentialsStoreMode::File,
+            Some("workspace-1"),
+            &ChatgptAuthTokensRefreshParams {
+                reason: codex_app_server_protocol::ChatgptAuthTokensRefreshReason::Unauthorized,
+                previous_account_id: Some("workspace-1".to_string()),
+            },
+        )
+        .expect("refresh response should resolve");
+
+        assert_eq!(response.chatgpt_account_id, "workspace-1");
+        assert_eq!(response.chatgpt_plan_type.as_deref(), Some("business"));
+        assert!(!response.access_token.is_empty());
+    }
+
+    #[test]
+    fn refresh_request_rejects_account_mismatch() {
+        let codex_home = TempDir::new().expect("tempdir");
+        write_chatgpt_auth(codex_home.path());
+
+        let err = resolve_chatgpt_auth_tokens_refresh_response(
+            codex_home.path(),
+            AuthCredentialsStoreMode::File,
+            Some("workspace-1"),
+            &ChatgptAuthTokensRefreshParams {
+                reason: codex_app_server_protocol::ChatgptAuthTokensRefreshReason::Unauthorized,
+                previous_account_id: Some("workspace-2".to_string()),
+            },
+        )
+        .expect_err("mismatched account should fail");
+
+        assert_eq!(
+            err,
+            "local ChatGPT auth refresh account mismatch: expected `workspace-2`, got `workspace-1`"
+        );
+    }
+}
+
 fn app_server_web_search_action_to_core(
     action: codex_app_server_protocol::WebSearchAction,
 ) -> Option<codex_protocol::models::WebSearchAction> {

codex-rs/tui_app_server/src/app/app_server_requests.rs

@@ -99,13 +99,7 @@ impl PendingAppServerRequests {
                         .to_string(),
                 })
             }
-            ServerRequest::ChatgptAuthTokensRefresh { request_id, .. } => {
-                Some(UnsupportedAppServerRequest {
-                    request_id: request_id.clone(),
-                    message: "ChatGPT auth token refresh is not available in app-server TUI yet."
-                        .to_string(),
-                })
-            }
+            ServerRequest::ChatgptAuthTokensRefresh { .. } => None,
             ServerRequest::ApplyPatchApproval { request_id, .. } => {
                 Some(UnsupportedAppServerRequest {
                     request_id: request_id.clone(),
@@ -608,6 +602,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn does_not_mark_chatgpt_auth_refresh_as_unsupported() {
+        let mut pending = PendingAppServerRequests::default();
+
+        assert_eq!(
+            pending.note_server_request(&ServerRequest::ChatgptAuthTokensRefresh {
+                request_id: AppServerRequestId::Integer(100),
+                params: codex_app_server_protocol::ChatgptAuthTokensRefreshParams {
+                    reason: codex_app_server_protocol::ChatgptAuthTokensRefreshReason::Unauthorized,
+                    previous_account_id: Some("workspace-1".to_string()),
+                },
+            }),
+            None
+        );
+    }
+
     #[test]
     fn rejects_invalid_patch_decisions_for_file_change_requests() {
         let mut pending = PendingAppServerRequests::default();

codex-rs/tui_app_server/src/lib.rs

@@ -100,6 +100,7 @@ pub mod insert_history;
 mod key_hint;
 mod line_truncation;
 pub mod live_wrap;
+mod local_chatgpt_auth;
 mod markdown;
 mod markdown_render;
 mod markdown_stream;