fix(mcp): preserve DML results and returning rows

agents/drivers/duckdb/src/query.rs

@@ -45,6 +45,22 @@ mod tests {
         assert_eq!(result.rows[2][0], serde_json::json!(99.999));
     }
 
+    #[test]
+    fn duckdb_execute_returns_dml_returning_rows() {
+        let con = duckdb::Connection::open_in_memory().expect("connect in-memory DuckDB");
+        con.execute_batch("CREATE TABLE users (id INTEGER, name VARCHAR)").expect("create table");
+
+        let inserted = duckdb_execute(&con, "INSERT INTO users VALUES (1, 'Ada') RETURNING id, name")
+            .expect("insert returning");
+        let updated = duckdb_execute(&con, "UPDATE users SET name = 'Ada Lovelace' RETURNING id, name")
+            .expect("update returning");
+        let deleted = duckdb_execute(&con, "DELETE FROM users RETURNING id, name").expect("delete returning");
+
+        assert_eq!(inserted.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada")]]);
+        assert_eq!(updated.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada Lovelace")]]);
+        assert_eq!(deleted.rows, vec![vec![serde_json::json!(1), serde_json::json!("Ada Lovelace")]]);
+    }
+
     #[test]
     fn duckdb_execute_returns_rows_for_from_first_query() {
         let con = duckdb::Connection::open_in_memory().expect("connect in-memory DuckDB");

agents/drivers/duckdb/src/sql.rs

@@ -202,6 +202,62 @@ pub fn starts_with_duckdb_result_sql_keyword(sql: &str) -> bool {
         .any(|keyword| {
             token.eq_ignore_ascii_case(keyword) || (*keyword == "DESCRIBE" && token.eq_ignore_ascii_case("DESC"))
         })
+        || matches!(token.to_ascii_uppercase().as_str(), "INSERT" | "UPDATE" | "DELETE" | "MERGE")
+            && contains_unquoted_sql_keyword(sql, "RETURNING")
+}
+
+fn contains_unquoted_sql_keyword(sql: &str, keyword: &str) -> bool {
+    let bytes = sql.as_bytes();
+    let mut index = 0;
+    let mut in_single = false;
+    let mut in_double = false;
+
+    while index < bytes.len() {
+        match bytes[index] {
+            b'\'' if !in_double => {
+                if in_single && bytes.get(index + 1) == Some(&b'\'') {
+                    index += 2;
+                    continue;
+                }
+                in_single = !in_single;
+                index += 1;
+            }
+            b'"' if !in_single => {
+                if in_double && bytes.get(index + 1) == Some(&b'"') {
+                    index += 2;
+                    continue;
+                }
+                in_double = !in_double;
+                index += 1;
+            }
+            b'-' if !in_single && !in_double && bytes.get(index + 1) == Some(&b'-') => {
+                index += 2;
+                while index < bytes.len() && bytes[index] != b'\n' {
+                    index += 1;
+                }
+            }
+            b'/' if !in_single && !in_double && bytes.get(index + 1) == Some(&b'*') => {
+                index += 2;
+                while index + 1 < bytes.len() && !(bytes[index] == b'*' && bytes[index + 1] == b'/') {
+                    index += 1;
+                }
+                index = (index + 2).min(bytes.len());
+            }
+            byte if !in_single && !in_double && (byte.is_ascii_alphabetic() || byte == b'_') => {
+                let start = index;
+                index += 1;
+                while index < bytes.len() && (bytes[index].is_ascii_alphanumeric() || bytes[index] == b'_') {
+                    index += 1;
+                }
+                if sql[start..index].eq_ignore_ascii_case(keyword) {
+                    return true;
+                }
+            }
+            _ => index += 1,
+        }
+    }
+
+    false
 }
 
 fn first_executable_sql_token(sql: &str) -> Option<&str> {
@@ -269,6 +325,8 @@ mod tests {
     fn detects_result_statements_after_comments() {
         assert!(starts_with_duckdb_result_sql_keyword("/* note */ WITH rows AS (SELECT 1) SELECT * FROM rows"));
         assert!(starts_with_duckdb_result_sql_keyword("DESC SELECT 1"));
+        assert!(starts_with_duckdb_result_sql_keyword("INSERT INTO items VALUES (1) RETURNING id"));
+        assert!(!starts_with_duckdb_result_sql_keyword("INSERT INTO items(note) VALUES ('RETURNING')"));
         assert!(!starts_with_duckdb_result_sql_keyword("INSERT INTO items VALUES (1)"));
     }
 }