fix: improve sandbox reliability under concurrent and long-lived workloads

  - Apply SQLite WAL mode, busy_timeout, and foreign_keys via
    SqliteConnectOptions so every pool connection gets the PRAGMAs,
    not just the first (fixes SQLITE_BUSY under concurrent access)
  - Don't consider sandboxes idle while exec sessions are active
  - Enable TCP keepalive on all upstream proxy connections to prevent
    NATs/firewalls from dropping idle SSE/streaming flows
  - Wrap smoltcp poll thread in catch_unwind to log panics instead of
    silently killing all guest networking
  - Log a warning when the connection limit is hit and a SYN is dropped

Author: HungKNguyen

Cargo.lock

@@ -12,7 +12,7 @@ use std::{
 };
 
 use microsandbox_migration::{Migrator, MigratorTrait};
-use sea_orm::{ConnectOptions, Database, DatabaseConnection};
+use sea_orm::DatabaseConnection;
 use tokio::sync::OnceCell;
 
 use crate::MicrosandboxResult;
@@ -126,24 +126,31 @@ async fn connect_and_migrate(
             db_path.display()
         ))
     })?;
-    let db_url = format!("sqlite://{db_path_str}?mode=rwc");
-
-    let mut opts = ConnectOptions::new(&db_url);
-    opts.max_connections(max_connections)
-        .connect_timeout(Duration::from_secs(connect_timeout_secs))
-        .sqlx_logging(false);
-
-    let conn = Database::connect(opts).await?;
-
-    // Enable WAL journal mode, busy timeout, and foreign key enforcement.
-    // WAL prevents SQLITE_BUSY when multiple processes (CLI + sandbox runtimes)
-    // access the same database concurrently.
-    use sea_orm::ConnectionTrait;
-    conn.execute(sea_orm::Statement::from_string(
-        sea_orm::DatabaseBackend::Sqlite,
-        microsandbox_utils::SQLITE_PRAGMAS,
-    ))
-    .await?;
+
+    // Use SqliteConnectOptions so WAL mode, busy_timeout, and foreign keys
+    // are applied to every connection the pool creates — not just the first.
+    // The old approach (executing PRAGMAs on the initial connection) left
+    // connections 2..N without busy_timeout, causing SQLITE_BUSY under
+    // concurrent sandbox access.
+    use sea_orm::sqlx::ConnectOptions as _;
+    use sea_orm::sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
+    use std::str::FromStr;
+
+    let sqlite_opts = SqliteConnectOptions::from_str(&format!("sqlite://{db_path_str}?mode=rwc"))
+        .map_err(|e| crate::MicrosandboxError::Custom(format!("database url: {e}")))?
+        .journal_mode(SqliteJournalMode::Wal)
+        .busy_timeout(Duration::from_secs(5))
+        .foreign_keys(true)
+        .disable_statement_logging();
+
+    let pool = SqlitePoolOptions::new()
+        .max_connections(max_connections)
+        .acquire_timeout(Duration::from_secs(connect_timeout_secs))
+        .connect_with(sqlite_opts)
+        .await
+        .map_err(|e| crate::MicrosandboxError::Custom(format!("database connect: {e}")))?;
+
+    let conn = sea_orm::SqlxSqliteConnector::from_sqlx_sqlite_pool(pool);
 
     Migrator::up(&conn, None).await?;
 

crates/microsandbox/lib/db/mod.rs

@@ -34,6 +34,7 @@ rustls = { workspace = true }
 rustls-native-certs = { workspace = true }
 rustls-pemfile = "2"
 serde = { workspace = true }
