Copy entry points and Jupyter data directories into ephemeral environments

crates/uv-fs/src/lib.rs

@@ -84,6 +84,8 @@ pub async fn read_to_string_transcode(path: impl AsRef<Path>) -> std::io::Result
 /// junction at the same path.
 ///
 /// Note that because junctions are used, the source must be a directory.
+///
+/// Changes to this function should be reflected in [`create_symlink`].
 #[cfg(windows)]
 pub fn replace_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
     // If the source is a file, we can't create a junction
@@ -138,6 +140,38 @@ pub fn replace_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io:
     }
 }
 
+/// Create a symlink at `dst` pointing to `src`.
+///
+/// On Windows, this uses the `junction` crate to create a junction point.
+///
+/// Note that because junctions are used, the source must be a directory.
+///
+/// Changes to this function should be reflected in [`replace_symlink`].
+#[cfg(windows)]
+pub fn create_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
+    // If the source is a file, we can't create a junction
+    if src.as_ref().is_file() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!(
+                "Cannot create a junction for {}: is not a directory",
+                src.as_ref().display()
+            ),
+        ));
+    }
+
+    junction::create(
+        dunce::simplified(src.as_ref()),
+        dunce::simplified(dst.as_ref()),
+    )
+}
+
+/// Create a symlink at `dst` pointing to `src`.
+#[cfg(unix)]
+pub fn create_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
+    fs_err::os::unix::fs::symlink(src.as_ref(), dst.as_ref())
+}
+
 #[cfg(unix)]
 pub fn remove_symlink(path: impl AsRef<Path>) -> std::io::Result<()> {
     fs_err::remove_file(path.as_ref())

crates/uv-trampoline-builder/src/lib.rs

@@ -41,6 +41,7 @@ const MAGIC_NUMBER_SIZE: usize = 4;
 pub struct Launcher {
     pub kind: LauncherKind,
     pub python_path: PathBuf,
+    payload: Vec<u8>,
 }
 
 impl Launcher {
@@ -109,11 +110,69 @@ impl Launcher {
             String::from_utf8(buffer).map_err(|err| Error::InvalidPath(err.utf8_error()))?,
         );
 
+        #[allow(clippy::cast_possible_truncation)]
+        let file_size = {
+            let raw_length = file
+                .seek(io::SeekFrom::End(0))
+                .map_err(|e| Error::InvalidLauncherSeek("size probe".into(), 0, e))?;
+
+            if raw_length > usize::MAX as u64 {
+                return Err(Error::InvalidDataLength(raw_length));
+            }
+
+            // SAFETY: Above we guarantee the length is less than uszie
+            raw_length as usize
+        };
+
+        // Read the payload
+        file.seek(io::SeekFrom::Start(0))
+            .map_err(|e| Error::InvalidLauncherSeek("rewind".into(), 0, e))?;
+        let payload_len =
+            file_size.saturating_sub(MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE + path_length);
+        let mut buffer = vec![0u8; payload_len];
+        file.read_exact(&mut buffer)
+            .map_err(|err| Error::InvalidLauncherRead("payload".into(), err))?;
+
         Ok(Some(Self {
             kind,
+            payload: buffer,
             python_path: path,
         }))
     }
+
+    pub fn write_to_file(self, file: &mut File) -> Result<(), Error> {
+        let python_path = self.python_path.simplified_display().to_string();
+
+        if python_path.len() > MAX_PATH_LENGTH as usize {
+            return Err(Error::InvalidPathLength(
+                u32::try_from(python_path.len()).expect("path length already checked"),
+            ));
+        }
+
+        let mut launcher: Vec<u8> = Vec::with_capacity(
+            self.payload.len() + python_path.len() + PATH_LENGTH_SIZE + MAGIC_NUMBER_SIZE,
+        );
+        launcher.extend_from_slice(&self.payload);
+        launcher.extend_from_slice(python_path.as_bytes());
+        launcher.extend_from_slice(
+            &u32::try_from(python_path.len())
+                .expect("file path should be smaller than 4GB")
+                .to_le_bytes(),
+        );
+        launcher.extend_from_slice(self.kind.magic_number());
+
+        file.write_all(&launcher)?;
+        Ok(())
+    }
+
+    #[must_use]
+    pub fn with_python_path(self, path: PathBuf) -> Self {
+        Self {
+            kind: self.kind,
+            payload: self.payload,
+            python_path: path,
+        }
+    }
 }
 
 /// The kind of trampoline launcher to create.
