[codex] guard native websocket state caches [ci changed_files]

codex-rs/Cargo.lock

@@ -2526,7 +2526,6 @@ dependencies = [
  "codex-app-server-protocol",
  "codex-apply-patch",
  "codex-async-utils",
- "codex-client",
  "codex-code-mode",
  "codex-config",
  "codex-connectors",

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 }

codex-rs/core/src/client.rs

@@ -185,6 +185,7 @@ struct ModelClientState {
     attestation_provider: Option<Arc<dyn AttestationProvider>>,
     http_state: Option<HttpStateContext>,
     disable_websockets: AtomicBool,
+    websocket_session_generation: AtomicU64,
     cached_websocket_session: StdMutex<WebsocketSession>,
 }
 
@@ -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);
     }
 }
 

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<ThreadId>,
+) -> 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<ThreadId>,
+    auth_manager: Option<Arc<AuthManager>>,
 ) -> 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<ThreadId>,
+    auth_manager: Option<Arc<AuthManager>>,
+    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);

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))