+socket2 = "0.5"
 smoltcp = { workspace = true, features = [
     "alloc",
     "async",

crates/network/Cargo.toml

@@ -213,17 +213,32 @@ impl SmoltcpNetwork {
             std::thread::Builder::new()
                 .name("smoltcp-poll".into())
                 .spawn(move || {
-                    stack::smoltcp_poll_loop(
-                        shared,
-                        poll_config,
-                        network_policy,
-                        dns_config,
-                        tls_state,
-                        published_ports,
-                        max_connections,
-                        tokio_handle,
-                        egress_handle,
-                    );
+                    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                        stack::smoltcp_poll_loop(
+                            shared,
+                            poll_config,
+                            network_policy,
+                            dns_config,
+                            tls_state,
+                            published_ports,
+                            max_connections,
+                            tokio_handle,
+                            egress_handle,
+                        );
+                    }));
+                    if let Err(e) = result {
+                        let msg = if let Some(s) = e.downcast_ref::<&str>() {
+                            s.to_string()
+                        } else if let Some(s) = e.downcast_ref::<String>() {
+                            s.clone()
+                        } else {
+                            "unknown panic".to_string()
+                        };
+                        tracing::error!(
+                            reason = %msg,
+                            "smoltcp poll thread panicked — all guest networking is dead"
+                        );
+                    }
                 })
                 .expect("failed to spawn smoltcp poll thread"),
         );

crates/network/lib/network.rs