@@ -177,6 +236,8 @@ pub enum Error {
     Io(#[from] io::Error),
     #[error("Only paths with a length up to 32KB are supported but found a length of {0} bytes")]
     InvalidPathLength(u32),
+    #[error("Only data with a length up to usize is supported but found a length of {0} bytes")]
+    InvalidDataLength(u64),
     #[error("Failed to parse executable path")]
     InvalidPath(#[source] Utf8Error),
     #[error("Failed to seek to {0} at offset {1}")]

crates/uv/src/commands/project/environment.rs

@@ -78,6 +78,20 @@ impl EphemeralEnvironment {
         )?;
         Ok(())
     }
+
+    /// Returns the path to the environment's scripts directory.
+    pub(crate) fn scripts(&self) -> &Path {
+        self.0.scripts()
+    }
+
+    /// Returns the path to the environment's Python executable.
+    pub(crate) fn sys_executable(&self) -> &Path {
+        self.0.interpreter().sys_executable()
+    }
+
+    pub(crate) fn sys_prefix(&self) -> &Path {
+        self.0.interpreter().sys_prefix()
+    }
 }
 
 /// A [`PythonEnvironment`] stored in the cache.

crates/uv/src/commands/project/run.rs

@@ -9,8 +9,9 @@ use anyhow::{Context, anyhow, bail};
 use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
+use thiserror::Error;
 use tokio::process::Command;
-use tracing::{debug, warn};
+use tracing::{debug, trace, warn};
 use url::Url;
 
 use uv_cache::Cache;
@@ -22,7 +23,7 @@ use uv_configuration::{
 };
 use uv_distribution_types::Requirement;
 use uv_fs::which::is_executable;
-use uv_fs::{PythonExt, Simplified};
+use uv_fs::{PythonExt, Simplified, create_symlink};
 use uv_installer::{SatisfiesResult, SitePackages};
 use uv_normalize::{DefaultExtras, DefaultGroups, PackageName};
 use uv_python::{
@@ -1071,6 +1072,67 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
                 requirements_site_packages.escape_for_python(),
             ))?;
 
+            // N.B. The order here matters — earlier interpreters take precedence over the
+            // later ones.
+            for interpreter in [requirements_env.interpreter(), &base_interpreter] {
+                // Copy each entrypoint from the base environments to the ephemeral environment,
+                // updating the Python executable target to ensure they run in the ephemeral
+                // environment.
+                for entry in fs_err::read_dir(interpreter.scripts())? {
+                    let entry = entry?;
+                    if !entry.file_type()?.is_file() {
+                        continue;
+                    }
+                    match copy_entrypoint(
+                        &entry.path(),
+                        &ephemeral_env.scripts().join(entry.file_name()),
+                        interpreter.sys_executable(),
+                        ephemeral_env.sys_executable(),
+                    ) {
+                        Ok(()) => {}
+                        // If the entrypoint already exists, skip it.
+                        Err(CopyEntrypointError::Io(err))
+                            if err.kind() == std::io::ErrorKind::AlreadyExists =>
+                        {
+                            trace!(
+                                "Skipping copy of entrypoint `{}`: already exists",
+                                &entry.path().display()
+                            );
+                        }
+                        Err(err) => return Err(err.into()),
+                    }
+                }
+
+                // Link data directories from the base environment to the ephemeral environment.
+                //
+                // This is critical for Jupyter Lab, which cannot operate without the files it
+                // writes to `<prefix>/share/jupyter`.
+                //
+                // See https://github.com/jupyterlab/jupyterlab/issues/17716
+                for dir in &["etc/jupyter", "share/jupyter"] {
+                    let source = interpreter.sys_prefix().join(dir);
+                    if !matches!(source.try_exists(), Ok(true)) {
+                        continue;
+                    }
+                    if !source.is_dir() {
+                        continue;
+                    }
+                    let target = ephemeral_env.sys_prefix().join(dir);
+                    if let Some(parent) = target.parent() {
+                        fs_err::create_dir_all(parent)?;
+                    }
+                    match create_symlink(&source, &target) {
+                        Ok(()) => trace!(
+                            "Created link for {} -> {}",
+                            target.user_display(),
+                            source.user_display()
+                        ),
+                        Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+                        Err(err) => return Err(err.into()),
+                    }
+                }
+            }
+
             // Write the `sys.prefix` of the parent environment to the `extends-environment` key of the `pyvenv.cfg`
             // file. This helps out static-analysis tools such as ty (see docs on
             // `CachedEnvironment::set_parent_environment`).
