diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 4c61447ac0..f9a4306eae 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -193,10 +193,14 @@ async fn mock_device_code_token_failure(server: &MockServer, status: u16) { async fn mock_device_code_oauth_token(server: &MockServer, id_token: &str) { Mock::given(method("POST")) .and(path("/oauth/token")) + .and(wiremock::matchers::body_string_contains( + "surface=codex_desktop_ssh", + )) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id_token": id_token, "access_token": "access-token-123", "refresh_token": "refresh-token-123", + "state": INTEGRITY_STATE, }))) .mount(server) .await; @@ -1142,7 +1146,11 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> ], ) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_client_info(client_info("codex_desktop_ssh")), + ) + .await??; let request_id = mcp.send_login_account_chatgpt_device_code_request().await?; let resp: JSONRPCResponse = timeout( @@ -1190,6 +1198,11 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> codex_home.path().join("auth.json").exists(), "auth.json should be created when device code login succeeds" ); + assert_eq!( + HttpStateStore::new(codex_home.path().to_path_buf()) + .get(HttpStateSurface::CodexDesktopSsh)?, + Some(INTEGRITY_STATE.to_string()), + ); Ok(()) } diff --git a/codex-rs/login/src/http_state.rs b/codex-rs/login/src/http_state.rs index afc5656840..060897a98d 100644 --- a/codex-rs/login/src/http_state.rs +++ b/codex-rs/login/src/http_state.rs @@ -28,42 +28,5 @@ pub(crate) fn clear_all(codex_home: &Path) { } #[cfg(test)] -mod tests { - use super::*; - use tempfile::TempDir; - - #[test] - fn replace_after_login_sets_and_clears_surface_state() { - let codex_home = TempDir::new().expect("tempdir"); - let store = HttpStateStore::new(codex_home.path().to_path_buf()); - store - .set(HttpStateSurface::CodexCli, "stale-cli-state".to_string()) - .expect("CLI state should store"); - - replace_after_login( - codex_home.path(), - HttpStateSurface::CodexDesktop, - Some("minted-state".to_string()), - ); - assert_eq!( - store - .get(HttpStateSurface::CodexCli) - .expect("CLI state should load"), - None, - ); - assert_eq!( - store - .get(HttpStateSurface::CodexDesktop) - .expect("state should load"), - Some("minted-state".to_string()), - ); - - replace_after_login(codex_home.path(), HttpStateSurface::CodexDesktop, None); - assert_eq!( - store - .get(HttpStateSurface::CodexDesktop) - .expect("state should load"), - None, - ); - } -} +#[path = "http_state_tests.rs"] +mod tests; diff --git a/codex-rs/login/src/http_state_tests.rs b/codex-rs/login/src/http_state_tests.rs new file mode 100644 index 0000000000..df68db1c9e --- /dev/null +++ b/codex-rs/login/src/http_state_tests.rs @@ -0,0 +1,38 @@ +use super::*; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +#[test] +fn replace_after_login_sets_and_clears_surface_state() { + let codex_home = TempDir::new().expect("tempdir"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set(HttpStateSurface::CodexCli, "stale-cli-state".to_string()) + .expect("CLI state should store"); + + replace_after_login( + codex_home.path(), + HttpStateSurface::CodexDesktop, + Some("minted-state".to_string()), + ); + assert_eq!( + store + .get(HttpStateSurface::CodexCli) + .expect("CLI state should load"), + None, + ); + assert_eq!( + store + .get(HttpStateSurface::CodexDesktop) + .expect("state should load"), + Some("minted-state".to_string()), + ); + + replace_after_login(codex_home.path(), HttpStateSurface::CodexDesktop, None); + assert_eq!( + store + .get(HttpStateSurface::CodexDesktop) + .expect("state should load"), + None, + ); +} diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index d1c5508a47..5f0354040b 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -94,6 +94,8 @@ impl ServerOptions { forced_chatgpt_workspace_id, codex_streamlined_login: false, cli_auth_credentials_store_mode, + // Direct `codex login` is a CLI request. App-server callers override + // this with the initiating native surface. http_state_surface: HttpStateSurface::CodexCli, } } diff --git a/codex-rs/login/tests/suite/device_code_login.rs b/codex-rs/login/tests/suite/device_code_login.rs index 5a738c9292..94231dd461 100644 --- a/codex-rs/login/tests/suite/device_code_login.rs +++ b/codex-rs/login/tests/suite/device_code_login.rs @@ -9,6 +9,7 @@ use codex_http_state::HttpStateSurface; use codex_login::ServerOptions; use codex_login::auth::load_auth_dot_json; use codex_login::run_device_code_login; +use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; use std::sync::atomic::AtomicUsize; diff --git a/codex-rs/login/tests/suite/login_server_e2e.rs b/codex-rs/login/tests/suite/login_server_e2e.rs index 0cf511b662..e531867d93 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -26,7 +26,10 @@ const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; // See spawn.rs for details -fn start_mock_issuer(chatgpt_account_id: &str) -> (SocketAddr, thread::JoinHandle<()>) { +fn start_mock_issuer( + chatgpt_account_id: &str, + http_state_surface: HttpStateSurface, +) -> (SocketAddr, thread::JoinHandle<()>) { // Bind to a random available port let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); let addr = listener.local_addr().unwrap(); @@ -40,6 +43,12 @@ fn start_mock_issuer(chatgpt_account_id: &str) -> (SocketAddr, thread::JoinHandl // Read body let mut body = String::new(); let _ = req.as_reader().read_to_string(&mut body); + if body.contains("grant_type=authorization_code") { + assert!( + body.contains(&format!("surface={http_state_surface}")), + "authorization-code exchange should identify the initiating surface" + ); + } // Build minimal JWT with plan=pro #[derive(serde::Serialize)] struct Header { @@ -95,7 +104,8 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { skip_if_no_network!(Ok(())); let chatgpt_account_id = "12345678-0000-0000-0000-000000000000"; - let (issuer_addr, issuer_handle) = start_mock_issuer(chatgpt_account_id); + let (issuer_addr, issuer_handle) = + start_mock_issuer(chatgpt_account_id, HttpStateSurface::CodexDesktop); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -178,7 +188,8 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { async fn creates_missing_codex_home_dir() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -222,7 +233,8 @@ async fn creates_missing_codex_home_dir() -> Result<()> { async fn login_server_includes_forced_workspaces_as_one_query_param() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -264,7 +276,8 @@ async fn login_server_includes_forced_workspaces_as_one_query_param() -> Result< async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_DISALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_DISALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -325,7 +338,8 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -394,7 +408,8 @@ async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error() async fn oauth_access_denied_unknown_reason_uses_generic_error_page() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -503,7 +518,8 @@ async fn falls_back_to_registered_fallback_port_when_default_port_is_in_use() -> }) }; - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -541,7 +557,8 @@ async fn falls_back_to_registered_fallback_port_when_default_port_is_in_use() -> async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + let (issuer_addr, _issuer_handle) = + start_mock_issuer(WORKSPACE_ID_ALLOWED, HttpStateSurface::CodexCli); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let first_tmp = tempdir()?;