@@ -54,6 +54,12 @@ async fn tcp_proxy_task(
     shared: Arc<SharedState>,
 ) -> io::Result<()> {
     let stream = TcpStream::connect(dst).await?;
+    // Enable TCP keepalive to prevent NATs/firewalls from dropping idle flows.
+    let sock = socket2::SockRef::from(&stream);
+    let keepalive = socket2::TcpKeepalive::new()
+        .with_time(std::time::Duration::from_secs(30))
+        .with_interval(std::time::Duration::from_secs(10));
+    let _ = sock.set_tcp_keepalive(&keepalive);
     let (mut server_rx, mut server_tx) = stream.into_split();
 
     let mut server_buf = vec![0u8; SERVER_READ_BUF_SIZE];

crates/network/lib/proxy.rs

@@ -273,8 +273,14 @@ pub fn smoltcp_poll_loop(
                             .evaluate_egress(dst, Protocol::Tcp)
                             .is_allow(),
                     };
-                    if allow && !conn_tracker.has_socket_for(&src, &dst) {
-                        conn_tracker.create_tcp_socket(src, dst, &mut sockets);
+                    if allow
+                        && !conn_tracker.has_socket_for(&src, &dst)
+                        && !conn_tracker.create_tcp_socket(src, dst, &mut sockets)
+                    {
+                        tracing::warn!(
+                            %src, %dst,
+                            "connection limit reached, dropping SYN (guest will see RST)"
+                        );
                     }
                     // Let smoltcp process — matching socket completes
                     // handshake, no socket means auto-RST.

crates/network/lib/stack.rs

@@ -150,7 +150,10 @@ async fn connect_upstream(dst: SocketAddr, sni: &str) -> io::Result<TcpStream> {
         let mut last_err = None;
         for addr in &addrs {
             match TcpStream::connect(addr).await {
-                Ok(stream) => return Ok(stream),
+                Ok(stream) => {
+                    set_tcp_keepalive(&stream);
+                    return Ok(stream);
+                }
                 Err(e) => {
                     tracing::debug!(addr = %addr, sni = %sni, error = %e, "upstream connect failed, trying next");
                     last_err = Some(e);
@@ -163,7 +166,23 @@ async fn connect_upstream(dst: SocketAddr, sni: &str) -> io::Result<TcpStream> {
     }
 
     // Fallback: use original destination IP from the guest's connection.
-    TcpStream::connect(dst).await
+    let stream = TcpStream::connect(dst).await?;
+    set_tcp_keepalive(&stream);
+    Ok(stream)
+}
+
+/// Enable TCP keepalive on an upstream connection.
+///
+/// Keeps long-lived connections (SSE streams, chunked responses) alive
+/// through NATs and firewalls that silently drop idle TCP flows.
+fn set_tcp_keepalive(stream: &TcpStream) {
+    let sock = socket2::SockRef::from(stream);
+    let keepalive = socket2::TcpKeepalive::new()
+        .with_time(std::time::Duration::from_secs(30))
+        .with_interval(std::time::Duration::from_secs(10));
+    if let Err(e) = sock.set_tcp_keepalive(&keepalive) {
+        tracing::debug!(error = %e, "failed to set TCP keepalive on upstream");
+    }
 }
 
 /// Bypass mode: plain TCP splice, no TLS termination.

crates/network/lib/tls/proxy.rs

@@ -42,14 +42,22 @@ impl HeartbeatReader {
 
     /// Check whether the sandbox is idle based on the heartbeat.
     ///
-    /// Returns `true` if `last_activity` is older than `timeout_secs`.
+    /// Returns `true` if `last_activity` is older than `timeout_secs` **and**
+    /// no exec sessions are running. A sandbox with active exec sessions is
+    /// never considered idle, even if the host hasn't sent serial data recently
+    /// (the guest may be making outbound API calls without host interaction).
+    ///
     /// Returns `false` if the heartbeat file doesn't exist (agent still booting).
     pub fn is_idle(&self, timeout_secs: u64) -> bool {
         let heartbeat = match self.read() {
             Some(hb) => hb,
             None => return false,
         };
 
+        if heartbeat.active_sessions > 0 {
+            return false;
+        }
+
         let elapsed = Utc::now()
             .signed_duration_since(heartbeat.last_activity)
             .num_seconds();

crates/runtime/lib/heartbeat.rs

@@ -14,7 +14,7 @@ use std::time::Duration;
 use microsandbox_db::entity::run as run_entity;
 use microsandbox_filesystem::{DynFileSystem, PassthroughConfig, PassthroughFs};
 use msb_krun::VmBuilder;
-use sea_orm::{ColumnTrait, ConnectOptions, Database, DatabaseConnection, EntityTrait, Set};
+use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, Set};
 use serde::Serialize;
 
 use crate::console::{AgentConsoleBackend, ConsoleSharedState};
@@ -684,25 +684,29 @@ async fn connect_db(
     db_path: &std::path::Path,
     connect_timeout_secs: u64,
 ) -> RuntimeResult<DatabaseConnection> {
+    use sea_orm::sqlx::ConnectOptions as _;
+    use sea_orm::sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
+    use std::str::FromStr;
+
     let url = format!("sqlite://{}?mode=rwc", db_path.display());
-    let opts = ConnectOptions::new(url)
+
+    // Use SqliteConnectOptions so WAL mode, busy_timeout, and foreign keys
+    // are set on every connection (not just the first in the pool).
+    let sqlite_opts = SqliteConnectOptions::from_str(&url)
+        .map_err(|e| RuntimeError::Custom(format!("database url: {e}")))?
+        .journal_mode(SqliteJournalMode::Wal)
+        .busy_timeout(Duration::from_secs(5))
+        .foreign_keys(true)
+        .disable_statement_logging();
+
+    let pool = SqlitePoolOptions::new()
         .max_connections(1)
-        .connect_timeout(Duration::from_secs(connect_timeout_secs))
-        .sqlx_logging(false)
-        .to_owned();
-    let db = Database::connect(opts)
+        .acquire_timeout(Duration::from_secs(connect_timeout_secs))
+        .connect_with(sqlite_opts)
         .await
         .map_err(|e| RuntimeError::Custom(format!("database connect: {e}")))?;
 
-    use sea_orm::ConnectionTrait;
-    db.execute(sea_orm::Statement::from_string(
-        sea_orm::DatabaseBackend::Sqlite,
-        microsandbox_utils::SQLITE_PRAGMAS,
-    ))
-    .await
-    .map_err(|e| RuntimeError::Custom(format!("database pragmas: {e}")))?;
-
-    Ok(db)
+    Ok(sea_orm::SqlxSqliteConnector::from_sqlx_sqlite_pool(pool))
 }
 
 /// Insert a run record into the database.