diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3ad472e33c..43ebc805db 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2529,7 +2529,6 @@ dependencies = [ "codex-app-server-protocol", "codex-apply-patch", "codex-async-utils", - "codex-client", "codex-code-mode", "codex-config", "codex-connectors", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 9139f394a1..25ebfda72f 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,7 +26,6 @@ chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } codex-api = { workspace = true } -codex-client = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 74dd9860d7..bef48bbc0a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -185,6 +185,7 @@ struct ModelClientState { attestation_provider: Option>, http_state: Option, disable_websockets: AtomicBool, + websocket_session_generation: AtomicU64, cached_websocket_session: StdMutex, } @@ -241,6 +242,7 @@ pub struct ModelClient { /// contract and can cause routing bugs. pub struct ModelClientSession { client: ModelClient, + websocket_session_generation: u64, websocket_session: WebsocketSession, /// Turn state for sticky routing. /// @@ -365,6 +367,7 @@ impl ModelClient { attestation_provider, http_state, disable_websockets: AtomicBool::new(false), + websocket_session_generation: AtomicU64::new(0), cached_websocket_session: StdMutex::new(WebsocketSession::default()), }), prompt_cache_key_override: None, @@ -390,9 +393,12 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + let (websocket_session_generation, websocket_session) = + self.take_cached_websocket_session(); ModelClientSession { client: self.clone(), - websocket_session: self.take_cached_websocket_session(), + websocket_session_generation, + websocket_session, turn_state: Arc::new(OnceLock::new()), } } @@ -426,7 +432,7 @@ impl ModelClient { .as_ref() .is_some_and(|state| state.set_surface(surface)) { - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } } @@ -434,12 +440,12 @@ impl ModelClient { self.state .window_generation .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn advance_window_generation(&self) { self.state.window_generation.fetch_add(1, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn current_window_id(&self) -> String { @@ -448,21 +454,45 @@ impl ModelClient { format!("{thread_id}:{window_generation}") } - fn take_cached_websocket_session(&self) -> WebsocketSession { + fn take_cached_websocket_session(&self) -> (u64, WebsocketSession) { let mut cached_websocket_session = self .state .cached_websocket_session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - std::mem::take(&mut *cached_websocket_session) + let generation = self + .state + .websocket_session_generation + .load(Ordering::Relaxed); + (generation, std::mem::take(&mut *cached_websocket_session)) } - fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) { - *self + fn invalidate_cached_websocket_session(&self) { + let mut cached_websocket_session = self .state .cached_websocket_session .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session; + .unwrap_or_else(std::sync::PoisonError::into_inner); + self.state + .websocket_session_generation + .fetch_add(1, Ordering::Relaxed); + *cached_websocket_session = WebsocketSession::default(); + } + + fn store_cached_websocket_session(&self, generation: u64, websocket_session: WebsocketSession) { + let mut cached_websocket_session = self + .state + .cached_websocket_session + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if self + .state + .websocket_session_generation + .load(Ordering::Relaxed) + == generation + { + *cached_websocket_session = websocket_session; + } } pub(crate) fn force_http_fallback( @@ -482,7 +512,7 @@ impl ModelClient { ); } - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); activated } @@ -997,7 +1027,7 @@ impl Drop for ModelClientSession { fn drop(&mut self) { let websocket_session = std::mem::take(&mut self.websocket_session); self.client - .store_cached_websocket_session(websocket_session); + .store_cached_websocket_session(self.websocket_session_generation, websocket_session); } } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index f8d904036a..419b28f2d2 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,12 +13,15 @@ use crate::GenerateAttestationFuture; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; +use codex_http_state::HttpStateStore; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; +use codex_model_provider_info::built_in_model_providers; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; use codex_protocol::SessionId; @@ -67,11 +70,32 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { fn test_model_client_with_parent( session_source: SessionSource, parent_thread_id: Option, +) -> ModelClient { + test_model_client_with_parent_and_auth( + session_source, + parent_thread_id, + /*auth_manager*/ None, + ) +} + +fn test_model_client_with_parent_and_auth( + session_source: SessionSource, + parent_thread_id: Option, + auth_manager: Option>, ) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); + test_model_client_with_provider(session_source, parent_thread_id, auth_manager, provider) +} + +fn test_model_client_with_provider( + session_source: SessionSource, + parent_thread_id: Option, + auth_manager: Option>, + provider: ModelProviderInfo, +) -> ModelClient { let thread_id = ThreadId::new(); ModelClient::new( - /*auth_manager*/ None, + auth_manager, thread_id.into(), thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), @@ -322,6 +346,77 @@ fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { ); } +#[tokio::test] +async fn current_client_setup_attaches_and_rotates_native_http_state() { + let codex_home = TempDir::new().expect("tempdir"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set(HttpStateSurface::CodexTui, "old-state".to_string()) + .expect("state should store"); + let auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + codex_home.path().to_path_buf(), + ); + let client = test_model_client_with_provider( + SessionSource::Cli, + /*parent_thread_id*/ None, + Some(auth_manager), + built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(), + ); + let setup = client + .current_client_setup() + .await + .expect("client setup should resolve"); + let request_url = "https://chatgpt.com/backend-api/codex/responses"; + let mut request_headers = http::HeaderMap::new(); + setup + .api_auth + .add_auth_headers_for_url(request_url, &mut request_headers); + assert_eq!( + request_headers + .get("x-oai-is") + .and_then(|value| value.to_str().ok()), + Some("old-state"), + ); + + let mut response_headers = http::HeaderMap::new(); + response_headers.insert( + "x-oai-is-update", + http::HeaderValue::from_static("new-state"), + ); + setup + .api_auth + .observe_response_headers(request_url, &request_headers, &response_headers); + assert_eq!( + store + .get(HttpStateSurface::CodexTui) + .expect("state should load"), + Some("new-state".to_string()), + ); +} + +#[test] +fn surface_change_discards_checked_out_websocket_session() { + let codex_home = TempDir::new().expect("tempdir"); + let auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + codex_home.path().to_path_buf(), + ); + let client = test_model_client_with_parent_and_auth( + SessionSource::Cli, + /*parent_thread_id*/ None, + Some(auth_manager), + ); + let mut session = client.new_session(); + session.websocket_session.last_response_from_untraced_warmup = true; + + client.set_http_state_surface(HttpStateSurface::CodexDesktop); + drop(session); + + let (_, cached_session) = client.take_cached_websocket_session(); + assert!(!cached_session.last_response_from_untraced_warmup); +} + #[tokio::test] async fn summarize_memories_returns_empty_for_empty_input() { let client = test_model_client(SessionSource::Cli); diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 34f480b861..ddf1c0bbcc 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -130,7 +130,9 @@ pub fn with_native_integrity_state( let Some(state) = state else { return auth_provider; }; - let Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_)) = auth else { + let Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_)) = + auth + else { return auth_provider; }; Arc::new(NativeIntegrityAuthProvider::new(auth_provider, state))