fix(sqlite): run remote batch saves in a worker transaction

crates/dbx-core/src/query/mod.rs

@@ -5306,6 +5306,9 @@ async fn exec_tx_sqlite_inner(
 ) -> Result<db::QueryResult, String> {
     let statements = statements.to_vec();
     let query_timeout = budget.query_timeout;
+    if let Some(worker) = pool.worker() {
+        return exec_tx_sqlite_worker_inner(worker, statements, start, query_timeout).await;
+    }
     tokio::task::spawn_blocking(move || {
         pool.with_connection(|conn| {
             conn.execute_batch("BEGIN").map_err(|e| format!("Failed to begin transaction: {}", e))?;
@@ -5476,6 +5479,55 @@ async fn exec_tx_sqlite_inner(
     .map_err(|e| e.to_string())?
 }
 
+async fn exec_tx_sqlite_worker_inner(
+    worker: Arc<db::sqlite_worker::SqliteWorkerClient>,
+    statements: Vec<String>,
+    start: std::time::Instant,
+    query_timeout: Option<Duration>,
+) -> Result<db::QueryResult, String> {
+    // Detached like the local spawn_blocking path, so a dropped caller cannot leave the
+    // worker connection inside an open transaction. One session keeps other requests out.
+    tokio::spawn(async move {
+        let mut session = worker.session().await;
+        session.query("BEGIN", None).await.map_err(|e| format!("Failed to begin transaction: {e}"))?;
+        // ponytail: the worker cannot interrupt a running statement, so the budget is only
+        // checked between statements; bounding one slow statement needs a worker interrupt op.
+        let within_budget = || match query_timeout {
+            Some(timeout) if start.elapsed() >= timeout => {
+                Err(format!("Query timed out after {} seconds", timeout.as_secs()))
+            }
+            _ => Ok(()),
+        };
+        let outcome: Result<db::QueryResult, String> = async {
+            let mut total_affected = 0;
+            for (i, sql) in statements.iter().enumerate() {
+                within_budget()?;
+                total_affected += session
+                    .query(sql, None)
+                    .await
+                    .map_err(|e| {
+                        query_error_with_omitted_sql_context(&format!("Statement {} failed: {}", i + 1, e), sql)
+                    })?
+                    .affected_rows;
+            }
+            within_budget()?;
+            let committed = session.query("COMMIT", None).await.map_err(|e| format!("COMMIT failed: {e}"))?;
+            Ok(db::QueryResult {
+                affected_rows: total_affected,
+                execution_time_ms: start.elapsed().as_millis(),
+                ..committed
+            })
+        }
+        .await;
+        if outcome.is_err() {
+            let _ = session.query("ROLLBACK", None).await;
+        }
+        outcome
+    })
+    .await
+    .map_err(|e| e.to_string())?
+}
+
 async fn exec_tx_explicit_inner(
     state: &AppState,
     pool_key: &str,
@@ -7304,6 +7356,61 @@ mod tests {
         assert!(error.contains("Statement 1 failed") && error.contains("interrupted_at"), "unexpected error: {error}");
     }
 
+    #[tokio::test]
+    async fn sqlite_ssh_worker_transaction_rolls_back_on_failure_and_commits_on_success() {
+        let remote = Arc::new(Mutex::new(rusqlite::Connection::open_in_memory().expect("open remote SQLite")));
+        remote.lock().unwrap().execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY)").expect("create table");
+        let (client_stream, worker_stream) = tokio::io::duplex(64 * 1024);
+        tokio::spawn(fake_sqlite_ssh_worker(worker_stream, remote.clone()));
+        let pool = db::sqlite::SqliteHandle::from_worker(Arc::new(
+            db::sqlite_worker::SqliteWorkerClient::from_test_stream(client_stream),
+        ));
+        let budget = DbOperationBudget::with_defaults();
+        let row_count =
+            || remote.lock().unwrap().query_row("SELECT COUNT(*) FROM t", [], |row| row.get::<_, i64>(0)).unwrap();
+
+        let error = exec_tx_sqlite_inner(
+            pool.clone(),
+            &["INSERT INTO t VALUES (1)".to_string(), "INSERT INTO missing VALUES (2)".to_string()],
+            std::time::Instant::now(),
+            &budget,
+        )
+        .await
+        .expect_err("second statement fails");
+        assert!(error.contains("Statement 2 failed") && error.contains("no such table: missing"), "{error}");
+        assert_eq!(row_count(), 0);
+
+        let result = exec_tx_sqlite_inner(
+            pool,
+            &["INSERT INTO t VALUES (1)".to_string(), "INSERT INTO t VALUES (2)".to_string()],
+            std::time::Instant::now(),
+            &budget,
+        )
+        .await
+        .expect("transaction commits");
+        assert_eq!(result.affected_rows, 2);
+        assert_eq!(row_count(), 2);
+    }
+
+    /// Answers SQLite worker JSONL requests from a real SQLite connection.
+    async fn fake_sqlite_ssh_worker(stream: tokio::io::DuplexStream, conn: Arc<Mutex<rusqlite::Connection>>) {
+        use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
+        let (reader, mut writer) = tokio::io::split(stream);
+        let mut lines = tokio::io::BufReader::new(reader).lines();
+        while let Ok(Some(line)) = lines.next_line().await {
+            let request: serde_json::Value = serde_json::from_str(&line).expect("worker request");
+            let mut response = {
+                let conn = conn.lock().unwrap();
+                match conn.execute_batch(request["sql"].as_str().unwrap_or_default()) {
+                    Ok(()) => serde_json::json!({ "affected_rows": conn.changes() }),
+                    Err(error) => serde_json::json!({ "error": error.to_string() }),
+                }
+            };
+            response["id"] = request["id"].clone();
+            writer.write_all(format!("{response}\n").as_bytes()).await.expect("worker response");
+        }
+    }
+
     #[tokio::test]
     async fn sqlite_external_interrupt_without_budget_does_not_panic() {
         // Regression: with query_timeout = None (MCP policy can set it to 0 =

crates/dbx-drivers/src/db/sqlite_worker.rs

@@ -102,8 +102,15 @@ enum WorkerIo {
     Closed,
 }
 
-impl SqliteWorkerClient {
-    pub async fn query(&self, sql: &str, max_rows: Option<usize>) -> Result<QueryResult, String> {
+/// Holds the worker I/O lock so consecutive requests, such as the statements of
+/// one transaction, cannot interleave with other callers on the worker's single connection.
+pub struct SqliteWorkerSession<'a> {
+    client: &'a SqliteWorkerClient,
+    io: tokio::sync::MutexGuard<'a, WorkerIo>,
+}
+
+impl SqliteWorkerSession<'_> {
+    pub async fn query(&mut self, sql: &str, max_rows: Option<usize>) -> Result<QueryResult, String> {
         match self.roundtrip(WorkerOp::Query { sql: sql.to_string(), max_rows }).await? {
             WorkerBody::Ok { columns, column_types, rows, affected_rows, truncated, .. } => Ok(QueryResult {
                 columns: columns.unwrap_or_default(),
@@ -125,6 +132,47 @@ impl SqliteWorkerClient {
         }
     }
 
+    async fn roundtrip(&mut self, op: WorkerOp) -> Result<WorkerBody, String> {
+        let id = self.client.next_id.fetch_add(1, Ordering::SeqCst);
+        let mut encoded = serde_json::to_vec(&WorkerRequest { id, op }).map_err(|e| e.to_string())?;
+        encoded.push(b'\n');
+        match &mut *self.io {
+            WorkerIo::Process { stdin, stdout, .. } => {
+                stdin.write_all(&encoded).await.map_err(|e| e.to_string())?;
+                stdin.flush().await.map_err(|e| e.to_string())?;
+                let mut line = String::new();
+                stdout.read_line(&mut line).await.map_err(|e| e.to_string())?;
+                parse_response(id, &line)
+            }
+            WorkerIo::Closed => Err("SQLite worker session is closed".to_string()),
+            WorkerIo::Ssh { stream, .. } => {
+                stream.write_all(&encoded).await.map_err(|e| e.to_string())?;
+                stream.flush().await.map_err(|e| e.to_string())?;
+                parse_response(id, &read_jsonl_line(stream).await?)
+            }
+        }
+    }
+}
+
+impl SqliteWorkerClient {
+    #[cfg(any(test, feature = "test-support"))]
+    pub fn from_test_stream(stream: impl AsyncRead + AsyncWrite + Send + Unpin + 'static) -> Self {
+        Self {
+            io: AsyncMutex::new(WorkerIo::Ssh { stream: BufReader::new(Box::pin(stream)) }),
+            next_id: AtomicU64::new(1),
+            ssh_session: None,
+            remove_remote_path: None,
+        }
+    }
+
+    pub async fn session(&self) -> SqliteWorkerSession<'_> {
+        SqliteWorkerSession { client: self, io: self.io.lock().await }
+    }
+
+    pub async fn query(&self, sql: &str, max_rows: Option<usize>) -> Result<QueryResult, String> {
+        self.session().await.query(sql, max_rows).await
+    }
+
     pub async fn backup(&self, dest: &str) -> Result<(), String> {
         match self.roundtrip(WorkerOp::Backup { dest: dest.to_string() }).await? {
             WorkerBody::Ok { .. } => Ok(()),
@@ -208,25 +256,7 @@ impl SqliteWorkerClient {
     }
 
     async fn roundtrip(&self, op: WorkerOp) -> Result<WorkerBody, String> {
-        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
-        let mut encoded = serde_json::to_vec(&WorkerRequest { id, op }).map_err(|e| e.to_string())?;
-        encoded.push(b'\n');
-        let mut io = self.io.lock().await;
-        match &mut *io {
-            WorkerIo::Process { stdin, stdout, .. } => {
-                stdin.write_all(&encoded).await.map_err(|e| e.to_string())?;
-                stdin.flush().await.map_err(|e| e.to_string())?;
-                let mut line = String::new();
-                stdout.read_line(&mut line).await.map_err(|e| e.to_string())?;
-                parse_response(id, &line)
-            }
-            WorkerIo::Closed => Err("SQLite worker session is closed".to_string()),
-            WorkerIo::Ssh { stream, .. } => {
-                stream.write_all(&encoded).await.map_err(|e| e.to_string())?;
-                stream.flush().await.map_err(|e| e.to_string())?;
-                parse_response(id, &read_jsonl_line(stream).await?)
-            }
-        }
+        self.session().await.roundtrip(op).await
     }
 }
 

crates/dbx-sqlite-worker/src/protocol.rs

@@ -28,6 +28,11 @@ pub struct WorkerResponse {
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[serde(untagged)]
 pub enum WorkerBody {
+    // Untagged variants are tried in order. `Ok` accepts any object because all of
+    // its fields are optional, so `Err` must come first or error payloads decode as success.
+    Err {
+        error: String,
+    },
     Ok {
         #[serde(default, skip_serializing_if = "Option::is_none")]
         columns: Option<Vec<String>>,
@@ -42,9 +47,6 @@ pub enum WorkerBody {
         #[serde(default, skip_serializing_if = "Option::is_none")]
         pong: Option<bool>,
     },
-    Err {
-        error: String,
-    },
 }
 
 impl WorkerBody {
@@ -97,4 +99,17 @@ mod tests {
         let decoded: WorkerRequest = serde_json::from_str(&encoded).unwrap();
         assert_eq!(decoded, request);
     }
+
+    #[test]
+    fn response_roundtrip_keeps_errors_distinct_from_success() {
+        for body in [
+            WorkerBody::err("no such table: t"),
+            WorkerBody::ok(),
+            WorkerBody::query(vec!["id".into()], Vec::new(), vec![vec![serde_json::json!(1)]], 0, false),
+        ] {
+            let encoded = serde_json::to_string(&WorkerResponse { id: 7, body: body.clone() }).unwrap();
+            let decoded: WorkerResponse = serde_json::from_str(&encoded).unwrap();
+            assert_eq!(decoded.body, body, "{encoded}");
+        }
+    }
 }