feat(agent): track execution context activity

apps/desktop/src/stores/queryStore.ts

@@ -1211,6 +1211,7 @@ export const useQueryStore = defineStore("query", () => {
     }
     const countSql = options.countSql;
     const clientSessionId = tabClientSessionId({ id: options.tabId }, "count");
+    const countExecutionId = `${options.executionId}:count`;
 
     if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) {
       setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount);
@@ -1220,7 +1221,7 @@ export const useQueryStore = defineStore("query", () => {
     void (async () => {
       try {
         console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() });
-        const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, undefined, {
+        const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, countExecutionId, {
           clientSessionId,
           timeoutSecs: options.timeoutSecs,
         });
@@ -2040,6 +2041,7 @@ export const useQueryStore = defineStore("query", () => {
       let executionTimeMs = 0;
       let offset = 0;
       const clientSessionId = tabClientSessionId(tab, "export");
+      const exportExecutionId = uuid();
 
       try {
         while (true) {
@@ -2055,7 +2057,7 @@ export const useQueryStore = defineStore("query", () => {
             limit: pageLimit,
             offset,
           });
-          const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, undefined, {
+          const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, exportExecutionId, {
             maxRows: pageLimit,
             fetchSize: pageLimit,
             clientSessionId,
@@ -2106,6 +2108,7 @@ export const useQueryStore = defineStore("query", () => {
     let offset = 0;
     let sessionId: string | undefined;
     const clientSessionId = tabClientSessionId(tab, "export");
+    const exportExecutionId = uuid();
 
     try {
       while (true) {
@@ -2126,7 +2129,7 @@ export const useQueryStore = defineStore("query", () => {
               timeoutSecs: queryTimeoutSecs,
             }
           : { maxRows: plan.pageLimit, fetchSize: plan.pageLimit, clientSessionId, timeoutSecs: queryTimeoutSecs };
-        const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, undefined, executionOptions);
+        const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, exportExecutionId, executionOptions);
         const result = results[0];
         if (!result) break;
         if (columns.length === 0) columns = result.columns;

crates/dbx-core/src/connection.rs

@@ -1,7 +1,7 @@
 use std::collections::{HashMap, HashSet};
 use std::path::PathBuf;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use tokio::sync::RwLock;
 use tokio::task::JoinHandle;
 
@@ -84,6 +84,7 @@ pub enum PoolKind {
 pub struct AppState {
     pub connections: Arc<RwLock<HashMap<String, PoolKind>>>,
     keepalive_tasks: Arc<RwLock<HashMap<String, JoinHandle<()>>>>,
+    pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
     pub configs: RwLock<HashMap<String, ConnectionConfig>>,
     pub running_queries: RunningQueries,
     pub tunnels: TunnelManager,
@@ -95,6 +96,40 @@ pub struct AppState {
     pub mq_registry: crate::mq::MqAdminRegistry,
 }
 
+#[derive(Clone, Copy)]
+struct PoolActivity {
+    last_used_at: Instant,
+}
+
+impl PoolActivity {
+    fn now() -> Self {
+        Self { last_used_at: Instant::now() }
+    }
+}
+
+pub struct PoolActivityTouch {
+    pool_key: String,
+    connections: Arc<RwLock<HashMap<String, PoolKind>>>,
+    pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
+}
+
+impl Drop for PoolActivityTouch {
+    fn drop(&mut self) {
+        let pool_key = self.pool_key.clone();
+        let connections = self.connections.clone();
+        let pool_activity = self.pool_activity.clone();
+        let Ok(handle) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+        handle.spawn(async move {
+            if !connections.read().await.contains_key(&pool_key) {
+                return;
+            }
+            pool_activity.write().await.insert(pool_key, PoolActivity::now());
+        });
+    }
+}
+
 pub fn metadata_connection_config(config: &ConnectionConfig) -> ConnectionConfig {
     let mut db_config = config.canonicalized();
     if database_capabilities::is_metadata_connection_scoped(&db_config.db_type) {
@@ -263,6 +298,7 @@ impl AppState {
         Self {
             connections: Arc::new(RwLock::new(HashMap::new())),
             keepalive_tasks: Arc::new(RwLock::new(HashMap::new())),
+            pool_activity: Arc::new(RwLock::new(HashMap::new())),
             configs: RwLock::new(HashMap::new()),
             running_queries: RunningQueries::default(),
             tunnels: TunnelManager::new(),
@@ -326,6 +362,7 @@ impl AppState {
 
     pub async fn insert_connection_pool(&self, pool_key: String, pool: PoolKind, config: &ConnectionConfig) {
         self.stop_keepalive_task(&pool_key).await;
+        self.pool_activity.write().await.insert(pool_key.clone(), PoolActivity::now());
         self.start_keepalive_task(&pool_key, &pool, config).await;
         let previous = self.connections.write().await.insert(pool_key, pool);
         if let Some(pool) = previous {
@@ -335,37 +372,77 @@ impl AppState {
 
     async fn start_keepalive_task(&self, pool_key: &str, pool: &PoolKind, config: &ConnectionConfig) {
         let interval_secs = config.keepalive_interval_secs;
-        if interval_secs == 0 {
+        let idle_timeout_secs = config.idle_timeout_secs;
+        let idle_cleanup_enabled = is_session_scoped_pool_key(pool_key) && idle_timeout_secs > 0;
+        let mut target = keepalive_target_from_pool(pool, config);
+        if interval_secs == 0 && !idle_cleanup_enabled {
             return;
         }
-        let Some(mut target) = keepalive_target_from_pool(pool, config) else {
+        if interval_secs > 0 && target.is_none() {
             log::debug!(
                 "Connection keepalive requested for '{pool_key}', but this database driver does not keep a pingable client handle."
             );
-            return;
+            if !idle_cleanup_enabled {
+                return;
+            }
         };
 
         let key = pool_key.to_string();
-        let interval = Duration::from_secs(interval_secs.max(1));
+        let interval = if interval_secs > 0 {
+            Duration::from_secs(interval_secs.max(1))
+        } else {
+            Duration::from_secs(idle_timeout_secs.min(60).max(1))
+        };
         let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1));
         let connections = self.connections.clone();
         let keepalive_tasks = self.keepalive_tasks.clone();
+        let pool_activity = self.pool_activity.clone();
+        let running_queries = self.running_queries.clone();
+        let idle_timeout = Duration::from_secs(idle_timeout_secs.max(1));
         let handle = tokio::spawn(async move {
             loop {
                 tokio::time::sleep(interval).await;
-                let result = tokio::time::timeout(timeout, ping_keepalive_target(&mut target, timeout)).await;
+
-                match result {
+                if running_queries.is_pool_active(&key) {
-                    Ok(Ok(())) => {}
+                    continue;
-                    Ok(Err(err)) => {
+                }
-                        log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool");
+
+                if idle_cleanup_enabled {
+                    let idle_for = {
+                        let activity = pool_activity.read().await;
+                        activity.get(&key).map(|activity| activity.last_used_at.elapsed())
+                    };
+                    if idle_for.is_some_and(|elapsed| elapsed >= idle_timeout) {
+                        log::info!(
+                            "Closing idle session-scoped connection pool '{key}' after {}s",
+                            idle_timeout.as_secs()
+                        );
                         keepalive_tasks.write().await.remove(&key);
+                        pool_activity.write().await.remove(&key);
                         let removed = connections.write().await.remove(&key);
                         if let Some(pool) = removed {
                             close_pool_kind_with_timeout(key, pool).await;
                         }
                         break;
                     }
-                    Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()),
+                }
+
+                if let Some(target) = target.as_mut() {
+                    let result = tokio::time::timeout(timeout, ping_keepalive_target(target, timeout)).await;
+                    match result {
+                        Ok(Ok(())) => {}
+                        Ok(Err(err)) => {
+                            log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool");
+                            keepalive_tasks.write().await.remove(&key);
+                            pool_activity.write().await.remove(&key);
+                            let removed = connections.write().await.remove(&key);
+                            if let Some(pool) = removed {
+                                close_pool_kind_with_timeout(key, pool).await;
+                            }
+                            break;
+                        }
+                        Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()),
+                    }
                 }
             }
         });
@@ -391,6 +468,18 @@ impl AppState {
         }
     }
 
+    pub async fn touch_pool_activity(&self, pool_key: &str) {
+        self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now());
+    }
+
+    pub fn pool_activity_touch(&self, pool_key: &str) -> PoolActivityTouch {
+        PoolActivityTouch {
+            pool_key: pool_key.to_string(),
+            connections: self.connections.clone(),
+            pool_activity: self.pool_activity.clone(),
+        }
+    }
+
     pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
         self.get_or_create_pool_for_session(connection_id, database, None).await
     }
@@ -413,6 +502,7 @@ impl AppState {
         if conns.contains_key(&pool_key) {
             drop(conns);
             if !self.remove_stale_connection_pool(&pool_key).await {
+                self.touch_pool_activity(&pool_key).await;
                 return Ok(pool_key);
             }
         } else {
@@ -904,6 +994,10 @@ impl AppState {
     }
 
     async fn remove_stale_connection_pool(&self, pool_key: &str) -> bool {
+        if self.running_queries.is_pool_active(pool_key) {
+            return false;
+        }
+
         let stale = {
             let connections = self.connections.read().await;
             let Some(pool) = connections.get(pool_key) else {
@@ -1085,6 +1179,7 @@ impl AppState {
         }
 
         self.stop_keepalive_task(pool_key).await;
+        self.pool_activity.write().await.remove(pool_key);
         let removed = self.connections.write().await.remove(pool_key);
         if let Some(pool) = removed {
             close_pool_kind(pool).await;
@@ -1115,6 +1210,7 @@ impl AppState {
             self.reset_connection_transport(connection_id).await;
         } else {
             self.stop_keepalive_task(&pool_key).await;
+            self.pool_activity.write().await.remove(&pool_key);
             let removed = self.connections.write().await.remove(&pool_key);
             if let Some(pool) = removed {
                 close_pool_kind(pool).await;
@@ -1143,6 +1239,7 @@ impl AppState {
             return Ok(false);
         }
         self.stop_keepalive_task(&pool_key).await;
+        self.pool_activity.write().await.remove(&pool_key);
         let removed = self.connections.write().await.remove(&pool_key);
         if let Some(pool) = removed {
             close_pool_kind(pool).await;
@@ -1154,6 +1251,7 @@ impl AppState {
 
     pub async fn remove_pool_by_key(&self, pool_key: &str) -> bool {
         self.stop_keepalive_task(pool_key).await;
+        self.pool_activity.write().await.remove(pool_key);
         let removed = self.connections.write().await.remove(pool_key);
         if let Some(pool) = removed {
             close_pool_kind(pool).await;
@@ -1182,6 +1280,12 @@ impl AppState {
             .cloned()
             .collect();
         self.stop_keepalive_tasks(&keys_to_remove).await;
+        {
+            let mut activity = self.pool_activity.write().await;
+            for key in &keys_to_remove {
+                activity.remove(key);
+            }
+        }
         let mut conns = self.connections.write().await;
         let mut removed = Vec::with_capacity(keys_to_remove.len());
         for key in keys_to_remove {
@@ -1472,6 +1576,12 @@ impl AppState {
         // Remove dead pools
         if !dead_keys.is_empty() {
             self.stop_keepalive_tasks(&dead_keys).await;
+            {
+                let mut activity = self.pool_activity.write().await;
+                for key in &dead_keys {
+                    activity.remove(key);
+                }
+            }
             let mut conns = self.connections.write().await;
             for key in &dead_keys {
                 if let Some(pool) = conns.remove(key) {
@@ -1517,6 +1627,12 @@ impl AppState {
             .cloned()
             .collect();
         self.stop_keepalive_tasks(&keys_to_remove).await;
+        {
+            let mut activity = self.pool_activity.write().await;
+            for key in &keys_to_remove {
+                activity.remove(key);
+            }
+        }
         let mut conns = self.connections.write().await;
         let mut removed = Vec::with_capacity(keys_to_remove.len());
         for key in keys_to_remove {
@@ -1545,6 +1661,12 @@ impl AppState {
             })
             .collect();
         self.stop_keepalive_tasks(&keys_to_remove).await;
+        {
+            let mut activity = self.pool_activity.write().await;
+            for key in &keys_to_remove {
+                activity.remove(key);
+            }
+        }
         let mut conns = self.connections.write().await;
         let mut removed = Vec::with_capacity(keys_to_remove.len());
         for key in keys_to_remove {
@@ -1697,6 +1819,10 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str
         .unwrap_or(base_pool_key)
 }
 
+fn is_session_scoped_pool_key(pool_key: &str) -> bool {
+    pool_key.contains(":session:")
+}
+
 pub(crate) fn config_for_pool_key<'a>(
     pool_key: &str,
     configs: &'a HashMap<String, ConnectionConfig>,
@@ -2763,6 +2889,23 @@ mod tests {
         );
     }
 
+    #[test]
+    fn session_scoped_pool_keys_are_sanitized_and_detected() {
+        let key = super::session_scoped_pool_key_for(
+            Some(DatabaseType::Mysql),
+            "mysql-conn:analytics".to_string(),
+            Some("tab-1:count"),
+        );
+
+        assert_eq!(key, "mysql-conn:analytics:session:tab-1_count");
+        assert!(super::is_session_scoped_pool_key(&key));
+        assert!(!super::is_session_scoped_pool_key("mysql-conn:analytics"));
+        assert_eq!(
+            super::session_scoped_pool_key_for(Some(DatabaseType::DuckDb), "duckdb-conn".to_string(), Some("tab-1")),
+            "duckdb-conn"
+        );
+    }
+
     #[test]
     fn mysql_direct_connections_skip_tcp_probe() {
         let mut config = mysql_config(Some("app"));
@@ -2884,6 +3027,36 @@ mod tests {
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    #[tokio::test]
+    async fn pool_activity_touch_updates_existing_pool_only() {
+        let (state, dir) = test_app_state().await;
+        let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
+        let pool_key = "conn:session:tab-1";
+
+        state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
+        state.pool_activity.write().await.insert(
+            pool_key.to_string(),
+            super::PoolActivity { last_used_at: std::time::Instant::now() - std::time::Duration::from_secs(10) },
+        );
+
+        {
+            let _touch = state.pool_activity_touch(pool_key);
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        let elapsed = state.pool_activity.read().await.get(pool_key).unwrap().last_used_at.elapsed();
+        assert!(elapsed < std::time::Duration::from_secs(10));
+
+        {
+            let _touch = state.pool_activity_touch(pool_key);
+            state.connections.write().await.remove(pool_key);
+            state.pool_activity.write().await.remove(pool_key);
+        }
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        assert!(!state.pool_activity.read().await.contains_key(pool_key));
+
+        let _ = std::fs::remove_dir_all(dir);
+    }
+
     #[cfg(feature = "duckdb-bundled")]
     #[tokio::test]
     async fn duckdb_client_session_reuses_base_pool_to_avoid_file_locks() {

crates/dbx-core/src/query.rs

@@ -805,6 +805,12 @@ pub async fn do_execute(
     cancel_token: Option<CancellationToken>,
     options: QueryExecutionOptions,
 ) -> Result<db::QueryResult, String> {
+    if let Some(execution_id) = options.execution_id.as_deref() {
+        state.running_queries.set_pool_key(execution_id, pool_key.to_string());
+    }
+    state.touch_pool_activity(pool_key).await;
+    let _activity_touch = state.pool_activity_touch(pool_key);
+
     let query_timeout = resolve_query_timeout(options.timeout_secs);
     let (_duckdb_attached_names, conn_name_if_readonly) = {
         let configs = state.configs.read().await;
@@ -1257,6 +1263,11 @@ async fn execute_postgres_drop_database(
     let pool_key = state
         .get_or_create_pool_for_session(connection_id, Some(admin_database), options.client_session_id.as_deref())
         .await?;
+    if let Some(execution_id) = options.execution_id.as_deref() {
+        state.running_queries.set_pool_key(execution_id, pool_key.clone());
+    }
+    state.touch_pool_activity(&pool_key).await;
+    let _activity_touch = state.pool_activity_touch(pool_key.as_str());
 
     if is_canceled(&cancel_token) {
         return Err(canceled_error());
@@ -1390,6 +1401,11 @@ pub async fn execute_multi_core_with_options(
             .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref())
             .await?
     };
+    if let Some(execution_id) = options.execution_id.as_deref() {
+        state.running_queries.set_pool_key(execution_id, pool_key.clone());
+    }
+    state.touch_pool_activity(&pool_key).await;
+    let _activity_touch = state.pool_activity_touch(pool_key.as_str());
 
     let is_sqlserver = {
         let connections = state.connections.read().await;

crates/dbx-core/src/query_cancel.rs

@@ -4,16 +4,64 @@ use tokio_util::sync::CancellationToken;
 
 type InterruptFn = Box<dyn Fn() + Send + 'static>;
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum RunningTaskKind {
+    Query,
+    Count,
+    Explain,
+    Export,
+    Unknown,
+}
+
+impl Default for RunningTaskKind {
+    fn default() -> Self {
+        Self::Unknown
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct RunningTaskMetadata {
+    pub kind: RunningTaskKind,
+    pub connection_id: Option<String>,
+    pub database: Option<String>,
+    pub client_session_id: Option<String>,
+}
+
+impl RunningTaskMetadata {
+    pub fn query(
+        connection_id: impl Into<String>,
+        database: impl Into<String>,
+        client_session_id: Option<String>,
+    ) -> Self {
+        let kind = task_kind_from_client_session_id(client_session_id.as_deref());
+        Self { kind, connection_id: Some(connection_id.into()), database: Some(database.into()), client_session_id }
+    }
+}
+
+#[derive(Clone)]
+struct RunningTask {
+    token: CancellationToken,
+    metadata: RunningTaskMetadata,
+    pool_key: Option<String>,
+}
+
 #[derive(Clone, Default)]
 pub struct RunningQueries {
-    inner: Arc<Mutex<HashMap<String, CancellationToken>>>,
+    inner: Arc<Mutex<HashMap<String, RunningTask>>>,
     interrupts: Arc<Mutex<HashMap<String, InterruptFn>>>,
 }
 
 impl RunningQueries {
     pub fn register(&self, execution_id: String) -> RegisteredQuery {
+        self.register_task(execution_id, RunningTaskMetadata::default())
+    }
+
+    pub fn register_task(&self, execution_id: String, metadata: RunningTaskMetadata) -> RegisteredQuery {
         let token = CancellationToken::new();
-        self.inner.lock().unwrap_or_else(|e| e.into_inner()).insert(execution_id.clone(), token.clone());
+        self.inner
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .insert(execution_id.clone(), RunningTask { token: token.clone(), metadata, pool_key: None });
 
         RegisteredQuery { execution_id, token, running_queries: self.clone() }
     }
@@ -23,7 +71,8 @@ impl RunningQueries {
     }
 
     pub fn cancel(&self, execution_id: &str) -> bool {
-        let token = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).cloned();
+        let token =
+            self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.token.clone());
         let interrupt = self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
 
         if let Some(interrupt) = interrupt {
@@ -37,17 +86,51 @@ impl RunningQueries {
         }
     }
 
+    pub fn set_pool_key(&self, execution_id: &str, pool_key: impl Into<String>) {
+        if let Some(task) = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get_mut(execution_id) {
+            task.pool_key = Some(pool_key.into());
+        }
+    }
+
+    pub fn is_pool_active(&self, pool_key: &str) -> bool {
+        self.inner.lock().unwrap_or_else(|e| e.into_inner()).values().any(|task| {
+            let _kind = task.metadata.kind;
+            task.pool_key.as_deref() == Some(pool_key)
+        })
+    }
+
     #[cfg(test)]
     pub fn has(&self, execution_id: &str) -> bool {
         self.inner.lock().unwrap_or_else(|e| e.into_inner()).contains_key(execution_id)
     }
 
+    #[cfg(test)]
+    pub fn task_kind(&self, execution_id: &str) -> Option<RunningTaskKind> {
+        self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.metadata.kind)
+    }
+
     fn remove(&self, execution_id: &str) {
         self.inner.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
         self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
     }
 }
 
+fn task_kind_from_client_session_id(client_session_id: Option<&str>) -> RunningTaskKind {
+    let Some(session_id) = client_session_id else {
+        return RunningTaskKind::Query;
+    };
+    let session_id = session_id.trim().to_ascii_lowercase();
+    if session_id.ends_with(":count") {
+        RunningTaskKind::Count
+    } else if session_id.ends_with(":explain") {
+        RunningTaskKind::Explain
+    } else if session_id.ends_with(":export") {
+        RunningTaskKind::Export
+    } else {
+        RunningTaskKind::Query
+    }
+}
+
 pub struct RegisteredQuery {
     execution_id: String,
     token: CancellationToken,
@@ -68,7 +151,7 @@ impl Drop for RegisteredQuery {
 
 #[cfg(test)]
 mod tests {
-    use super::RunningQueries;
+    use super::{RunningQueries, RunningTaskKind, RunningTaskMetadata};
 
     #[test]
     fn cancel_marks_registered_query_as_cancelled() {
@@ -103,4 +186,20 @@ mod tests {
 
         assert!(!running.has("exec-1"));
     }
+
+    #[test]
+    fn register_task_tracks_kind_and_pool_activity() {
+        let running = RunningQueries::default();
+        let registered = running.register_task(
+            "exec-1".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())),
+        );
+
+        running.set_pool_key("exec-1", "conn-1:main:session:tab-1_export");
+
+        assert_eq!(running.task_kind("exec-1"), Some(RunningTaskKind::Export));
+        assert!(running.is_pool_active("conn-1:main:session:tab-1_export"));
+        drop(registered);
+        assert!(!running.is_pool_active("conn-1:main:session:tab-1_export"));
+    }
 }

crates/dbx-web/src/routes/query.rs

@@ -6,6 +6,7 @@ use serde::Deserialize;
 
 use crate::error::AppError;
 use crate::state::WebState;
+use dbx_core::query_cancel::RunningTaskMetadata;
 
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -269,7 +270,10 @@ pub async fn execute_query(
 ) -> Result<Json<serde_json::Value>, AppError> {
     let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
 
-    let registered = state.app.running_queries.register(execution_id.clone());
+    let registered = state.app.running_queries.register_task(
+        execution_id.clone(),
+        RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()),
+    );
     let cancel_token = registered.token();
 
     let result = dbx_core::query::execute_sql_statement_with_options(
@@ -302,7 +306,10 @@ pub async fn execute_multi(
 ) -> Result<Json<serde_json::Value>, AppError> {
     let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
 
-    let registered = state.app.running_queries.register(execution_id.clone());
+    let registered = state.app.running_queries.register_task(
+        execution_id.clone(),
+        RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()),
+    );
     let cancel_token = registered.token();
 
     let result = dbx_core::query::execute_multi_core_with_options(

src-tauri/src/commands/query.rs

@@ -5,6 +5,7 @@ use tauri::State;
 use crate::commands::connection::AppState;
 use dbx_core::db;
 use dbx_core::models::connection::DatabaseType;
+use dbx_core::query_cancel::RunningTaskMetadata;
 use dbx_core::sql::split_sql_statements;
 
 #[tauri::command]
@@ -23,8 +24,13 @@ pub async fn execute_query(
     client_session_id: Option<String>,
     timeout_secs: Option<u64>,
 ) -> Result<db::QueryResult, String> {
-    let registered_query =
+    let execution_id = execution_id.filter(|id| !id.trim().is_empty());
-        execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone()));
+    let registered_query = execution_id.as_ref().map(|id| {
+        state.running_queries.register_task(
+            id.clone(),
+            RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()),
+        )
+    });
     let cancel_token = registered_query.as_ref().map(|query| query.token());
 
     dbx_core::query::execute_sql_statement_with_options(
@@ -41,7 +47,7 @@ pub async fn execute_query(
             result_session_id,
             client_session_id,
             timeout_secs,
-            execution_id: execution_id.filter(|id| !id.trim().is_empty()),
+            execution_id,
         },
     )
     .await
@@ -63,10 +69,15 @@ pub async fn execute_multi(
     client_session_id: Option<String>,
     timeout_secs: Option<u64>,
 ) -> Result<Vec<db::QueryResult>, String> {
-    let registered_query =
+    let execution_id = execution_id.filter(|id| !id.trim().is_empty());
-        execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone()));
+    let registered_query = execution_id.as_ref().map(|id| {
+        state.running_queries.register_task(
+            id.clone(),
+            RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()),
+        )
+    });
     let cancel_token = registered_query.as_ref().map(|query| query.token());
-    let trace_id = execution_id.clone().as_deref().unwrap_or("no-execution-id").to_string();
+    let trace_id = execution_id.as_deref().unwrap_or("no-execution-id").to_string();
     let started_at = Instant::now();
     log::info!(
         "[query][execute_multi:start] trace_id={} connection_id={} database={} schema={:?} sql={}",
@@ -91,7 +102,7 @@ pub async fn execute_multi(
             result_session_id,
             client_session_id,
             timeout_secs,
-            execution_id: execution_id.filter(|id| !id.trim().is_empty()),
+            execution_id,
         },
     )
     .await;