From 3cbc28407692b9e09184cf87e624a41bd231b375 Mon Sep 17 00:00:00 2001
From: Chris Tate <chris@ctate.dev>
Date: Fri, 6 Mar 2026 12:46:03 -0600
Subject: [PATCH] fix: persist auth cookies on close in native mode (#650)

(cherry picked from commit b7e7a2548e54464e306deccc332dd69b3ba32d86)
---
 cli/src/native/browser.rs    |   9 ++-
 cli/src/native/cdp/chrome.rs |  18 ++++++
 cli/src/native/e2e_tests.rs  | 112 +++++++++++++++++++++++++++++++++++
 3 files changed, 136 insertions(+), 3 deletions(-)

diff --git a/cli/src/native/browser.rs b/cli/src/native/browser.rs
index c518f3e..c69036c 100644
--- a/cli/src/native/browser.rs
+++ b/cli/src/native/browser.rs
@@ -507,9 +507,12 @@ impl BrowserManager {
             .send_command_no_params("Browser.close", None)
             .await;
 
-        // Kill Chrome process if we own it
-        if let Some(ref mut chrome) = self.chrome_process {
-            chrome.kill();
+        if let Some(mut chrome) = self.chrome_process.take() {
+            let timeout = std::time::Duration::from_secs(5);
+            let _ = tokio::task::spawn_blocking(move || {
+                chrome.wait_or_kill(timeout);
+            })
+            .await;
         }
 
         Ok(())
diff --git a/cli/src/native/cdp/chrome.rs b/cli/src/native/cdp/chrome.rs
index 8b49359..d340169 100644
--- a/cli/src/native/cdp/chrome.rs
+++ b/cli/src/native/cdp/chrome.rs
@@ -16,6 +16,24 @@ impl ChromeProcess {
         let _ = self.child.kill();
         let _ = self.child.wait();
     }
+
+    /// Wait for Chrome to exit on its own (after Browser.close CDP command),
+    /// falling back to kill() if it doesn't exit within the timeout.
+    /// This allows Chrome to flush cookies and other state to the user-data-dir.
+    pub fn wait_or_kill(&mut self, timeout: Duration) {
+        let start = std::time::Instant::now();
+        let poll_interval = Duration::from_millis(50);
+
+        while start.elapsed() < timeout {
+            match self.child.try_wait() {
+                Ok(Some(_)) => return,
+                Ok(None) => std::thread::sleep(poll_interval),
+                Err(_) => break,
+            }
+        }
+
+        self.kill();
+    }
 }
 
 impl Drop for ChromeProcess {
diff --git a/cli/src/native/e2e_tests.rs b/cli/src/native/e2e_tests.rs
index 7b3a84e..eb4820d 100644
--- a/cli/src/native/e2e_tests.rs
+++ b/cli/src/native/e2e_tests.rs
@@ -1293,3 +1293,115 @@ async fn e2e_error_handling() {
     let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
     assert_success(&resp);
 }
+
+// ---------------------------------------------------------------------------
+// Profile cookie persistence across restarts
+// ---------------------------------------------------------------------------
+
+#[tokio::test]
+#[ignore]
+async fn e2e_profile_cookie_persistence() {
+    let profile_dir = std::env::temp_dir().join(format!(
+        "agent-browser-e2e-profile-{}",
+        uuid::Uuid::new_v4()
+    ));
+
+    // Session 1: launch with profile, set a cookie, close
+    {
+        let mut state = DaemonState::new();
+
+        let resp = execute_command(
+            &json!({
+                "id": "1",
+                "action": "launch",
+                "headless": true,
+                "profile": profile_dir.to_str().unwrap()
+            }),
+            &mut state,
+        )
+        .await;
+        assert_success(&resp);
+
+        let resp = execute_command(
+            &json!({ "id": "2", "action": "navigate", "url": "https://example.com" }),
+            &mut state,
+        )
+        .await;
+        assert_success(&resp);
+
+        let resp = execute_command(
+            &json!({
+                "id": "3",
+                "action": "cookies_set",
+                "name": "persist_test",
+                "value": "should_survive_restart",
+                "domain": ".example.com",
+                "path": "/",
+                "expires": 2000000000
+            }),
+            &mut state,
+        )
+        .await;
+        assert_success(&resp);
+
+        // Verify cookie is set
+        let resp =
+            execute_command(&json!({ "id": "4", "action": "cookies_get" }), &mut state).await;
+        assert_success(&resp);
+        let cookies = get_data(&resp)["cookies"].as_array().unwrap();
+        let found = cookies
+            .iter()
+            .any(|c| c["name"] == "persist_test" && c["value"] == "should_survive_restart");
+        assert!(found, "Cookie should exist before close");
+
+        let resp = execute_command(&json!({ "id": "5", "action": "close" }), &mut state).await;
+        assert_success(&resp);
+    }
+
+    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+
+    // Session 2: reopen with the same profile, verify cookie persisted
+    {
+        let mut state = DaemonState::new();
+
+        let resp = execute_command(
+            &json!({
+                "id": "10",
+                "action": "launch",
+                "headless": true,
+                "profile": profile_dir.to_str().unwrap()
+            }),
+            &mut state,
+        )
+        .await;
+        assert_success(&resp);
+
+        let resp = execute_command(
+            &json!({ "id": "11", "action": "navigate", "url": "https://example.com" }),
+            &mut state,
+        )
+        .await;
+        assert_success(&resp);
+
+        let resp =
+            execute_command(&json!({ "id": "12", "action": "cookies_get" }), &mut state).await;
+        assert_success(&resp);
+        let cookies = get_data(&resp)["cookies"].as_array().unwrap();
+        let found = cookies
+            .iter()
+            .any(|c| c["name"] == "persist_test" && c["value"] == "should_survive_restart");
+        assert!(
+            found,
+            "Cookie should persist across restart with --profile. Cookies found: {:?}",
+            cookies
+                .iter()
+                .map(|c| c["name"].as_str().unwrap_or("?"))
+                .collect::<Vec<_>>()
+        );
+
+        let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
+        assert_success(&resp);
+    }
+
+    let _ = std::fs::remove_dir_all(&profile_dir);
+}