@@ -1669,3 +1731,92 @@ fn read_recursion_depth_from_environment_variable() -> anyhow::Result<u32> {
         .parse::<u32>()
         .with_context(|| format!("invalid value for {}", EnvVars::UV_RUN_RECURSION_DEPTH))
 }
+
+#[derive(Error, Debug)]
+enum CopyEntrypointError {
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+    #[cfg(windows)]
+    #[error(transparent)]
+    Trampoline(#[from] uv_trampoline_builder::Error),
+}
+
+/// Create a copy of the entrypoint at `source` at `target`, if it has a Python shebang, replacing
+/// the previous Python executable with a new one.
+///
+/// This is a no-op if the target already exists.
+///
+/// Note on Windows, the entrypoints do not use shebangs and require a rewrite of the trampoline.
+#[cfg(unix)]
+fn copy_entrypoint(
+    source: &Path,
+    target: &Path,
+    previous_executable: &Path,
+    python_executable: &Path,
+) -> Result<(), CopyEntrypointError> {
+    use std::io::Write;
+    use std::os::unix::fs::PermissionsExt;
+
+    use fs_err::os::unix::fs::OpenOptionsExt;
+
+    let contents = fs_err::read_to_string(source)?;
+
+    let Some(contents) = contents
+        // Check for a relative path or relocatable shebang
+        .strip_prefix(
+            r#"#!/bin/sh
+'''exec' "$(dirname -- "$(realpath -- "$0")")"/'python' "$0" "$@"
+' '''
+"#,
+        )
+        // Or an absolute path shebang
+        .or_else(|| contents.strip_prefix(&format!("#!{}\n", previous_executable.display())))
+    else {
+        // If it's not a Python shebang, we'll skip it
+        trace!(
+            "Skipping copy of entrypoint `{}`: does not start with expected shebang",
+            source.user_display()
+        );
+        return Ok(());
+    };
+
+    let contents = format!("#!{}\n{}", python_executable.display(), contents);
+    let mode = fs_err::metadata(source)?.permissions().mode();
+    let mut file = fs_err::OpenOptions::new()
+        .create_new(true)
+        .write(true)
+        .mode(mode)
+        .open(target)?;
+    file.write_all(contents.as_bytes())?;
+
+    trace!("Updated entrypoint at {}", target.user_display());
+
+    Ok(())
+}
+
+/// Create a copy of the entrypoint at `source` at `target`, if it's a Python script launcher,
+/// replacing the target Python executable with a new one.
+#[cfg(windows)]
+fn copy_entrypoint(
+    source: &Path,
+    target: &Path,
+    _previous_executable: &Path,
+    python_executable: &Path,
+) -> Result<(), CopyEntrypointError> {
+    use uv_trampoline_builder::Launcher;
+
+    let Some(launcher) = Launcher::try_from_path(source)? else {
+        return Ok(());
+    };
+
+    let launcher = launcher.with_python_path(python_executable.to_path_buf());
+    let mut file = fs_err::OpenOptions::new()
+        .create_new(true)
+        .write(true)
+        .open(target)?;
+    launcher.write_to_file(&mut file)?;
+
+    trace!("Updated entrypoint at {}", target.user_display());
+
+    Ok(())
+}