fix(duckdb): checkpoint wal on worker shutdown

.gitignore

@@ -22,6 +22,7 @@ src-tauri/webview2-fixed-runtime/
 plugins/jdbc/target/
 plugins/jdbc/dependency-reduced-pom.xml
 plugins/jdbc/lib/*.jar
+agents/drivers/duckdb/target/
 plugins/jdbc/.gradle/
 plugins/jdbc/build/
 plugins/jdbc/dist/

agents/drivers/duckdb/src/connection.rs

@@ -23,7 +23,6 @@ impl DuckDbConnection {
     }
 
     // Preserve PoisonError ownership so close_connection can still release the contained DuckDB handle.
-    #[cfg(test)]
     #[allow(clippy::result_large_err)]
     fn into_inner(self) -> std::sync::LockResult<duckdb::Connection> {
         self.connection.into_inner()
@@ -116,8 +115,7 @@ pub fn is_memory_database_path(path: &str) -> bool {
 /// Unlike relying on Drop, this calls `duckdb_disconnect` synchronously
 /// so the file handle is released before this function returns.
 /// On Windows this prevents "file already in use" errors when reconnecting.
-#[cfg(test)]
-fn close_connection(con: Arc<DuckDbConnection>) {
+pub(crate) fn close_connection(con: Arc<DuckDbConnection>) {
     match Arc::try_unwrap(con) {
         Ok(handle) => match handle.into_inner() {
             Ok(conn) => {

agents/drivers/duckdb/src/runtime.rs

@@ -148,6 +148,27 @@ impl DuckDbWorkerSession {
         Ok(connection.interrupt_handle())
     }
 
+    /// Takes and closes the connection so DuckDB's shutdown checkpoint runs and
+    /// the database WAL is removed before this worker exits. Returns false —
+    /// leaving the connection untouched — when the connection is poisoned
+    /// (dropping it aborts the process) so the caller falls back to an
+    /// immediate exit and the WAL replays on the next open.
+    fn close_connection_for_shutdown(&mut self) -> bool {
+        if self.connection.is_none() {
+            return true;
+        }
+        if !self.is_connection_healthy() {
+            return false;
+        }
+        match self.connection.take() {
+            Some(connection) => {
+                crate::connection::close_connection(connection);
+                true
+            }
+            None => true,
+        }
+    }
+
     /// Probes whether the connection is still usable after an execute error.
     ///
     /// duckdb-rs 1.10503.1 has a bug where `Connection::prepare()` failing at the
@@ -205,6 +226,21 @@ struct WorkerHandleResult {
 }
 
 impl DuckDbWorkerRuntime {
+    /// Closes the session connection when the worker is idle and the connection
+    /// is healthy, so the DuckDB shutdown checkpoint removes the WAL. Returns
+    /// false when a query is running (the session is locked or an interrupt is
+    /// registered) or the connection is poisoned; callers then keep the legacy
+    /// immediate-exit behavior and rely on WAL replay at the next open.
+    fn close_session_for_shutdown(&self) -> bool {
+        if self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()).is_some() {
+            return false;
+        }
+        match self.session.try_lock() {
+            Ok(mut session) => session.close_connection_for_shutdown(),
+            Err(_) => false,
+        }
+    }
+
     async fn handle_request(
         &self,
         request: DuckDbWorkerRequest,
@@ -327,10 +363,18 @@ impl DuckDbWorkerRuntime {
                     shutdown: false,
                 }
             }
-            DuckDbWorkerMethod::Shutdown => WorkerHandleResult {
-                response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "shutdown": true }))),
-                shutdown: true,
-            },
+            DuckDbWorkerMethod::Shutdown => {
+                // Close the connection before acking: the checkpoint that removes
+                // the WAL must finish before the parent observes this response
+                // (or the process exit it waits for). Busy or poisoned sessions
+                // skip the close; the parent's kill fallback and WAL replay on
+                // the next open cover them.
+                self.close_session_for_shutdown();
+                WorkerHandleResult {
+                    response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "shutdown": true }))),
+                    shutdown: true,
+                }
+            }
         }
     }
 
@@ -377,7 +421,15 @@ pub async fn run_stdio_worker() -> Result<(), String> {
     loop {
         let line = lines.next_line().await.map_err(|e| e.to_string())?;
         let Some(line) = line else {
-            std::process::exit(0);
+            // Parent closed stdin (exit or crash). Idle workers checkpoint and
+            // remove their WAL on the normal return path; a worker still running
+            // a query must not wait for it (nor drop a poisoned connection),
+            // so it keeps the legacy immediate exit and the WAL replays on the
+            // next open.
+            if !runtime.close_session_for_shutdown() {
+                std::process::exit(0);
+            }
+            break;
         };
         if line.trim().is_empty() {
             continue;

agents/drivers/duckdb/tests/duckdb_worker_process.rs

@@ -499,6 +499,52 @@ async fn worker_process_retries_connect_after_transient_file_lock() {
     assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]);
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn worker_shutdown_checkpoints_and_removes_wal() {
+    let _guard = duckdb_worker_process_test_guard().await;
+    let executable = PathBuf::from(env!("CARGO_BIN_EXE_dbx-duckdb-driver"));
+    let db_path = temp_duckdb_path();
+    let wal_path = wal_path_for(&db_path);
+    let _ = std::fs::remove_file(&db_path);
+    let _ = std::fs::remove_file(&wal_path);
+
+    let client =
+        DuckDbWorkerClient::open_with_executable(executable.clone(), db_path.to_string_lossy().to_string(), Vec::new(), None)
+            .await
+            .expect("worker process connects");
+    client
+        .execute(
+            None,
+            "CREATE TABLE events (id INTEGER); INSERT INTO events VALUES (7);".to_string(),
+            Some(10),
+            None,
+            Some(Duration::from_secs(5)),
+        )
+        .await
+        .expect("write events");
+
+    // Committed writes sit in the WAL while the connection is open.
+    assert!(wal_path.exists(), "committed writes should leave a WAL file next to the database");
+
+    client.shutdown().await;
+
+    assert!(!wal_path.exists(), "shutdown should checkpoint and remove the WAL file");
+
+    // The checkpointed data survives a fresh open of the database file.
+    let client =
+        DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new(), None)
+            .await
+            .expect("reopen after shutdown");
+    let result = client
+        .execute(None, "SELECT id FROM events".to_string(), Some(10), None, Some(Duration::from_secs(5)))
+        .await
+        .expect("read events back");
+    assert_eq!(result.rows, vec![vec![serde_json::json!(7)]]);
+    client.shutdown().await;
+    assert!(!wal_path.exists(), "second shutdown should also remove the WAL file");
+    let _ = std::fs::remove_file(&db_path);
+}
+
 async fn write_worker_request(stdin: &mut tokio::process::ChildStdin, request: DuckDbWorkerRequest) {
     let line = serde_json::to_string(&request).expect("request json");
     stdin.write_all(line.as_bytes()).await.expect("write request");
@@ -517,6 +563,10 @@ async fn read_worker_response(
     serde_json::from_str(&line).expect("parse worker response")
 }
 
+fn wal_path_for(db_path: &PathBuf) -> PathBuf {
+    PathBuf::from(format!("{}.wal", db_path.display()))
+}
+
 fn temp_duckdb_path() -> PathBuf {
     let suffix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
     let pid = std::process::id();

crates/dbx-core/src/db/duckdb_worker_process.rs

@@ -28,6 +28,10 @@ const DUCKDB_WORKER_POISONED_CODE: &str = "duckdb_worker_poisoned";
 const DUCKDB_WORKER_REQUEST_TIMEOUT_CODE: &str = "duckdb_worker_request_timeout";
 const DEFAULT_WORKER_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
 const DEFAULT_WORKER_KILL_WAIT: Duration = Duration::from_secs(3);
+/// Grace period for a shutdown-acked worker to exit by itself (session drop
+/// plus DuckDB shutdown checkpoint). Kept below the 3s pool close timeout in
+/// `connection.rs` so pool close cancellation cannot cut the checkpoint short.
+const WORKER_SHUTDOWN_EXIT_WAIT: Duration = Duration::from_millis(2500);
 const DEFAULT_WORKER_START_WAIT: Duration = Duration::from_secs(5);
 pub const DUCKDB_DRIVER_PATH_ENV: &str = "DBX_DUCKDB_DRIVER_PATH";
 type PendingRequests = Arc<Mutex<HashMap<String, PendingRequest>>>;
@@ -362,7 +366,61 @@ impl DuckDbWorkerClient {
                 Some(self.inner.request_timeout),
             )
             .await;
-        self.kill().await;
+        self.wait_for_exit_or_kill().await;
+    }
+
+    /// Lets a shutdown-acked worker exit on its own so its session drop runs
+    /// the DuckDB shutdown checkpoint that removes the database WAL; SIGKILL
+    /// stays the fallback for stuck workers. Killing immediately after the
+    /// ack would truncate that checkpoint and leave the WAL file behind.
+    /// The wait must fit inside the caller's pool close timeout so the
+    /// detached close task is not cancelled mid-checkpoint.
+    async fn wait_for_exit_or_kill(&self) {
+        let (child, generation) = {
+            let mut state = self.inner.state.lock().await;
+            let child = state.child.take();
+            let generation = state.generation;
+            state.stdin = None;
+            state.connected = false;
+            (child, generation)
+        };
+
+        let Some(mut child) = child else { return };
+        match tokio::time::timeout(WORKER_SHUTDOWN_EXIT_WAIT, child.wait()).await {
+            Ok(Ok(status)) => {
+                log::info!("[duckdb-worker:shutdown:exit] status={status}");
+                self.fail_pending_for_generation(
+                    generation,
+                    "duckdb_worker_killed",
+                    "DuckDB worker process was killed",
+                )
+                .await;
+                return;
+            }
+            Ok(Err(err)) => {
+                log::warn!("[duckdb-worker:shutdown:wait-failed] error={err}");
+                self.fail_pending_for_generation(
+                    generation,
+                    "duckdb_worker_killed",
+                    "DuckDB worker process was killed",
+                )
+                .await;
+                return;
+            }
+            Err(_) => {
+                log::warn!("[duckdb-worker:shutdown:wait-timeout] wait_ms={}", WORKER_SHUTDOWN_EXIT_WAIT.as_millis());
+            }
+        }
+
+        let _ = child.start_kill();
+        match tokio::time::timeout(DEFAULT_WORKER_KILL_WAIT, child.wait()).await {
+            Ok(Ok(status)) => log::info!("[duckdb-worker:kill:exit] status={status}"),
+            Ok(Err(err)) => log::warn!("[duckdb-worker:kill:wait-failed] error={err}"),
+            Err(_) => {
+                log::warn!("[duckdb-worker:kill:wait-timeout] wait_ms={}", DEFAULT_WORKER_KILL_WAIT.as_millis())
+            }
+        }
+        self.fail_pending_for_generation(generation, "duckdb_worker_killed", "DuckDB worker process was killed").await;
     }
 
     pub async fn kill(&self) {