e84e39940b
This PR does two things: 1. add a new function in core that maps the core-internal plan type to the external plan type; 2. implement account/read that get account status (v2 of `getAuthStatus`).
493 lines
15 KiB
Rust
493 lines
15 KiB
Rust
use anyhow::Result;
|
|
use anyhow::bail;
|
|
use app_test_support::McpProcess;
|
|
use app_test_support::to_response;
|
|
|
|
use app_test_support::ChatGptAuthFixture;
|
|
use app_test_support::write_chatgpt_auth;
|
|
use codex_app_server_protocol::Account;
|
|
use codex_app_server_protocol::AuthMode;
|
|
use codex_app_server_protocol::CancelLoginAccountParams;
|
|
use codex_app_server_protocol::CancelLoginAccountResponse;
|
|
use codex_app_server_protocol::GetAccountParams;
|
|
use codex_app_server_protocol::GetAccountResponse;
|
|
use codex_app_server_protocol::JSONRPCError;
|
|
use codex_app_server_protocol::JSONRPCResponse;
|
|
use codex_app_server_protocol::LoginAccountResponse;
|
|
use codex_app_server_protocol::LogoutAccountResponse;
|
|
use codex_app_server_protocol::RequestId;
|
|
use codex_app_server_protocol::ServerNotification;
|
|
use codex_core::auth::AuthCredentialsStoreMode;
|
|
use codex_login::login_with_api_key;
|
|
use codex_protocol::account::PlanType as AccountPlanType;
|
|
use pretty_assertions::assert_eq;
|
|
use serial_test::serial;
|
|
use std::path::Path;
|
|
use std::time::Duration;
|
|
use tempfile::TempDir;
|
|
use tokio::time::timeout;
|
|
|
|
const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
|
|
|
// Helper to create a minimal config.toml for the app server
|
|
#[derive(Default)]
|
|
struct CreateConfigTomlParams {
|
|
forced_method: Option<String>,
|
|
forced_workspace_id: Option<String>,
|
|
requires_openai_auth: Option<bool>,
|
|
}
|
|
|
|
fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std::io::Result<()> {
|
|
let config_toml = codex_home.join("config.toml");
|
|
let forced_line = if let Some(method) = params.forced_method {
|
|
format!("forced_login_method = \"{method}\"\n")
|
|
} else {
|
|
String::new()
|
|
};
|
|
let forced_workspace_line = if let Some(ws) = params.forced_workspace_id {
|
|
format!("forced_chatgpt_workspace_id = \"{ws}\"\n")
|
|
} else {
|
|
String::new()
|
|
};
|
|
let requires_line = match params.requires_openai_auth {
|
|
Some(true) => "requires_openai_auth = true\n".to_string(),
|
|
Some(false) => String::new(),
|
|
None => String::new(),
|
|
};
|
|
let contents = format!(
|
|
r#"
|
|
model = "mock-model"
|
|
approval_policy = "never"
|
|
sandbox_mode = "danger-full-access"
|
|
{forced_line}
|
|
{forced_workspace_line}
|
|
|
|
model_provider = "mock_provider"
|
|
|
|
[model_providers.mock_provider]
|
|
name = "Mock provider for test"
|
|
base_url = "http://127.0.0.1:0/v1"
|
|
wire_api = "chat"
|
|
request_max_retries = 0
|
|
stream_max_retries = 0
|
|
{requires_line}
|
|
"#
|
|
);
|
|
std::fs::write(config_toml, contents)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn logout_account_removes_auth_and_notifies() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
|
|
|
|
login_with_api_key(
|
|
codex_home.path(),
|
|
"sk-test-key",
|
|
AuthCredentialsStoreMode::File,
|
|
)?;
|
|
assert!(codex_home.path().join("auth.json").exists());
|
|
|
|
let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let id = mcp.send_logout_account_request().await?;
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(id)),
|
|
)
|
|
.await??;
|
|
let _ok: LogoutAccountResponse = to_response(resp)?;
|
|
|
|
let note = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_notification_message("account/updated"),
|
|
)
|
|
.await??;
|
|
let parsed: ServerNotification = note.try_into()?;
|
|
let ServerNotification::AccountUpdated(payload) = parsed else {
|
|
bail!("unexpected notification: {parsed:?}");
|
|
};
|
|
assert!(
|
|
payload.auth_mode.is_none(),
|
|
"auth_method should be None after logout"
|
|
);
|
|
|
|
assert!(
|
|
!codex_home.path().join("auth.json").exists(),
|
|
"auth.json should be deleted"
|
|
);
|
|
|
|
let get_id = mcp
|
|
.send_get_account_request(GetAccountParams {
|
|
refresh_token: false,
|
|
})
|
|
.await?;
|
|
let get_resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(get_id)),
|
|
)
|
|
.await??;
|
|
let account: GetAccountResponse = to_response(get_resp)?;
|
|
assert_eq!(account.account, None);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn login_account_api_key_succeeds_and_notifies() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let req_id = mcp
|
|
.send_login_account_api_key_request("sk-test-key")
|
|
.await?;
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(req_id)),
|
|
)
|
|
.await??;
|
|
let login: LoginAccountResponse = to_response(resp)?;
|
|
assert_eq!(login, LoginAccountResponse::ApiKey {});
|
|
|
|
let note = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_notification_message("account/login/completed"),
|
|
)
|
|
.await??;
|
|
let parsed: ServerNotification = note.try_into()?;
|
|
let ServerNotification::AccountLoginCompleted(payload) = parsed else {
|
|
bail!("unexpected notification: {parsed:?}");
|
|
};
|
|
pretty_assertions::assert_eq!(payload.login_id, None);
|
|
pretty_assertions::assert_eq!(payload.success, true);
|
|
pretty_assertions::assert_eq!(payload.error, None);
|
|
|
|
let note = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_notification_message("account/updated"),
|
|
)
|
|
.await??;
|
|
let parsed: ServerNotification = note.try_into()?;
|
|
let ServerNotification::AccountUpdated(payload) = parsed else {
|
|
bail!("unexpected notification: {parsed:?}");
|
|
};
|
|
pretty_assertions::assert_eq!(payload.auth_mode, Some(AuthMode::ApiKey));
|
|
|
|
assert!(codex_home.path().join("auth.json").exists());
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
forced_method: Some("chatgpt".to_string()),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let request_id = mcp
|
|
.send_login_account_api_key_request("sk-test-key")
|
|
.await?;
|
|
let err: JSONRPCError = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_error_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
|
|
assert_eq!(
|
|
err.error.message,
|
|
"API key login is disabled. Use ChatGPT login instead."
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn login_account_chatgpt_rejected_when_forced_api() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
forced_method: Some("api".to_string()),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let request_id = mcp.send_login_account_chatgpt_request().await?;
|
|
let err: JSONRPCError = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_error_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
|
|
assert_eq!(
|
|
err.error.message,
|
|
"ChatGPT login is disabled. Use API key login instead."
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
// Serialize tests that launch the login server since it binds to a fixed port.
|
|
#[serial(login_port)]
|
|
async fn login_account_chatgpt_start() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let request_id = mcp.send_login_account_chatgpt_request().await?;
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
|
|
let login: LoginAccountResponse = to_response(resp)?;
|
|
let LoginAccountResponse::Chatgpt { login_id, auth_url } = login else {
|
|
bail!("unexpected login response: {login:?}");
|
|
};
|
|
assert!(
|
|
auth_url.contains("redirect_uri=http%3A%2F%2Flocalhost"),
|
|
"auth_url should contain a redirect_uri to localhost"
|
|
);
|
|
|
|
let cancel_id = mcp
|
|
.send_cancel_login_account_request(CancelLoginAccountParams {
|
|
login_id: login_id.clone(),
|
|
})
|
|
.await?;
|
|
let cancel_resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)),
|
|
)
|
|
.await??;
|
|
let _ok: CancelLoginAccountResponse = to_response(cancel_resp)?;
|
|
|
|
let note = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_notification_message("account/login/completed"),
|
|
)
|
|
.await??;
|
|
let parsed: ServerNotification = note.try_into()?;
|
|
let ServerNotification::AccountLoginCompleted(payload) = parsed else {
|
|
bail!("unexpected notification: {parsed:?}");
|
|
};
|
|
pretty_assertions::assert_eq!(payload.login_id, Some(login_id));
|
|
pretty_assertions::assert_eq!(payload.success, false);
|
|
assert!(
|
|
payload.error.is_some(),
|
|
"expected a non-empty error on cancel"
|
|
);
|
|
|
|
let maybe_updated = timeout(
|
|
Duration::from_millis(500),
|
|
mcp.read_stream_until_notification_message("account/updated"),
|
|
)
|
|
.await;
|
|
assert!(
|
|
maybe_updated.is_err(),
|
|
"account/updated should not be emitted when login is cancelled"
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
// Serialize tests that launch the login server since it binds to a fixed port.
|
|
#[serial(login_port)]
|
|
async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
forced_workspace_id: Some("ws-forced".to_string()),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let request_id = mcp.send_login_account_chatgpt_request().await?;
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
|
|
let login: LoginAccountResponse = to_response(resp)?;
|
|
let LoginAccountResponse::Chatgpt { auth_url, .. } = login else {
|
|
bail!("unexpected login response: {login:?}");
|
|
};
|
|
assert!(
|
|
auth_url.contains("allowed_workspace_id=ws-forced"),
|
|
"auth URL should include forced workspace"
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_account_no_auth() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
requires_openai_auth: Some(true),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let params = GetAccountParams {
|
|
refresh_token: false,
|
|
};
|
|
let request_id = mcp.send_get_account_request(params).await?;
|
|
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
let account: GetAccountResponse = to_response(resp)?;
|
|
|
|
assert_eq!(account.account, None, "expected no account");
|
|
assert_eq!(account.requires_openai_auth, true);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_account_with_api_key() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
requires_openai_auth: Some(true),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let req_id = mcp
|
|
.send_login_account_api_key_request("sk-test-key")
|
|
.await?;
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(req_id)),
|
|
)
|
|
.await??;
|
|
let _login_ok = to_response::<LoginAccountResponse>(resp)?;
|
|
|
|
let params = GetAccountParams {
|
|
refresh_token: false,
|
|
};
|
|
let request_id = mcp.send_get_account_request(params).await?;
|
|
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
let received: GetAccountResponse = to_response(resp)?;
|
|
|
|
let expected = GetAccountResponse {
|
|
account: Some(Account::ApiKey {}),
|
|
requires_openai_auth: true,
|
|
};
|
|
assert_eq!(received, expected);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_account_when_auth_not_required() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
requires_openai_auth: Some(false),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new(codex_home.path()).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let params = GetAccountParams {
|
|
refresh_token: false,
|
|
};
|
|
let request_id = mcp.send_get_account_request(params).await?;
|
|
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
let received: GetAccountResponse = to_response(resp)?;
|
|
|
|
let expected = GetAccountResponse {
|
|
account: None,
|
|
requires_openai_auth: false,
|
|
};
|
|
assert_eq!(received, expected);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_account_with_chatgpt() -> Result<()> {
|
|
let codex_home = TempDir::new()?;
|
|
create_config_toml(
|
|
codex_home.path(),
|
|
CreateConfigTomlParams {
|
|
requires_openai_auth: Some(true),
|
|
..Default::default()
|
|
},
|
|
)?;
|
|
write_chatgpt_auth(
|
|
codex_home.path(),
|
|
ChatGptAuthFixture::new("access-chatgpt")
|
|
.email("user@example.com")
|
|
.plan_type("pro"),
|
|
AuthCredentialsStoreMode::File,
|
|
)?;
|
|
|
|
let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
|
|
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
|
|
|
|
let params = GetAccountParams {
|
|
refresh_token: false,
|
|
};
|
|
let request_id = mcp.send_get_account_request(params).await?;
|
|
|
|
let resp: JSONRPCResponse = timeout(
|
|
DEFAULT_READ_TIMEOUT,
|
|
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
|
|
)
|
|
.await??;
|
|
let received: GetAccountResponse = to_response(resp)?;
|
|
|
|
let expected = GetAccountResponse {
|
|
account: Some(Account::Chatgpt {
|
|
email: "user@example.com".to_string(),
|
|
plan_type: AccountPlanType::Pro,
|
|
}),
|
|
requires_openai_auth: true,
|
|
};
|
|
assert_eq!(received, expected);
|
|
Ok(())
|
|
}
|