[pull] main from NVIDIA:main

.github/workflows/windows-msvc.yml

@@ -6,7 +6,7 @@ jobs:
     runs-on: windows-2025
     steps:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
-      - uses: jdx/mise-action@7e36c90d9ab29c415a2384db3006f3ec8a8cc654 # v4.2.4
+      - uses: jdx/mise-action@c2a87611a18de5b3828c5652fe268e992400cb5c # v4.3.0
         with:
           install: false
           experimental: true
@@ -31,7 +31,7 @@ jobs:
     if: false # flip to true once the runner is online
     steps:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
-      - uses: jdx/mise-action@7e36c90d9ab29c415a2384db3006f3ec8a8cc654 # v4.2.4
+      - uses: jdx/mise-action@c2a87611a18de5b3828c5652fe268e992400cb5c # v4.3.0
         with:
           install: false
           experimental: true

crates/openshell-server/src/lib.rs

@@ -62,7 +62,10 @@ use std::net::SocketAddr;
 use std::path::Path;
 #[cfg(test)]
 use std::sync::LazyLock;
-use std::sync::{Arc, Mutex};
+use std::sync::{
+    Arc, Mutex,
+    atomic::{AtomicBool, Ordering},
+};
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::watch;
@@ -295,6 +298,10 @@ pub struct ServerState {
     /// query session state to surface supervisor readiness.
     pub supervisor_sessions: Arc<supervisor_session::SupervisorSessionRegistry>,
 
+    /// Set once graceful gateway shutdown begins so stream handlers can
+    /// distinguish expected transport closes from runtime failures.
+    pub(crate) gateway_shutting_down: AtomicBool,
+
     /// Validated built-in and operator-registered supervisor middleware.
     pub middleware_registry: Arc<MiddlewareRegistry>,
 
@@ -411,6 +418,7 @@ impl ServerState {
             ssh_connections_by_sandbox: Mutex::new(HashMap::new()),
             settings_mutex: tokio::sync::Mutex::new(()),
             supervisor_sessions,
+            gateway_shutting_down: AtomicBool::new(false),
             extension_mint_limiter: auth::extension_mint_limit::ExtensionMintLimiter::default(),
             middleware_registry: Arc::new(MiddlewareRegistry::default()),
             oidc_cache,
@@ -802,6 +810,7 @@ pub(crate) async fn run_server(
 
     shutdown_signal().await;
     info!("Shutdown signal received; stopping gateway");
+    state.gateway_shutting_down.store(true, Ordering::Release);
     let _ = shutdown_tx.send(true);
 
     for task in listener_tasks {

crates/openshell-server/src/supervisor_session.rs

@@ -3,6 +3,7 @@
 
 use std::collections::HashMap;
 use std::pin::Pin;
+use std::sync::atomic::Ordering;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
@@ -666,9 +667,23 @@ async fn expected_transport_close_during_session_teardown(
     let session_no_longer_current = !state
         .supervisor_sessions
         .is_current_session(sandbox_id, session_id);
+    expected_transport_close_during_session_state(
+        status,
+        state.gateway_shutting_down.load(Ordering::Acquire),
+        session_no_longer_current,
+        sandbox_is_terminating_or_gone(state, sandbox_id).await,
+    )
+}
+
+fn expected_transport_close_during_session_state(
+    status: &Status,
+    gateway_shutting_down: bool,
+    session_no_longer_current: bool,
+    sandbox_terminating_or_gone: bool,
+) -> bool {
     expected_transport_close_during_shutdown(
         status,
-        session_no_longer_current || sandbox_is_terminating_or_gone(state, sandbox_id).await,
+        gateway_shutting_down || session_no_longer_current || sandbox_terminating_or_gone,
     )
 }
 
@@ -1447,6 +1462,16 @@ mod tests {
         assert!(!expected_transport_close_during_shutdown(&status, true));
     }
 
+    #[test]
+    fn gateway_shutdown_makes_session_transport_close_nonfatal() {
+        let status =
+            Status::unknown("h2 protocol error: error reading a body from connection: broken pipe");
+
+        assert!(expected_transport_close_during_session_state(
+            &status, true, false, false,
+        ));
+    }
+
     #[test]
     fn sandbox_proto_terminating_detects_deleting_phase() {
         let mut sandbox = sandbox_record("sbx-1", "sandbox-one");

crates/openshell-supervisor-network/src/policy_dns/resolver.rs

@@ -315,6 +315,23 @@ mod tests {
     use openshell_core::net::set_tcp_nodelay_best_effort;
     use tokio::net::TcpListener;
 
+    async fn bind_dns_test_server() -> (UdpSocket, TcpListener, SocketAddr) {
+        const MAX_BIND_ATTEMPTS: usize = 100;
+
+        for _ in 0..MAX_BIND_ATTEMPTS {
+            // Port zero only guarantees availability for the protocol being bound.
+            let tcp = TcpListener::bind("127.0.0.1:0").await.unwrap();
+            let server = tcp.local_addr().unwrap();
+            match UdpSocket::bind(server).await {
+                Ok(udp) => return (udp, tcp, server),
+                Err(error) if error.kind() == std::io::ErrorKind::AddrInUse => {}
+                Err(error) => panic!("failed to bind DNS test UDP socket: {error}"),
+            }
+        }
+
+        panic!("failed to bind DNS test TCP and UDP sockets to the same port");
+    }
+
     #[test]
     fn resolver_address_deduplication_preserves_answer_order() {
         let mut addresses = vec![
@@ -383,9 +400,7 @@ mod tests {
 
     #[tokio::test]
     async fn truncated_udp_retries_over_tcp_and_follows_cname() {
-        let udp = UdpSocket::bind("127.0.0.1:0").await.unwrap();
-        let server = udp.local_addr().unwrap();
-        let tcp = TcpListener::bind(server).await.unwrap();
+        let (udp, tcp, server) = bind_dns_test_server().await;
 
         let udp_task = tokio::spawn(async move {
             let mut wire = [0_u8; MAX_DNS_MESSAGE_BYTES];