feat(core): optimize query resource lifecycle

crates/dbx-core/src/connection.rs

@@ -4,6 +4,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::sync::{Mutex, RwLock};
+#[cfg(feature = "duckdb-bundled")]
 use tokio::task::JoinHandle;
 
 use mysql_async::prelude::Queryable;
@@ -29,6 +30,7 @@ use crate::path_utils::expand_tilde;
 use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv};
 use crate::query_cancel::RunningQueries;
 use crate::storage::{normalize_duckdb_worker_max_processes, Storage, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT};
+use crate::task_supervisor::TaskSupervisor;
 
 pub const JDBC_PLUGIN_NOT_INSTALLED: &str =
     "JDBC plugin is not installed. Install the optional JDBC plugin to use this connection.";
@@ -175,7 +177,7 @@ macro_rules! agent_connection_pool_database_type {
 
 pub struct AppState {
     pub connections: Arc<RwLock<HashMap<String, PoolKind>>>,
-    keepalive_tasks: Arc<RwLock<HashMap<String, JoinHandle<()>>>>,
+    task_supervisor: TaskSupervisor,
     pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
     connection_attempts: RwLock<HashMap<String, ConnectionAttemptState>>,
     pub configs: RwLock<HashMap<String, ConnectionConfig>>,
@@ -219,6 +221,7 @@ pub struct PoolActivityTouch {
     pool_key: String,
     connections: Arc<RwLock<HashMap<String, PoolKind>>>,
     pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
+    task_supervisor: TaskSupervisor,
 }
 
 impl Drop for PoolActivityTouch {
@@ -226,10 +229,7 @@ impl Drop for PoolActivityTouch {
         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 {
+        self.task_supervisor.spawn_replace(format!("pool-activity:{pool_key}"), move |_| async move {
             if !connections.read().await.contains_key(&pool_key) {
                 return;
             }
@@ -569,7 +569,7 @@ impl AppState {
     ) -> Self {
         Self {
             connections: Arc::new(RwLock::new(HashMap::new())),
-            keepalive_tasks: Arc::new(RwLock::new(HashMap::new())),
+            task_supervisor: TaskSupervisor::new(),
             pool_activity: Arc::new(RwLock::new(HashMap::new())),
             connection_attempts: RwLock::new(HashMap::new()),
             configs: RwLock::new(HashMap::new()),
@@ -926,13 +926,15 @@ impl AppState {
         let interval = Duration::from_secs(interval_secs.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 cancel_contexts = self.postgres_cancel_contexts.clone();
         let running_queries = self.running_queries.clone();
-        let handle = tokio::spawn(async move {
+        self.task_supervisor.spawn_replace(format!("keepalive:{pool_key}"), move |shutdown| async move {
             loop {
-                tokio::time::sleep(interval).await;
+                tokio::select! {
+                    _ = shutdown.cancelled() => break,
+                    _ = tokio::time::sleep(interval) => {}
+                }
 
                 if running_queries.is_pool_active(&key) {
                     continue;
@@ -944,7 +946,6 @@ impl AppState {
                         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);
                             cancel_contexts.write().await.remove(&key);
                             let removed = connections.write().await.remove(&key);
@@ -958,7 +959,6 @@ impl AppState {
                                 "Connection keepalive timed out for '{key}' after {}s; invalidating pool",
                                 timeout.as_secs()
                             );
-                            keepalive_tasks.write().await.remove(&key);
                             pool_activity.write().await.remove(&key);
                             cancel_contexts.write().await.remove(&key);
                             let removed = connections.write().await.remove(&key);
@@ -971,26 +971,15 @@ impl AppState {
                 }
             }
         });
-        let previous = self.keepalive_tasks.write().await.insert(pool_key.to_string(), handle);
-        if let Some(previous) = previous {
-            previous.abort();
-        }
     }
 
     async fn stop_keepalive_task(&self, pool_key: &str) {
-        let task = self.keepalive_tasks.write().await.remove(pool_key);
-        if let Some(task) = task {
-            task.abort();
-        }
+        self.task_supervisor.stop(&format!("keepalive:{pool_key}"));
     }
 
     async fn stop_keepalive_tasks(&self, pool_keys: &[String]) {
-        let mut tasks = self.keepalive_tasks.write().await;
-        for pool_key in pool_keys {
-            if let Some(task) = tasks.remove(pool_key) {
-                task.abort();
-            }
-        }
+        let keys: Vec<String> = pool_keys.iter().map(|pool_key| format!("keepalive:{pool_key}")).collect();
+        self.task_supervisor.stop_many(keys.iter().map(String::as_str));
     }
 
     pub async fn touch_pool_activity(&self, pool_key: &str) {
@@ -1007,9 +996,20 @@ impl AppState {
             pool_key: pool_key.to_string(),
             connections: self.connections.clone(),
             pool_activity: self.pool_activity.clone(),
+            task_supervisor: self.task_supervisor.clone(),
         }
     }
 
+    pub async fn shutdown_background_tasks(&self, deadline: Duration) {
+        self.running_queries.cancel_all();
+        self.task_supervisor.shutdown(deadline).await;
+    }
+
+    #[cfg(test)]
+    pub fn supervised_task_count(&self) -> usize {
+        self.task_supervisor.active_count()
+    }
+
     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
     }
@@ -2207,16 +2207,18 @@ impl AppState {
     #[cfg(feature = "duckdb-bundled")]
     pub fn spawn_duckdb_pool_cleanup(&self, pool_key: String, con: DuckDbHandle) {
         let connections = self.connections.clone();
-        let keepalive_tasks = self.keepalive_tasks.clone();
+        let supervisor = self.task_supervisor.clone();
         let pool_activity = self.pool_activity.clone();
         let postgres_cancel_contexts = self.postgres_cancel_contexts.clone();
-        tokio::spawn(async move {
+        let task_key = format!("duckdb-cleanup:{pool_key}");
+        supervisor.clone().spawn_replace(task_key, move |shutdown| async move {
             while Arc::strong_count(&con) > 2 {
-                tokio::time::sleep(Duration::from_millis(50)).await;
-            }
-            if let Some(handle) = keepalive_tasks.write().await.remove(&pool_key) {
-                handle.abort();
+                tokio::select! {
+                    _ = shutdown.cancelled() => return,
+                    _ = tokio::time::sleep(Duration::from_millis(50)) => {}
+                }
             }
+            supervisor.stop(&format!("keepalive:{pool_key}"));
             pool_activity.write().await.remove(&pool_key);
             postgres_cancel_contexts.write().await.remove(&pool_key);
             let removed = {
@@ -2241,20 +2243,28 @@ impl AppState {
         &self,
         pool_key: String,
         con: DuckDbHandle,
-        task: JoinHandle<Result<db::QueryResult, String>>,
+        mut task: JoinHandle<Result<db::QueryResult, String>>,
     ) {
         let connections = self.connections.clone();
-        let keepalive_tasks = self.keepalive_tasks.clone();
+        let supervisor = self.task_supervisor.clone();
         let pool_activity = self.pool_activity.clone();
         let postgres_cancel_contexts = self.postgres_cancel_contexts.clone();
-        tokio::spawn(async move {
-            let _ = task.await;
+        let task_key = format!("duckdb-draining:{pool_key}");
+        supervisor.clone().spawn_replace(task_key, move |shutdown| async move {
+            tokio::select! {
+                _ = shutdown.cancelled() => {
+                    task.abort();
+                    return;
+                }
+                _ = &mut task => {}
+            }
             while Arc::strong_count(&con) > 2 {
-                tokio::time::sleep(Duration::from_millis(50)).await;
-            }
-            if let Some(handle) = keepalive_tasks.write().await.remove(&pool_key) {
-                handle.abort();
+                tokio::select! {
+                    _ = shutdown.cancelled() => return,
+                    _ = tokio::time::sleep(Duration::from_millis(50)) => {}
+                }
             }
+            supervisor.stop(&format!("keepalive:{pool_key}"));
             pool_activity.write().await.remove(&pool_key);
             postgres_cancel_contexts.write().await.remove(&pool_key);
             let removed = {
@@ -3288,26 +3298,6 @@ fn native_postgres_url_config(config: &ConnectionConfig) -> Option<ConnectionCon
     }
 }
 
-#[cfg(feature = "duckdb-bundled")]
-fn duckdb_paths_match(left: &str, right: &str) -> bool {
-    let left = expand_tilde(left);
-    let right = expand_tilde(right);
-
-    if db::duckdb_driver::is_memory_database_path(&left) || db::duckdb_driver::is_memory_database_path(&right) {
-        return left.trim().eq_ignore_ascii_case(right.trim());
-    }
-
-    if let (Ok(left_path), Ok(right_path)) = (std::fs::canonicalize(&left), std::fs::canonicalize(&right)) {
-        return left_path == right_path;
-    }
-
-    if cfg!(windows) {
-        left.eq_ignore_ascii_case(&right)
-    } else {
-        left == right
-    }
-}
-
 pub async fn probe_connection_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> Result<(), String> {
     if !uses_tcp_probe(config, host, port) {
         return Ok(());
@@ -4671,7 +4661,7 @@ mod tests {
         tokio::time::sleep(std::time::Duration::from_millis(20)).await;
 
         assert!(state.connections.read().await.contains_key(pool_key));
-        assert!(!state.keepalive_tasks.read().await.contains_key(pool_key));
+        assert_eq!(state.supervised_task_count(), 0);
 
         let _ = std::fs::remove_dir_all(dir);
     }

crates/dbx-core/src/lib.rs

@@ -60,6 +60,7 @@ pub mod storage;
 pub mod table_export;
 pub mod table_import;
 pub mod table_structure_sql;
+pub mod task_supervisor;
 pub mod text_export;
 pub mod token_usage;
 pub mod transfer;

crates/dbx-core/src/query_cancel.rs

@@ -1,7 +1,15 @@
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use tokio_util::sync::CancellationToken;
 
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct RunningQueryDiagnostics {
+    pub active_execution_ids: Vec<String>,
+    pub active_by_connection: HashMap<String, usize>,
+    pub interrupt_registrations: usize,
+}
+
 type InterruptFn = Box<dyn Fn() + Send + 'static>;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -25,6 +33,7 @@ pub struct RunningTaskMetadata {
     pub connection_id: Option<String>,
     pub database: Option<String>,
     pub client_session_id: Option<String>,
+    pub owner_scope: Option<String>,
 }
 
 impl RunningTaskMetadata {
@@ -34,12 +43,25 @@ impl RunningTaskMetadata {
         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 }
+        let owner_scope = client_session_id.as_deref().and_then(owner_scope_from_client_session_id);
+        Self {
+            kind,
+            connection_id: Some(connection_id.into()),
+            database: Some(database.into()),
+            client_session_id,
+            owner_scope,
+        }
+    }
+
+    pub fn with_owner_scope(mut self, owner_scope: impl Into<String>) -> Self {
+        self.owner_scope = Some(owner_scope.into());
+        self
     }
 }
 
 #[derive(Clone)]
 struct RunningTask {
+    registration_id: u64,
     token: CancellationToken,
     metadata: RunningTaskMetadata,
     pool_key: Option<String>,
@@ -49,6 +71,7 @@ struct RunningTask {
 pub struct RunningQueries {
     inner: Arc<Mutex<HashMap<String, RunningTask>>>,
     interrupts: Arc<Mutex<HashMap<String, InterruptFn>>>,
+    next_registration_id: Arc<AtomicU64>,
 }
 
 impl RunningQueries {
@@ -58,12 +81,19 @@ impl RunningQueries {
 
     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(), RunningTask { token: token.clone(), metadata, pool_key: None });
+        let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed) + 1;
+        let previous = self.inner.lock().unwrap_or_else(|e| e.into_inner()).insert(
+            execution_id.clone(),
+            RunningTask { registration_id, token: token.clone(), metadata, pool_key: None },
+        );
+        if let Some(previous) = previous {
+            previous.token.cancel();
+            if let Some(interrupt) = self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(&execution_id) {
+                interrupt();
+            }
+        }
 
-        RegisteredQuery { execution_id, token, running_queries: self.clone() }
+        RegisteredQuery { execution_id, registration_id, token, running_queries: self.clone() }
     }
 
     pub fn register_interrupt(&self, execution_id: &str, interrupt: impl Fn() + Send + 'static) {
@@ -86,6 +116,52 @@ impl RunningQueries {
         }
     }
 
+    pub fn cancel_connection(&self, connection_id: &str) -> usize {
+        self.cancel_matching(|task| task.metadata.connection_id.as_deref() == Some(connection_id))
+    }
+
+    pub fn cancel_client_session(&self, client_session_id: &str) -> usize {
+        self.cancel_matching(|task| task.metadata.client_session_id.as_deref() == Some(client_session_id))
+    }
+
+    pub fn cancel_owner_scope(&self, owner_scope: &str) -> usize {
+        self.cancel_matching(|task| task.metadata.owner_scope.as_deref() == Some(owner_scope))
+    }
+
+    pub fn cancel_all(&self) -> usize {
+        self.cancel_matching(|_| true)
+    }
+
+    pub fn diagnostics(&self) -> RunningQueryDiagnostics {
+        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
+        let mut active_execution_ids = inner.keys().cloned().collect::<Vec<_>>();
+        active_execution_ids.sort();
+        let mut active_by_connection = HashMap::new();
+        for task in inner.values() {
+            if let Some(connection_id) = &task.metadata.connection_id {
+                *active_by_connection.entry(connection_id.clone()).or_insert(0) += 1;
+            }
+        }
+        drop(inner);
+        RunningQueryDiagnostics {
+            active_execution_ids,
+            active_by_connection,
+            interrupt_registrations: self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).len(),
+        }
+    }
+
+    fn cancel_matching(&self, predicate: impl Fn(&RunningTask) -> bool) -> usize {
+        let execution_ids: Vec<String> = self
+            .inner
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .iter()
+            .filter(|(_, task)| predicate(task))
+            .map(|(execution_id, _)| execution_id.clone())
+            .collect();
+        execution_ids.iter().filter(|execution_id| self.cancel(execution_id)).count()
+    }
+
     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());
@@ -109,9 +185,27 @@ impl RunningQueries {
         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);
+    #[cfg(test)]
+    pub fn registration_counts(&self) -> (usize, usize) {
+        (
+            self.inner.lock().unwrap_or_else(|e| e.into_inner()).len(),
+            self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).len(),
+        )
+    }
+
+    fn remove(&self, execution_id: &str, registration_id: u64) {
+        let removed = {
+            let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
+            if inner.get(execution_id).is_some_and(|task| task.registration_id == registration_id) {
+                inner.remove(execution_id);
+                true
+            } else {
+                false
+            }
+        };
+        if removed {
+            self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
+        }
     }
 }
 
@@ -131,8 +225,14 @@ fn task_kind_from_client_session_id(client_session_id: Option<&str>) -> RunningT
     }
 }
 
+fn owner_scope_from_client_session_id(client_session_id: &str) -> Option<String> {
+    let owner = client_session_id.trim().split(':').next()?.trim();
+    (!owner.is_empty()).then(|| owner.to_string())
+}
+
 pub struct RegisteredQuery {
     execution_id: String,
+    registration_id: u64,
     token: CancellationToken,
     running_queries: RunningQueries,
 }
@@ -145,7 +245,7 @@ impl RegisteredQuery {
 
 impl Drop for RegisteredQuery {
     fn drop(&mut self) {
-        self.running_queries.remove(&self.execution_id);
+        self.running_queries.remove(&self.execution_id, self.registration_id);
     }
 }
 
@@ -187,6 +287,61 @@ mod tests {
         assert!(!running.has("exec-1"));
     }
 
+    #[test]
+    fn stale_registration_drop_does_not_remove_replacement() {
+        let running = RunningQueries::default();
+        let first = running.register("exec-1".to_string());
+        let second = running.register("exec-1".to_string());
+
+        assert!(first.token().is_cancelled());
+        drop(first);
+        assert!(running.has("exec-1"));
+        drop(second);
+        assert!(!running.has("exec-1"));
+    }
+
+    #[test]
+    fn scoped_cancellation_only_cancels_matching_tasks() {
+        let running = RunningQueries::default();
+        let first = running.register_task(
+            "exec-1".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())),
+        );
+        let second = running.register_task(
+            "exec-2".to_string(),
+            RunningTaskMetadata::query("conn-2", "main", Some("tab-2".to_string())),
+        );
+
+        assert_eq!(running.cancel_connection("conn-1"), 1);
+        assert!(first.token().is_cancelled());
+        assert!(!second.token().is_cancelled());
+        assert_eq!(running.cancel_client_session("tab-2"), 1);
+        assert!(second.token().is_cancelled());
+    }
+
+    #[test]
+    fn query_metadata_derives_stable_owner_scope() {
+        let metadata = RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string()));
+        assert_eq!(metadata.owner_scope.as_deref(), Some("tab-1"));
+    }
+
+    #[test]
+    fn owner_scope_cancellation_matches_related_task_kinds() {
+        let running = RunningQueries::default();
+        let query = running.register_task(
+            "exec-query".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())),
+        );
+        let export = running.register_task(
+            "exec-export".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())),
+        );
+
+        assert_eq!(running.cancel_owner_scope("tab-1"), 2);
+        assert!(query.token().is_cancelled());
+        assert!(export.token().is_cancelled());
+    }
+
     #[test]
     fn register_task_tracks_kind_and_pool_activity() {
         let running = RunningQueries::default();
@@ -202,4 +357,48 @@ mod tests {
         drop(registered);
         assert!(!running.is_pool_active("conn-1:main:session:tab-1_export"));
     }
+
+    #[test]
+    fn unwind_removes_task_and_interrupt_registrations() {
+        let running = RunningQueries::default();
+        let unwind_running = running.clone();
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
+            let _registered = unwind_running.register("exec-unwind".to_string());
+            unwind_running.register_interrupt("exec-unwind", || {});
+            panic!("simulate driver unwind");
+        }));
+
+        assert!(result.is_err());
+        assert_eq!(running.registration_counts(), (0, 0));
+    }
+
+    #[test]
+    fn completion_after_cancel_removes_all_registrations() {
+        let running = RunningQueries::default();
+        let registered = running.register("exec-cancel".to_string());
+        running.register_interrupt("exec-cancel", || {});
+
+        assert!(running.cancel("exec-cancel"));
+        drop(registered);
+
+        assert_eq!(running.registration_counts(), (0, 0));
+    }
+
+    #[test]
+    fn diagnostics_expose_ids_and_scoped_counts_without_query_text() {
+        let running = RunningQueries::default();
+        let _first = running.register_task(
+            "exec-2".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())),
+        );
+        let _second = running.register_task(
+            "exec-1".to_string(),
+            RunningTaskMetadata::query("conn-1", "main", Some("tab-2".to_string())),
+        );
+
+        let diagnostics = running.diagnostics();
+
+        assert_eq!(diagnostics.active_execution_ids, vec!["exec-1", "exec-2"]);
+        assert_eq!(diagnostics.active_by_connection.get("conn-1"), Some(&2));
+    }
 }

crates/dbx-core/src/storage.rs

@@ -1,5 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use log::warn;
 use rusqlite::{params, params_from_iter, Connection, DatabaseName, OpenFlags, OptionalExtension, ToSql};
@@ -102,6 +103,7 @@ pub struct Storage {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct TabRuntimeCacheEntry {
     pub key: String,
     pub payload: Vec<u8>,
@@ -109,6 +111,32 @@ pub struct TabRuntimeCacheEntry {
     pub column_count: i64,
     pub byte_size: i64,
     pub updated_at: String,
+    pub created_at: i64,
+    pub last_accessed_at: i64,
+    pub owner_id: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TabRuntimeCacheMetadata {
+    pub key: String,
+    pub row_count: i64,
+    pub column_count: i64,
+    pub byte_size: i64,
+    pub updated_at: String,
+    pub created_at: i64,
+    pub last_accessed_at: i64,
+    pub owner_id: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TabRuntimeCachePruneResult {
+    pub deleted_entries: usize,
+    pub deleted_bytes: i64,
+    pub orphan_deletions: usize,
+    pub remaining_entries: usize,
+    pub remaining_bytes: i64,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -260,7 +288,10 @@ const SCHEMA_STATEMENTS: &[&str] = &[
         row_count INTEGER NOT NULL DEFAULT 0,
         column_count INTEGER NOT NULL DEFAULT 0,
         byte_size INTEGER NOT NULL DEFAULT 0,
-        updated_at TEXT NOT NULL
+        updated_at TEXT NOT NULL,
+        created_at INTEGER NOT NULL DEFAULT 0,
+        last_accessed_at INTEGER NOT NULL DEFAULT 0,
+        owner_id TEXT
     )",
     "CREATE TABLE IF NOT EXISTS mq_token_records (
         id TEXT PRIMARY KEY,
@@ -319,6 +350,7 @@ impl Storage {
             }
             ensure_history_columns_sync(conn)?;
             ensure_saved_sql_columns_sync(conn)?;
+            ensure_tab_runtime_cache_columns_sync(conn)?;
             Ok(())
         })
     }
@@ -432,6 +464,26 @@ fn ensure_saved_sql_columns_sync(conn: &Connection) -> Result<(), String> {
     Ok(())
 }
 
+fn ensure_tab_runtime_cache_columns_sync(conn: &Connection) -> Result<(), String> {
+    const COLUMNS: &[(&str, &str)] = &[
+        ("created_at", "INTEGER NOT NULL DEFAULT 0"),
+        ("last_accessed_at", "INTEGER NOT NULL DEFAULT 0"),
+        ("owner_id", "TEXT"),
+    ];
+    ensure_table_columns(conn, "tab_runtime_cache", COLUMNS)?;
+    let now = unix_timestamp_millis();
+    // Legacy rows must receive a grace period instead of being treated as ancient crash leftovers.
+    conn.execute("UPDATE tab_runtime_cache SET created_at = ?1 WHERE created_at = 0", [now])
+        .map_err(|e| e.to_string())?;
+    conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = created_at WHERE last_accessed_at = 0", [])
+        .map_err(|e| e.to_string())?;
+    Ok(())
+}
+
+fn unix_timestamp_millis() -> i64 {
+    SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis().min(i64::MAX as u128) as i64
+}
+
 fn ensure_table_columns(conn: &Connection, table_name: &str, columns: &[(&str, &str)]) -> Result<(), String> {
     let mut stmt =
         conn.prepare(&format!("SELECT name FROM pragma_table_info('{table_name}')")).map_err(|e| e.to_string())?;
@@ -2082,18 +2134,21 @@ impl Storage {
         payload: Vec<u8>,
         row_count: i64,
         column_count: i64,
+        owner_id: Option<String>,
     ) -> Result<(), String> {
         let key = key.to_string();
         let byte_size = payload.len() as i64;
+        let now = unix_timestamp_millis();
         self.with_conn(move |conn| {
             conn.execute(
                 "INSERT INTO tab_runtime_cache \
-                 (cache_key, payload, row_count, column_count, byte_size, updated_at) \
-                 VALUES (?1, ?2, ?3, ?4, ?5, datetime('now')) \
+                 (cache_key, payload, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id) \
+                 VALUES (?1, ?2, ?3, ?4, ?5, datetime('now'), ?6, ?6, ?7) \
                  ON CONFLICT(cache_key) DO UPDATE SET \
                  payload = excluded.payload, row_count = excluded.row_count, column_count = excluded.column_count, \
-                 byte_size = excluded.byte_size, updated_at = excluded.updated_at",
-                params![key, payload, row_count, column_count, byte_size],
+                 byte_size = excluded.byte_size, updated_at = excluded.updated_at, \
+                 last_accessed_at = excluded.last_accessed_at, owner_id = excluded.owner_id",
+                params![key, payload, row_count, column_count, byte_size, now, owner_id],
             )
             .map(|_| ())
             .map_err(|e| e.to_string())
@@ -2103,11 +2158,13 @@ impl Storage {
 
     pub async fn load_tab_runtime_cache(&self, key: &str) -> Result<Option<TabRuntimeCacheEntry>, String> {
         let key = key.to_string();
+        let now = unix_timestamp_millis();
         self.with_conn(move |conn| {
-            conn.query_row(
-                "SELECT cache_key, payload, row_count, column_count, byte_size, updated_at \
+            let entry = conn
+                .query_row(
+                "SELECT cache_key, payload, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id \
                  FROM tab_runtime_cache WHERE cache_key = ?1",
-                [key],
+                [&key],
                 |row| {
                     Ok(TabRuntimeCacheEntry {
                         key: row.get(0)?,
@@ -2116,11 +2173,134 @@ impl Storage {
                         column_count: row.get(3)?,
                         byte_size: row.get(4)?,
                         updated_at: row.get(5)?,
+                        created_at: row.get(6)?,
+                        last_accessed_at: now,
+                        owner_id: row.get(8)?,
                     })
                 },
             )
             .optional()
-            .map_err(|e| e.to_string())
+            .map_err(|e| e.to_string())?;
+            if entry.is_some() {
+                conn.execute(
+                    "UPDATE tab_runtime_cache SET last_accessed_at = ?2 WHERE cache_key = ?1",
+                    params![key, now],
+                )
+                .map_err(|e| e.to_string())?;
+            }
+            Ok(entry)
+        })
+        .await
+    }
+
+    pub async fn list_tab_runtime_cache_metadata(&self) -> Result<Vec<TabRuntimeCacheMetadata>, String> {
+        self.with_conn(move |conn| {
+            let mut statement = conn
+                .prepare(
+                    "SELECT cache_key, row_count, column_count, byte_size, updated_at, created_at, last_accessed_at, owner_id \
+                     FROM tab_runtime_cache ORDER BY last_accessed_at ASC, cache_key ASC",
+                )
+                .map_err(|e| e.to_string())?;
+            let metadata = statement
+                .query_map([], |row| {
+                    Ok(TabRuntimeCacheMetadata {
+                        key: row.get(0)?,
+                        row_count: row.get(1)?,
+                        column_count: row.get(2)?,
+                        byte_size: row.get(3)?,
+                        updated_at: row.get(4)?,
+                        created_at: row.get(5)?,
+                        last_accessed_at: row.get(6)?,
+                        owner_id: row.get(7)?,
+                    })
+                })
+                .map_err(|e| e.to_string())?
+                .collect::<Result<Vec<_>, _>>()
+                .map_err(|e| e.to_string())?;
+            Ok(metadata)
+        })
+        .await
+    }
+
+    pub async fn prune_tab_runtime_cache(
+        &self,
+        live_keys: Vec<String>,
+        max_bytes: i64,
+        orphan_grace_ms: i64,
+        max_age_ms: Option<i64>,
+    ) -> Result<TabRuntimeCachePruneResult, String> {
+        let now = unix_timestamp_millis();
+        self.with_conn(move |conn| {
+            let live_keys: HashSet<String> = live_keys.into_iter().collect();
+            let mut statement = conn
+                .prepare(
+                    "SELECT cache_key, byte_size, created_at, last_accessed_at \
+                     FROM tab_runtime_cache ORDER BY last_accessed_at ASC, cache_key ASC",
+                )
+                .map_err(|e| e.to_string())?;
+            let entries = statement
+                .query_map([], |row| {
+                    Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?, row.get::<_, i64>(2)?, row.get::<_, i64>(3)?))
+                })
+                .map_err(|e| e.to_string())?
+                .collect::<Result<Vec<_>, _>>()
+                .map_err(|e| e.to_string())?;
+            drop(statement);
+
+            let mut total_bytes = entries.iter().map(|(_, bytes, _, _)| *bytes).sum::<i64>();
+            let mut deleted = HashSet::new();
+            let mut orphan_deletions = 0usize;
+            for (key, bytes, created_at, last_accessed_at) in &entries {
+                if live_keys.contains(key) {
+                    continue;
+                }
+                let orphan_expired = now.saturating_sub(*created_at) >= orphan_grace_ms.max(0);
+                let age_expired =
+                    max_age_ms.is_some_and(|max_age| now.saturating_sub(*last_accessed_at) >= max_age.max(0));
+                if orphan_expired || age_expired {
+                    deleted.insert(key.clone());
+                    total_bytes = total_bytes.saturating_sub(*bytes);
+                    if orphan_expired {
+                        orphan_deletions += 1;
+                    }
+                }
+            }
+
+            for (key, bytes, _, _) in &entries {
+                if total_bytes <= max_bytes.max(0) {
+                    break;
+                }
+                if live_keys.contains(key) || deleted.contains(key) {
+                    continue;
+                }
+                deleted.insert(key.clone());
+                total_bytes = total_bytes.saturating_sub(*bytes);
+            }
+
+            let deleted_bytes =
+                entries.iter().filter(|(key, _, _, _)| deleted.contains(key)).map(|(_, bytes, _, _)| *bytes).sum();
+            let transaction = conn.transaction().map_err(|e| e.to_string())?;
+            for key in &deleted {
+                transaction
+                    .execute("DELETE FROM tab_runtime_cache WHERE cache_key = ?1", [key])
+                    .map_err(|e| e.to_string())?;
+            }
+            transaction.commit().map_err(|e| e.to_string())?;
+            Ok(TabRuntimeCachePruneResult {
+                deleted_entries: deleted.len(),
+                deleted_bytes,
+                orphan_deletions,
+                remaining_entries: entries.len().saturating_sub(deleted.len()),
+                remaining_bytes: total_bytes,
+            })
+        })
+        .await
+    }
+
+    pub async fn delete_tab_runtime_cache_owner(&self, owner_id: &str) -> Result<usize, String> {
+        let owner_id = owner_id.to_string();
+        self.with_conn(move |conn| {
+            conn.execute("DELETE FROM tab_runtime_cache WHERE owner_id = ?1", [owner_id]).map_err(|e| e.to_string())
         })
         .await
     }
@@ -2535,6 +2715,7 @@ mod tests {
         ConnectionConfig, DatabaseConnectionInfo, DatabaseType, SshTunnelConfig, TransportLayerConfig,
     };
     use crate::saved_sql::SavedSqlFile;
+    use rusqlite::Connection;
     use std::time::{SystemTime, UNIX_EPOCH};
 
     fn temp_db_path(name: &str) -> std::path::PathBuf {
@@ -3256,7 +3437,10 @@ mod tests {
         let path = temp_db_path("tab-runtime-cache");
         let storage = Storage::open(&path).await.unwrap();
 
-        storage.save_tab_runtime_cache("tab:1:result", vec![1, 2, 3, 4], 10, 3).await.unwrap();
+        storage
+            .save_tab_runtime_cache("tab:1:result", vec![1, 2, 3, 4], 10, 3, Some("connection-1".to_string()))
+            .await
+            .unwrap();
         let entry = storage.load_tab_runtime_cache("tab:1:result").await.unwrap().unwrap();
 
         assert_eq!(entry.key, "tab:1:result");
@@ -3264,11 +3448,86 @@ mod tests {
         assert_eq!(entry.row_count, 10);
         assert_eq!(entry.column_count, 3);
         assert_eq!(entry.byte_size, 4);
+        assert_eq!(entry.owner_id.as_deref(), Some("connection-1"));
+        assert!(entry.created_at > 0);
+        assert!(entry.last_accessed_at >= entry.created_at);
 
         storage.delete_tab_runtime_cache("tab:1:result").await.unwrap();
         assert_eq!(storage.load_tab_runtime_cache("tab:1:result").await.unwrap(), None);
     }
 
+    #[tokio::test]
+    async fn tab_runtime_cache_pruning_retains_live_entries_and_enforces_byte_budget() {
+        let path = temp_db_path("tab-runtime-cache-prune");
+        let storage = Storage::open(&path).await.unwrap();
+        for (key, size) in [("live", 6usize), ("old", 5), ("new", 4)] {
+            storage.save_tab_runtime_cache(key, vec![1; size], 1, 1, Some("connection-1".to_string())).await.unwrap();
+        }
+        storage
+            .with_conn(|conn| {
+                conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = 1 WHERE cache_key = 'old'", [])
+                    .map_err(|e| e.to_string())?;
+                conn.execute("UPDATE tab_runtime_cache SET last_accessed_at = 2 WHERE cache_key = 'new'", [])
+                    .map_err(|e| e.to_string())?;
+                Ok(())
+            })
+            .await
+            .unwrap();
+
+        let result = storage.prune_tab_runtime_cache(vec!["live".to_string()], 10, i64::MAX, None).await.unwrap();
+
+        assert_eq!(result.deleted_entries, 1);
+        assert!(storage.load_tab_runtime_cache("live").await.unwrap().is_some());
+        assert!(storage.load_tab_runtime_cache("old").await.unwrap().is_none());
+        assert!(storage.load_tab_runtime_cache("new").await.unwrap().is_some());
+    }
+
+    #[tokio::test]
+    async fn tab_runtime_cache_pruning_respects_orphan_grace_period() {
+        let path = temp_db_path("tab-runtime-cache-orphan-grace");
+        let storage = Storage::open(&path).await.unwrap();
+        storage.save_tab_runtime_cache("fresh", vec![1], 1, 1, None).await.unwrap();
+        storage.save_tab_runtime_cache("crash-leftover", vec![2], 1, 1, None).await.unwrap();
+        storage
+            .with_conn(|conn| {
+                conn.execute("UPDATE tab_runtime_cache SET created_at = 1 WHERE cache_key = 'crash-leftover'", [])
+                    .map(|_| ())
+                    .map_err(|e| e.to_string())
+            })
+            .await
+            .unwrap();
+
+        let result = storage.prune_tab_runtime_cache(Vec::new(), 1024, 60_000, None).await.unwrap();
+
+        assert_eq!(result.orphan_deletions, 1);
+        assert!(storage.load_tab_runtime_cache("fresh").await.unwrap().is_some());
+        assert!(storage.load_tab_runtime_cache("crash-leftover").await.unwrap().is_none());
+    }
+
+    #[tokio::test]
+    async fn tab_runtime_cache_migrates_legacy_schema_without_dropping_entries() {
+        let path = temp_db_path("tab-runtime-cache-legacy-schema");
+        {
+            let connection = Connection::open(&path).unwrap();
+            connection
+                .execute(
+                    "CREATE TABLE tab_runtime_cache (cache_key TEXT PRIMARY KEY, payload BLOB NOT NULL, row_count INTEGER NOT NULL DEFAULT 0, column_count INTEGER NOT NULL DEFAULT 0, byte_size INTEGER NOT NULL DEFAULT 0, updated_at TEXT NOT NULL)",
+                    [],
+                )
+                .unwrap();
+            connection
+                .execute("INSERT INTO tab_runtime_cache VALUES ('legacy', X'0102', 1, 1, 2, '2026-01-01')", [])
+                .unwrap();
+        }
+
+        let storage = Storage::open(&path).await.unwrap();
+        let entry = storage.load_tab_runtime_cache("legacy").await.unwrap().unwrap();
+
+        assert_eq!(entry.payload, vec![1, 2]);
+        assert!(entry.created_at > 0);
+        assert!(entry.last_accessed_at >= entry.created_at);
+    }
+
     #[tokio::test]
     async fn saved_sql_summary_omits_sql_text_and_loads_file_on_demand() {
         let path = temp_db_path("saved-sql-summary");

crates/dbx-core/src/task_supervisor.rs

@@ -0,0 +1,111 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
+
+#[derive(Clone, Default)]
+pub struct TaskSupervisor {
+    accepting: Arc<AtomicBool>,
+    cancellation: CancellationToken,
+    tasks: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
+}
+
+impl TaskSupervisor {
+    pub fn new() -> Self {
+        Self {
+            accepting: Arc::new(AtomicBool::new(true)),
+            cancellation: CancellationToken::new(),
+            tasks: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    pub fn spawn_replace<F, Fut>(&self, key: impl Into<String>, task: F) -> bool
+    where
+        F: FnOnce(CancellationToken) -> Fut,
+        Fut: Future<Output = ()> + Send + 'static,
+    {
+        if !self.accepting.load(Ordering::Acquire) {
+            return false;
+        }
+        let handle = tokio::spawn(task(self.cancellation.child_token()));
+        if let Some(previous) = self.tasks.lock().unwrap_or_else(|error| error.into_inner()).insert(key.into(), handle)
+        {
+            previous.abort();
+        }
+        true
+    }
+
+    pub fn stop(&self, key: &str) {
+        if let Some(task) = self.tasks.lock().unwrap_or_else(|error| error.into_inner()).remove(key) {
+            task.abort();
+        }
+    }
+
+    pub fn stop_many<'a>(&self, keys: impl IntoIterator<Item = &'a str>) {
+        for key in keys {
+            self.stop(key);
+        }
+    }
+
+    pub fn active_count(&self) -> usize {
+        self.tasks.lock().unwrap_or_else(|error| error.into_inner()).len()
+    }
+
+    pub async fn shutdown(&self, deadline: Duration) {
+        self.accepting.store(false, Ordering::Release);
+        self.cancellation.cancel();
+        let mut handles: Vec<JoinHandle<()>> =
+            self.tasks.lock().unwrap_or_else(|error| error.into_inner()).drain().map(|(_, handle)| handle).collect();
+        let wait = async {
+            for handle in &mut handles {
+                let _ = handle.await;
+            }
+        };
+        if tokio::time::timeout(deadline, wait).await.is_err() {
+            for handle in handles {
+                handle.abort();
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::TaskSupervisor;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    #[tokio::test]
+    async fn replacement_aborts_previous_task() {
+        let supervisor = TaskSupervisor::new();
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let flag = cancelled.clone();
+        supervisor.spawn_replace("task", move |token| async move {
+            token.cancelled().await;
+            flag.store(true, Ordering::SeqCst);
+        });
+        supervisor.spawn_replace("task", |_| async {});
+        tokio::task::yield_now().await;
+        assert_eq!(supervisor.active_count(), 1);
+        assert!(!cancelled.load(Ordering::SeqCst));
+    }
+
+    #[tokio::test]
+    async fn shutdown_cancels_tasks_and_rejects_new_work() {
+        let supervisor = TaskSupervisor::new();
+        let cancelled = Arc::new(AtomicBool::new(false));
+        let flag = cancelled.clone();
+        supervisor.spawn_replace("task", move |token| async move {
+            token.cancelled().await;
+            flag.store(true, Ordering::SeqCst);
+        });
+        supervisor.shutdown(Duration::from_secs(1)).await;
+        assert!(cancelled.load(Ordering::SeqCst));
+        assert_eq!(supervisor.active_count(), 0);
+        assert!(!supervisor.spawn_replace("late", |_| async {}));
+    }
+}