[App-server] Implement v2 for account/login/start and account/login/completed (#6183)
This PR implements `account/login/start` and `account/login/completed`. Instead of having separate endpoints for login with chatgpt and api, we have a single enum handling different login methods. For sync auth methods like sign in with api key, we still send a `completed` notification back to be compatible with the async login flow.
This commit is contained in:
Celia Chen
GitHub
@@ -4,6 +4,7 @@ use crate::fuzzy_file_search::run_fuzzy_file_search;
|
||||
use crate::models::supported_models;
|
||||
use crate::outgoing_message::OutgoingMessageSender;
|
||||
use crate::outgoing_message::OutgoingNotification;
|
||||
use codex_app_server_protocol::AccountLoginCompletedNotification;
|
||||
use codex_app_server_protocol::AccountRateLimitsUpdatedNotification;
|
||||
use codex_app_server_protocol::AccountUpdatedNotification;
|
||||
use codex_app_server_protocol::AddConversationListenerParams;
|
||||
@@ -36,6 +37,7 @@ use codex_app_server_protocol::InterruptConversationResponse;
|
||||
use codex_app_server_protocol::JSONRPCErrorError;
|
||||
use codex_app_server_protocol::ListConversationsParams;
|
||||
use codex_app_server_protocol::ListConversationsResponse;
|
||||
use codex_app_server_protocol::LoginAccountParams;
|
||||
use codex_app_server_protocol::LoginApiKeyParams;
|
||||
use codex_app_server_protocol::LoginApiKeyResponse;
|
||||
use codex_app_server_protocol::LoginChatGptCompleteNotification;
|
||||
@@ -222,12 +224,8 @@ impl CodexMessageProcessor {
|
||||
ClientRequest::ModelList { request_id, params } => {
|
||||
self.list_models(request_id, params).await;
|
||||
}
|
||||
ClientRequest::LoginAccount {
|
||||
request_id,
|
||||
params: _,
|
||||
} => {
|
||||
self.send_unimplemented_error(request_id, "account/login")
|
||||
.await;
|
||||
ClientRequest::LoginAccount { request_id, params } => {
|
||||
self.login_v2(request_id, params).await;
|
||||
}
|
||||
ClientRequest::LogoutAccount {
|
||||
request_id,
|
||||
@@ -267,13 +265,13 @@ impl CodexMessageProcessor {
|
||||
self.git_diff_to_origin(request_id, params.cwd).await;
|
||||
}
|
||||
ClientRequest::LoginApiKey { request_id, params } => {
|
||||
self.login_api_key(request_id, params).await;
|
||||
self.login_api_key_v1(request_id, params).await;
|
||||
}
|
||||
ClientRequest::LoginChatGpt {
|
||||
request_id,
|
||||
params: _,
|
||||
} => {
|
||||
self.login_chatgpt(request_id).await;
|
||||
self.login_chatgpt_v1(request_id).await;
|
||||
}
|
||||
ClientRequest::CancelLoginChatGpt { request_id, params } => {
|
||||
self.cancel_login_chatgpt(request_id, params.login_id).await;
|
||||
@@ -335,20 +333,34 @@ impl CodexMessageProcessor {
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
}
|
||||
|
||||
async fn login_api_key(&mut self, request_id: RequestId, params: LoginApiKeyParams) {
|
||||
async fn login_v2(&mut self, request_id: RequestId, params: LoginAccountParams) {
|
||||
match params {
|
||||
LoginAccountParams::ApiKey { api_key } => {
|
||||
self.login_api_key_v2(request_id, LoginApiKeyParams { api_key })
|
||||
.await;
|
||||
}
|
||||
LoginAccountParams::Chatgpt => {
|
||||
self.login_chatgpt_v2(request_id).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn login_api_key_common(
|
||||
&mut self,
|
||||
params: &LoginApiKeyParams,
|
||||
) -> std::result::Result<(), JSONRPCErrorError> {
|
||||
if matches!(
|
||||
self.config.forced_login_method,
|
||||
Some(ForcedLoginMethod::Chatgpt)
|
||||
) {
|
||||
let error = JSONRPCErrorError {
|
||||
return Err(JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "API key login is disabled. Use ChatGPT login instead.".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
return;
|
||||
});
|
||||
}
|
||||
|
||||
// Cancel any active login attempt.
|
||||
{
|
||||
let mut guard = self.active_login.lock().await;
|
||||
if let Some(active) = guard.take() {
|
||||
@@ -363,6 +375,19 @@ impl CodexMessageProcessor {
|
||||
) {
|
||||
Ok(()) => {
|
||||
self.auth_manager.reload();
|
||||
Ok(())
|
||||
}
|
||||
Err(err) => Err(JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to save api key: {err}"),
|
||||
data: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn login_api_key_v1(&mut self, request_id: RequestId, params: LoginApiKeyParams) {
|
||||
match self.login_api_key_common(¶ms).await {
|
||||
Ok(()) => {
|
||||
self.outgoing
|
||||
.send_response(request_id, LoginApiKeyResponse {})
|
||||
.await;
|
||||
@@ -374,31 +399,57 @@ impl CodexMessageProcessor {
|
||||
.send_server_notification(ServerNotification::AuthStatusChange(payload))
|
||||
.await;
|
||||
}
|
||||
Err(err) => {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to save api key: {err}"),
|
||||
data: None,
|
||||
};
|
||||
Err(error) => {
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn login_chatgpt(&mut self, request_id: RequestId) {
|
||||
async fn login_api_key_v2(&mut self, request_id: RequestId, params: LoginApiKeyParams) {
|
||||
match self.login_api_key_common(¶ms).await {
|
||||
Ok(()) => {
|
||||
let response = codex_app_server_protocol::LoginAccountResponse::ApiKey {};
|
||||
self.outgoing.send_response(request_id, response).await;
|
||||
|
||||
let payload_login_completed = AccountLoginCompletedNotification {
|
||||
login_id: None,
|
||||
success: true,
|
||||
error: None,
|
||||
};
|
||||
self.outgoing
|
||||
.send_server_notification(ServerNotification::AccountLoginCompleted(
|
||||
payload_login_completed,
|
||||
))
|
||||
.await;
|
||||
|
||||
let payload_v2 = AccountUpdatedNotification {
|
||||
auth_mode: self.auth_manager.auth().map(|auth| auth.mode),
|
||||
};
|
||||
self.outgoing
|
||||
.send_server_notification(ServerNotification::AccountUpdated(payload_v2))
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build options for a ChatGPT login attempt; performs validation.
|
||||
async fn login_chatgpt_common(
|
||||
&self,
|
||||
) -> std::result::Result<LoginServerOptions, JSONRPCErrorError> {
|
||||
let config = self.config.as_ref();
|
||||
|
||||
if matches!(config.forced_login_method, Some(ForcedLoginMethod::Api)) {
|
||||
let error = JSONRPCErrorError {
|
||||
return Err(JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "ChatGPT login is disabled. Use API key login instead.".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
return;
|
||||
});
|
||||
}
|
||||
|
||||
let opts = LoginServerOptions {
|
||||
Ok(LoginServerOptions {
|
||||
open_browser: false,
|
||||
..LoginServerOptions::new(
|
||||
config.codex_home.clone(),
|
||||
@@ -406,99 +457,190 @@ impl CodexMessageProcessor {
|
||||
config.forced_chatgpt_workspace_id.clone(),
|
||||
config.cli_auth_credentials_store_mode,
|
||||
)
|
||||
};
|
||||
})
|
||||
}
|
||||
|
||||
enum LoginChatGptReply {
|
||||
Response(LoginChatGptResponse),
|
||||
Error(JSONRPCErrorError),
|
||||
}
|
||||
// Deprecated in favor of login_chatgpt_v2.
|
||||
async fn login_chatgpt_v1(&mut self, request_id: RequestId) {
|
||||
match self.login_chatgpt_common().await {
|
||||
Ok(opts) => match run_login_server(opts) {
|
||||
Ok(server) => {
|
||||
let login_id = Uuid::new_v4();
|
||||
let shutdown_handle = server.cancel_handle();
|
||||
|
||||
let reply = match run_login_server(opts) {
|
||||
Ok(server) => {
|
||||
let login_id = Uuid::new_v4();
|
||||
let shutdown_handle = server.cancel_handle();
|
||||
|
||||
// Replace active login if present.
|
||||
{
|
||||
let mut guard = self.active_login.lock().await;
|
||||
if let Some(existing) = guard.take() {
|
||||
existing.drop();
|
||||
}
|
||||
*guard = Some(ActiveLogin {
|
||||
shutdown_handle: shutdown_handle.clone(),
|
||||
login_id,
|
||||
});
|
||||
}
|
||||
|
||||
let response = LoginChatGptResponse {
|
||||
login_id,
|
||||
auth_url: server.auth_url.clone(),
|
||||
};
|
||||
|
||||
// Spawn background task to monitor completion.
|
||||
let outgoing_clone = self.outgoing.clone();
|
||||
let active_login = self.active_login.clone();
|
||||
let auth_manager = self.auth_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
let (success, error_msg) = match tokio::time::timeout(
|
||||
LOGIN_CHATGPT_TIMEOUT,
|
||||
server.block_until_done(),
|
||||
)
|
||||
.await
|
||||
// Replace active login if present.
|
||||
{
|
||||
Ok(Ok(())) => (true, None),
|
||||
Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))),
|
||||
Err(_elapsed) => {
|
||||
// Timeout: cancel server and report
|
||||
shutdown_handle.shutdown();
|
||||
(false, Some("Login timed out".to_string()))
|
||||
let mut guard = self.active_login.lock().await;
|
||||
if let Some(existing) = guard.take() {
|
||||
existing.drop();
|
||||
}
|
||||
};
|
||||
let payload = LoginChatGptCompleteNotification {
|
||||
login_id,
|
||||
success,
|
||||
error: error_msg,
|
||||
};
|
||||
outgoing_clone
|
||||
.send_server_notification(ServerNotification::LoginChatGptComplete(payload))
|
||||
.await;
|
||||
*guard = Some(ActiveLogin {
|
||||
shutdown_handle: shutdown_handle.clone(),
|
||||
login_id,
|
||||
});
|
||||
}
|
||||
|
||||
// Send an auth status change notification.
|
||||
if success {
|
||||
// Update in-memory auth cache now that login completed.
|
||||
auth_manager.reload();
|
||||
// Spawn background task to monitor completion.
|
||||
let outgoing_clone = self.outgoing.clone();
|
||||
let active_login = self.active_login.clone();
|
||||
let auth_manager = self.auth_manager.clone();
|
||||
let auth_url = server.auth_url.clone();
|
||||
tokio::spawn(async move {
|
||||
let (success, error_msg) = match tokio::time::timeout(
|
||||
LOGIN_CHATGPT_TIMEOUT,
|
||||
server.block_until_done(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => (true, None),
|
||||
Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))),
|
||||
Err(_elapsed) => {
|
||||
shutdown_handle.shutdown();
|
||||
(false, Some("Login timed out".to_string()))
|
||||
}
|
||||
};
|
||||
|
||||
// Notify clients with the actual current auth mode.
|
||||
let current_auth_method = auth_manager.auth().map(|a| a.mode);
|
||||
let payload = AuthStatusChangeNotification {
|
||||
auth_method: current_auth_method,
|
||||
let payload = LoginChatGptCompleteNotification {
|
||||
login_id,
|
||||
success,
|
||||
error: error_msg.clone(),
|
||||
};
|
||||
outgoing_clone
|
||||
.send_server_notification(ServerNotification::AuthStatusChange(payload))
|
||||
.send_server_notification(ServerNotification::LoginChatGptComplete(
|
||||
payload,
|
||||
))
|
||||
.await;
|
||||
|
||||
if success {
|
||||
auth_manager.reload();
|
||||
|
||||
// Notify clients with the actual current auth mode.
|
||||
let current_auth_method = auth_manager.auth().map(|a| a.mode);
|
||||
let payload = AuthStatusChangeNotification {
|
||||
auth_method: current_auth_method,
|
||||
};
|
||||
outgoing_clone
|
||||
.send_server_notification(ServerNotification::AuthStatusChange(
|
||||
payload,
|
||||
))
|
||||
.await;
|
||||
}
|
||||
|
||||
// Clear the active login if it matches this attempt. It may have been replaced or cancelled.
|
||||
let mut guard = active_login.lock().await;
|
||||
if guard.as_ref().map(|l| l.login_id) == Some(login_id) {
|
||||
*guard = None;
|
||||
}
|
||||
});
|
||||
|
||||
let response = LoginChatGptResponse { login_id, auth_url };
|
||||
self.outgoing.send_response(request_id, response).await;
|
||||
}
|
||||
Err(err) => {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to start login server: {err}"),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
self.outgoing.send_error(request_id, err).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn login_chatgpt_v2(&mut self, request_id: RequestId) {
|
||||
match self.login_chatgpt_common().await {
|
||||
Ok(opts) => match run_login_server(opts) {
|
||||
Ok(server) => {
|
||||
let login_id = Uuid::new_v4();
|
||||
let shutdown_handle = server.cancel_handle();
|
||||
|
||||
// Replace active login if present.
|
||||
{
|
||||
let mut guard = self.active_login.lock().await;
|
||||
if let Some(existing) = guard.take() {
|
||||
existing.drop();
|
||||
}
|
||||
*guard = Some(ActiveLogin {
|
||||
shutdown_handle: shutdown_handle.clone(),
|
||||
login_id,
|
||||
});
|
||||
}
|
||||
|
||||
// Clear the active login if it matches this attempt. It may have been replaced or cancelled.
|
||||
let mut guard = active_login.lock().await;
|
||||
if guard.as_ref().map(|l| l.login_id) == Some(login_id) {
|
||||
*guard = None;
|
||||
}
|
||||
});
|
||||
// Spawn background task to monitor completion.
|
||||
let outgoing_clone = self.outgoing.clone();
|
||||
let active_login = self.active_login.clone();
|
||||
let auth_manager = self.auth_manager.clone();
|
||||
let auth_url = server.auth_url.clone();
|
||||
tokio::spawn(async move {
|
||||
let (success, error_msg) = match tokio::time::timeout(
|
||||
LOGIN_CHATGPT_TIMEOUT,
|
||||
server.block_until_done(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => (true, None),
|
||||
Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))),
|
||||
Err(_elapsed) => {
|
||||
shutdown_handle.shutdown();
|
||||
(false, Some("Login timed out".to_string()))
|
||||
}
|
||||
};
|
||||
|
||||
LoginChatGptReply::Response(response)
|
||||
}
|
||||
Err(err) => LoginChatGptReply::Error(JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to start login server: {err}"),
|
||||
data: None,
|
||||
}),
|
||||
};
|
||||
let payload_v2 = AccountLoginCompletedNotification {
|
||||
login_id: Some(login_id.to_string()),
|
||||
success,
|
||||
error: error_msg,
|
||||
};
|
||||
outgoing_clone
|
||||
.send_server_notification(ServerNotification::AccountLoginCompleted(
|
||||
payload_v2,
|
||||
))
|
||||
.await;
|
||||
|
||||
match reply {
|
||||
LoginChatGptReply::Response(resp) => {
|
||||
self.outgoing.send_response(request_id, resp).await
|
||||
if success {
|
||||
auth_manager.reload();
|
||||
|
||||
// Notify clients with the actual current auth mode.
|
||||
let current_auth_method = auth_manager.auth().map(|a| a.mode);
|
||||
let payload_v2 = AccountUpdatedNotification {
|
||||
auth_mode: current_auth_method,
|
||||
};
|
||||
outgoing_clone
|
||||
.send_server_notification(ServerNotification::AccountUpdated(
|
||||
payload_v2,
|
||||
))
|
||||
.await;
|
||||
}
|
||||
|
||||
// Clear the active login if it matches this attempt. It may have been replaced or cancelled.
|
||||
let mut guard = active_login.lock().await;
|
||||
if guard.as_ref().map(|l| l.login_id) == Some(login_id) {
|
||||
*guard = None;
|
||||
}
|
||||
});
|
||||
|
||||
let response = codex_app_server_protocol::LoginAccountResponse::Chatgpt {
|
||||
login_id: login_id.to_string(),
|
||||
auth_url,
|
||||
};
|
||||
self.outgoing.send_response(request_id, response).await;
|
||||
}
|
||||
Err(err) => {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to start login server: {err}"),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
self.outgoing.send_error(request_id, err).await;
|
||||
}
|
||||
LoginChatGptReply::Error(err) => self.outgoing.send_error(request_id, err).await,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -581,7 +723,7 @@ impl CodexMessageProcessor {
|
||||
.await;
|
||||
|
||||
let payload_v2 = AccountUpdatedNotification {
|
||||
auth_method: current_auth_method,
|
||||
auth_mode: current_auth_method,
|
||||
};
|
||||
self.outgoing
|
||||
.send_server_notification(ServerNotification::AccountUpdated(payload_v2))
|
||||
|
||||
Reference in New Issue
Block a user