Initial impl from Cursor/Composer

- attempts to call winapi directly as unsafe block to allow
  deregistering ctrl+c handler

Author: rajatarya

hf_xet/Cargo.lock

@@ -42,10 +42,12 @@ tracing = "0.1"
 # Unix-specific dependencies
 [target.'cfg(unix)'.dependencies]
 signal-hook = "0.3"
+signal-hook-registry = "1.4"
 
 # Windows-specific dependencies
 [target.'cfg(windows)'.dependencies]
 ctrlc = "3.4"
+winapi = { version = "0.3", features = ["consoleapi", "wincon", "errhandlingapi"] }
 
 [features]
 default = ["no-default-cache"]                           # By default, hf_xet disables the disk cache. 

hf_xet/Cargo.toml

@@ -17,30 +17,85 @@ lazy_static! {
 }
 
 #[cfg(unix)]
-fn install_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
+lazy_static! {
+    static ref SIGINT_HANDLER_REGISTRATION: RwLock<Option<signal_hook_registry::SigId>> = RwLock::new(None);
+}
+
+#[cfg(windows)]
+lazy_static! {
+    static ref SIGINT_HANDLER_INSTALLED: AtomicBool = AtomicBool::new(false);
+}
+
+#[cfg(unix)]
+fn install_sigint_handler() -> Result<signal_hook_registry::SigId, MultithreadedRuntimeError> {
     use signal_hook::consts::SIGINT;
-    use signal_hook::flag;
 
     // Register the SIGINT handler to set our atomic flag.
-    // Using `signal_hook::flag::register` allows us to set the atomic flag when SIGINT is received.
-    flag::register(SIGINT, SIGINT_DETECTED.clone()).map_err(|e| {
+    // Using signal_hook_registry directly to get a SigId that we can unregister later.
+    let flag = SIGINT_DETECTED.clone();
+    let signal_id = unsafe {
+        signal_hook_registry::register(SIGINT, move || {
+            flag.store(true, Ordering::SeqCst);
+        })
+    }
+    .map_err(|e| {
         MultithreadedRuntimeError::Other(format!("Initialization Error: Unable to register SIGINT handler {e:?}"))
     })?;
 
-    Ok(())
+    Ok(signal_id)
+}
+
+#[cfg(windows)]
+extern "system" fn console_ctrl_handler(
+    ctrl_type: winapi::shared::minwindef::DWORD,
+) -> winapi::shared::minwindef::BOOL {
+    use winapi::um::wincon;
+
+    // Only handle CTRL_C_EVENT
+    if ctrl_type == wincon::CTRL_C_EVENT {
+        // Check if we have active operations
+        let has_active_ops = {
+            let guard = MULTITHREADED_RUNTIME.read().unwrap();
+            if let Some((runtime_pid, ref runtime)) = *guard {
+                runtime_pid == std::process::id() && runtime.external_executor_count() > 0
+            } else {
+                false
+            }
+        };
+
+        if has_active_ops {
+            // We have active operations, handle it ourselves
+            SIGINT_DETECTED.store(true, Ordering::SeqCst);
+            winapi::shared::minwindef::TRUE
+        } else {
+            // No active operations, let Python's handler (or default) handle it
+            // Return FALSE to continue handler chain
+            winapi::shared::minwindef::FALSE
+        }
+    } else {
+        // For other control events, let default handler process them
+        winapi::shared::minwindef::FALSE
+    }
 }
 
 #[cfg(windows)]
 fn install_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
-    // On Windows, use ctrlc crate.
-    // This sets a callback to run on Ctrl-C:
-    let sigint_detected_flag = SIGINT_DETECTED.clone();
-    ctrlc::set_handler(move || {
-        sigint_detected_flag.store(true, Ordering::SeqCst);
-    })
-    .map_err(|e| {
-        MultithreadedRuntimeError::Other(format!("Initialization Error: Unable to register SIGINT handler {e:?}"))
-    })?;
+    use winapi::um::consoleapi::SetConsoleCtrlHandler;
+    use winapi::um::wincon::CTRL_C_EVENT;
+
+    // Install our handler using Windows API directly (instead of ctrlc)
+    // Our handler checks if operations are active:
+    // - If active: handles Ctrl+C and returns TRUE (stops propagation)
+    // - If not active: returns FALSE (allows Python's handler or default to handle it)
+    // This way we don't need to save/restore Python's handler - we just delegate to it
+    unsafe {
+        if SetConsoleCtrlHandler(Some(console_ctrl_handler), winapi::shared::minwindef::TRUE) == 0 {
+            let error = winapi::um::errhandlingapi::GetLastError();
+            return Err(MultithreadedRuntimeError::Other(format!(
+                "Initialization Error: Unable to register SIGINT handler. Windows error: {error}"
+            )));
+        }
+    }
     Ok(())
 }
 
@@ -54,7 +109,20 @@ fn check_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
     let pid = std::process::id();
 
     if stored_pid == pid {
-        return Ok(());
+        // Check if handler is already registered
+        #[cfg(unix)]
+        {
+            let reg = SIGINT_HANDLER_REGISTRATION.read().unwrap();
+            if reg.is_some() {
+                return Ok(());
+            }
+        }
+        #[cfg(windows)]
+        {
+            if SIGINT_HANDLER_INSTALLED.load(Ordering::SeqCst) {
+                return Ok(());
+            }
+        }
     }
 
     // Need to install it; acquire a lock to do so.
@@ -63,10 +131,33 @@ fn check_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
     // If another thread beat us to it while we're waiting for the lock.
     let stored_pid = SIGINT_HANDLER_INSTALL_PID.0.load(Ordering::SeqCst);
     if stored_pid == pid {
-        return Ok(());
+        #[cfg(unix)]
+        {
+            let reg = SIGINT_HANDLER_REGISTRATION.read().unwrap();
+            if reg.is_some() {
+                return Ok(());
+            }
+        }
+        #[cfg(windows)]
+        {
+            if SIGINT_HANDLER_INSTALLED.load(Ordering::SeqCst) {
+                return Ok(());
+            }
+        }
+    }
+
+    #[cfg(unix)]
+    {
+        let signal_id = install_sigint_handler()?;
+        let mut reg = SIGINT_HANDLER_REGISTRATION.write().unwrap();
+        *reg = Some(signal_id);
     }
 
-    install_sigint_handler()?;
+    #[cfg(windows)]
+    {
+        install_sigint_handler()?;
+        SIGINT_HANDLER_INSTALLED.store(true, Ordering::SeqCst);
+    }
 
     // Finally, store that we have installed it successfully.
     SIGINT_HANDLER_INSTALL_PID.0.store(pid, Ordering::SeqCst);
@@ -94,6 +185,65 @@ fn in_sigint_shutdown() -> bool {
     SIGINT_DETECTED.load(Ordering::Relaxed)
 }
 
+#[cfg(unix)]
+fn restore_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
+    use signal_hook_registry;
+
+    let mut reg = SIGINT_HANDLER_REGISTRATION.write().unwrap();
+    if let Some(signal_id) = reg.take() {
+        signal_hook_registry::unregister(signal_id);
+    }
+    Ok(())
+}
+
+#[cfg(windows)]
+fn restore_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
+    use winapi::um::consoleapi::SetConsoleCtrlHandler;
+
+    // Remove our handler by unregistering it
+    // This allows Python's handler (if any) or the default handler to take over
+    unsafe {
+        if SetConsoleCtrlHandler(Some(console_ctrl_handler), winapi::shared::minwindef::FALSE) == 0 {
+            let error = winapi::um::errhandlingapi::GetLastError();
+            // Log but don't fail - handler removal is best effort
+            if cfg!(debug_assertions) {
+                eprintln!("[debug] Warning: Failed to unregister console control handler. Windows error: {error}");
+            }
+        }
+    }
+    SIGINT_HANDLER_INSTALLED.store(false, Ordering::SeqCst);
+    Ok(())
+}
+
+fn maybe_restore_sigint_handler() -> Result<(), MultithreadedRuntimeError> {
+    // Check if runtime exists and has no active operations
+    let guard = MULTITHREADED_RUNTIME.read().unwrap();
+    if let Some((runtime_pid, ref runtime)) = *guard
+        && runtime_pid == std::process::id()
+        && runtime.external_executor_count() == 0
+    {
+        // Check if handler is installed
+        let should_restore = {
+            #[cfg(unix)]
+            {
+                let reg = SIGINT_HANDLER_REGISTRATION.read().unwrap();
+                reg.is_some()
+            }
+            #[cfg(windows)]
+            {
+                SIGINT_HANDLER_INSTALLED.load(Ordering::SeqCst)
+            }
+        };
+
+        if should_restore {
+            // No active operations, restore handler
+            drop(guard);
+            restore_sigint_handler()?;
+        }
+    }
+    Ok(())
+}
+
 fn signal_check_background_loop() {
     const SIGNAL_CHECK_INTERVAL: Duration = Duration::from_millis(250);
 
@@ -223,6 +373,10 @@ where
         return Err(PyKeyboardInterrupt::new_err(()));
     }
 
+    // After operation completes, check if we should restore the signal handler
+    // This allows Python's original handler to be restored when no operations are active
+    let _ = maybe_restore_sigint_handler();
+
     // Now return the result.
     result
 }