Fix vendored uv drift-keep after uv remove (#1287)

crates/socket-patch-cli/CLI_CONTRACT.md

@@ -822,7 +822,11 @@ worse, lets a warm cache silently serve unpatched bytes):
   that no longer exists at all — the user removed the dependency — is not drift: it warns
   `vendor_lock_entry_removed` and the artifact and entry are kept unless every wired file that exists
   was read and none mentions the uuid in any spelling (an unreadable lock keeps them), so `rollback` / `remove` / `scan --prune` clean up
-  after `npm uninstall` / `yarn remove` / `pnpm remove` / `bun remove`), removes the artifacts, prunes the
+  after `npm uninstall` / `yarn remove` / `pnpm remove` / `bun remove`; in uv projects and PEP 723
+  script locks the same holds after `uv remove` of the package, and after `uv remove` of the parent
+  of a TRANSITIVE vendored package, whose `[tool.uv]` override + source and `[manifest]` records
+  uv leaves in place: those are socket-patch's own records, which the revert restores, so only a
+  reference outside them counts as drift, #1287), removes the artifacts, prunes the
   ledger, sweeps orphan uuid dirs, and (v5.0) prunes the now-empty `.socket/vendor/<eco>/` and
   `.socket/vendor/` levels — `.socket/` itself is removed by the lock guard when nothing else is
   left. It works without a manifest: with no manifest and no ledger it is a clean exit-0 no-op.

crates/socket-patch-cli/tests/e2e_vendor_pypi_build.rs

@@ -1216,6 +1216,24 @@ fn uv_vendor_revert_after_uv_remove() {
     );
 }
 
+/// #1287: six vendored TRANSITIVELY (through `[tool.uv]
+/// override-dependencies` + sources, here beside a user
+/// `constraint-dependencies` pin), then `uv remove python-dateutil` drops
+/// six's `[[package]]` unit but leaves the `[tool.uv]` lines and the lock's
+/// `[manifest]` records socket-patch wrote. `vendor --revert` must read the
+/// vanished unit as removed and unwind the rest, instead of drift-keeping
+/// everything (which left `vendor --check` red with a prune remedy that
+/// changed nothing).
+#[test]
+#[serial_test::serial]
+fn uv_vendor_revert_after_uv_remove_of_the_transitive_parent() {
+    uv_relock_then_revert(
+        "uv-remove-parent",
+        "[project]\nname = \"vendor-capstone\"\nversion = \"0.1.0\"\nrequires-python = \">=3.9\"\ndependencies = [\"python-dateutil==2.8.2\", \"attrs>=20\"]\n\n[tool.uv]\nconstraint-dependencies = [\"six==1.16.0\"]\n",
+        &["remove", "-q", "python-dateutil"],
+    );
+}
+
 /// #821: six in a PEP 735 dev group, then `uv add --dev zipp` rewrites the
 /// whole `requires-dev` group line.
 #[test]

crates/socket-patch-cli/tests/mode_migration_pypi.rs

@@ -852,6 +852,67 @@ fn uv_remove_script_six(root: &Path) {
 /// `vendor --check` stayed red and its own `scan --prune` remedy looped.
 #[tokio::test]
 async fn script_lock_unwinds_after_uv_remove_script() {
+    assert_script_lock_unwinds(stage_script_lock, uv_remove_script_six).await;
+}
+
+/// A PEP 723 script whose six arrives only TRANSITIVELY (through
+/// python-dateutil), with its `.py.lock`; returns its wiring files.
+fn stage_transitive_script_lock(root: &Path) -> &'static [&'static str] {
+    std::fs::write(
+        root.join("job.py"),
+        "# /// script\n# requires-python = \">=3.9\"\n# dependencies = [\"python-dateutil==2.8.2\"]\n# ///\nimport six\n",
+    )
+    .unwrap();
+    std::fs::write(
+        root.join("job.py.lock"),
+        format!(
+            "version = 1\nrevision = 3\nrequires-python = \">=3.9\"\n\n[manifest]\nrequirements = [{{ name = \"python-dateutil\", specifier = \"==2.8.2\" }}]\n\n[[package]]\nname = \"python-dateutil\"\nversion = \"2.8.2\"\nsource = {{ registry = \"https://pypi.org/simple\" }}\ndependencies = [{{ name = \"six\" }}]\nwheels = [{{ url = \"https://files.pythonhosted.org/python_dateutil-2.8.2-py2.py3-none-any.whl\", hash = \"sha256:{}\" }}]\n\n[[package]]\nname = \"six\"\nversion = \"1.16.0\"\nsource = {{ registry = \"https://pypi.org/simple\" }}\nwheels = [{{ url = \"https://files.pythonhosted.org/six-1.16.0-py2.py3-none-any.whl\", hash = \"sha256:{WHEEL_SHA}\" }}]\n",
+            "d".repeat(64)
+        ),
+    )
+    .unwrap();
+    &["job.py", "job.py.lock"]
+}
+
+/// What `uv remove --script job.py python-dateutil` leaves of the vendored
+/// [`stage_transitive_script_lock`] (checked against uv 0.11.19): the
+/// dependency and both lock units go, while the script's `[tool.uv]`
+/// override + source and the lock's `[manifest] overrides` socket-patch
+/// wrote stay, still naming the vendored wheel.
+fn uv_remove_script_parent(root: &Path) {
+    let script = std::fs::read_to_string(root.join("job.py")).unwrap();
+    std::fs::write(
+        root.join("job.py"),
+        script.replace("\"python-dateutil==2.8.2\"", ""),
+    )
+    .unwrap();
+    let lock = std::fs::read_to_string(root.join("job.py.lock")).unwrap();
+    let overrides = lock
+        .lines()
+        .find(|line| line.starts_with("overrides = "))
+        .expect("the vendored lock carries the override");
+    std::fs::write(
+        root.join("job.py.lock"),
+        format!(
+            "version = 1\nrevision = 3\nrequires-python = \">=3.9\"\n\n[manifest]\n{overrides}\n"
+        ),
+    )
+    .unwrap();
+}
+
+/// #1287: the script lane of a vendored TRANSITIVE package whose parent
+/// `uv remove --script` dropped: every unwind retires the entry instead of
+/// drift-keeping it, and `vendor --check` turns green.
+#[tokio::test]
+async fn transitive_script_lock_unwinds_after_uv_remove_of_its_parent() {
+    assert_script_lock_unwinds(stage_transitive_script_lock, uv_remove_script_parent).await;
+}
+
+/// Vendor the script lock `stage` writes, apply `remove` (a `uv remove
+/// --script`), then every unwind must retire the entry (see
+/// [`script_lock_unwinds_after_uv_remove_script`]). Files the unwind
+/// restores must no longer name the vendored wheel.
+async fn assert_script_lock_unwinds(stage: StageFn, remove: fn(&Path)) {
     let server = MockServer::start().await;
     mount_hosted_api(&server, true).await;
     let uri = server.uri();
@@ -876,13 +937,16 @@ async fn script_lock_unwinds_after_uv_remove_script() {
         hosted_scan_args(&uri),
     ] {
         let (_tmp, root) = project();
-        let files = stage_script_lock(&root);
+        let files = stage(&root);
         vendor_project(&root, files);
-        uv_remove_script_six(&root);
+        remove(&root);
         let removed: Vec<String> = files
             .iter()
             .map(|f| std::fs::read_to_string(root.join(f)).unwrap())
             .collect();
+        // What remains after the unwind: the user's own content, with any
+        // surviving socket-patch wiring gone.
+        let transitive = removed.iter().any(|text| text.contains(UUID));
         let (code, env) = run_cli(&root, &["vendor", "--check"], &[]);
         assert_eq!(code, 1, "{unwind:?}: the removal is flagged first: {env:#}");
 
@@ -915,11 +979,15 @@ async fn script_lock_unwinds_after_uv_remove_script() {
             std::fs::read_to_string(root.join(".socket/vendor/state.json")).unwrap_or_default();
         assert!(!ledger.contains(UUID), "{unwind:?}: {ledger}");
         for (f, text) in files.iter().zip(&removed) {
-            assert_eq!(
-                &std::fs::read_to_string(root.join(f)).unwrap(),
-                text,
-                "{unwind:?}: {f} stays as uv left it"
-            );
+            let after = std::fs::read_to_string(root.join(f)).unwrap();
+            if transitive {
+                assert!(
+                    !after.contains(UUID) && !after.contains("override"),
+                    "{unwind:?}: {f} is unwired:\n{after}"
+                );
+            } else {
+                assert_eq!(&after, text, "{unwind:?}: {f} stays as uv left it");
+            }
         }
         // `vendor --revert` and `rollback` keep the manifest record, so
         // check then reports the patch as not vendored; the unwinds that

crates/socket-patch-core/src/vendor/pypi_lock.rs

@@ -720,6 +720,74 @@ pub(super) fn still_references_artifact(restored: &str, original: &str, uuid: &s
     restored.contains(&needle) && !original.contains(&needle)
 }
 
+/// Whether `name` (PEP 503) left the live lock `text` altogether: no unit
+/// of that name, and no unit lists it among its `dependencies` (#1287).
+/// For a script lock, `script` is the live script, which must not declare
+/// it either. That is a dependency the user dropped (`uv remove --script`
+/// of the parent of a transitive package), not a hand edit of its unit.
+fn package_vanished(text: &str, name: &str, script: Option<&str>) -> bool {
+    let Ok(doc) = text.parse::<DocumentMut>() else {
+        return false;
+    };
+    let (collection, _) = crate::utils::python_lock::lock_package_collection(&doc);
+    let names = |item: Option<&Item>| -> bool {
+        item.and_then(Item::as_array).is_some_and(|deps| {
+            deps.iter().any(|dep| {
+                dep.as_inline_table()
+                    .and_then(|t| t.get("name"))
+                    .and_then(Value::as_str)
+                    .is_some_and(|n| canonicalize_pypi_name(n) == name)
+            })
+        })
+    };
+    let in_lock = doc
+        .get(collection)
+        .and_then(Item::as_array_of_tables)
+        .is_some_and(|packages| {
+            packages.iter().any(|unit| {
+                unit.get("name")
+                    .and_then(Item::as_str)
+                    .is_some_and(|n| canonicalize_pypi_name(n) == name)
+                    || names(unit.get("dependencies"))
+            })
+        });
+    let declared = script.is_some_and(|script| {
+        script_metadata(script).ok().is_some_and(|(_, metadata)| {
+            metadata.parse::<DocumentMut>().ok().is_some_and(|doc| {
+                doc.get("dependencies")
+                    .and_then(Item::as_array)
+                    .is_some_and(|deps| {
+                        deps.iter().filter_map(Value::as_str).any(|spec| {
+                            canonicalize_pypi_name(crate::vendor::common::pep508_name(spec)) == name
+                        })
+                    })
+            })
+        })
+    });
+    !in_lock && !declared
+}
+
+/// `text` with every lock unit named `name` (PEP 503) dropped: the vendored
+/// unit of a dependency that has since left the lock, which the document
+/// restore then no longer tries to pair (#1287).
+fn without_package(text: &str, name: &str) -> Result<String, String> {
+    let mut doc: DocumentMut = text
+        .parse()
+        .map_err(|error| format!("invalid recorded TOML: {error}"))?;
+    let (collection, _) = crate::utils::python_lock::lock_package_collection(&doc);
+    if let Some(packages) = doc
+        .get_mut(collection)
+        .and_then(Item::as_array_of_tables_mut)
+    {
+        packages.retain(|unit| {
+            unit.get("name")
+                .and_then(Item::as_str)
+                .is_none_or(|n| canonicalize_pypi_name(n) != name)
+        });
+    }
+    Ok(doc.to_string())
+}
+
 pub(super) async fn revert_python_locks(
     entry: &VendorEntry,
     root: &Path,
@@ -752,6 +820,25 @@ pub(super) async fn revert_python_locks(
             Err((_, error)) => return RevertOutcome::failed(error),
         }
     }
+    // The vendored package's PEP 503 name, and each script lock's live
+    // script (read once, before any record is reverted).
+    let package_name = entry
+        .base_purl
+        .strip_prefix("pkg:pypi/")
+        .and_then(|rest| rest.rsplit_once('@'))
+        .map(|(name, _)| canonicalize_pypi_name(name));
+    let mut script_texts = std::collections::BTreeMap::new();
+    for record in entry.wiring.iter().filter(|r| r.kind == KIND) {
+        if let Some(script) = record
+            .file
+            .strip_suffix(".lock")
+            .filter(|s| s.ends_with(".py"))
+        {
+            if let Ok(text) = read_file(&root.join(script)).await {
+                script_texts.insert(record.file.as_str(), text);
+            }
+        }
+    }
     for record in entry.wiring.iter().rev() {
         if !allowed_file(&record.file, &record.kind) {
             warnings.push(VendorWarning::new(
@@ -788,6 +875,40 @@ pub(super) async fn revert_python_locks(
             Ok(live) => live,
             Err((_, error)) => return RevertOutcome::failed(error),
         };
+        // #1287: the vendored package's own unit is gone from the lock along
+        // with every dependent (`uv remove` of the parent of a transitive
+        // package), while the overrides it was wired through survive. That
+        // unit has nothing left to restore: drop it from both recorded
+        // documents so the rest of the record reverts normally.
+        let (original_owned, new_owned);
+        let (original, new) = if record.kind == KIND
+            && package_name.as_deref().is_some_and(|name| {
+                package_vanished(
+                    &live,
+                    name,
+                    script_texts.get(record.file.as_str()).map(String::as_str),
+                )
+            }) {
+            let name = package_name.as_deref().unwrap_or_default();
+            match (without_package(original, name), without_package(new, name)) {
+                (Ok(o), Ok(n)) => {
+                    warnings.push(VendorWarning::new(
+                        super::LOCK_ENTRY_REMOVED_CODE,
+                        format!(
+                            "{}: the {name} unit no longer exists and nothing depends on it \
+                             (the dependency was removed); its other wiring is restored",
+                            record.file
+                        ),
+                    ));
+                    original_owned = o;
+                    new_owned = n;
+                    (original_owned.as_str(), new_owned.as_str())
+                }
+                (Err(error), _) | (_, Err(error)) => return RevertOutcome::failed(error),
+            }
+        } else {
+            (original, new)
+        };
         let restored = if record.kind == SCRIPT_KIND {
             (|| {
                 if live == new || live == original {
@@ -1710,6 +1831,71 @@ mod tests {
         (script.to_string(), lock.to_string())
     }
 
+    /// A script whose `one` arrives only TRANSITIVELY (through `parent`),
+    /// vendored: wired through the script's `[tool.uv]` override + source
+    /// and the lock's `[manifest] overrides`.
+    async fn vendor_transitive_script_pair(root: &Path) -> (VendorEntry, String, String) {
+        let lock = "version = 1\nrevision = 3\nrequires-python = \">=3.9\"\n\n[manifest]\nrequirements = [{name = \"attrs\", specifier = \">=20\"}, {name = \"parent\", specifier = \"==1\"}]\n\n[[package]]\nname = \"attrs\"\nversion = \"25.3.0\"\nsource = {registry = \"https://pypi.org/simple\"}\n\n[[package]]\nname = \"one\"\nversion = \"1\"\nsource = {registry = \"https://pypi.org/simple\"}\n\n[[package]]\nname = \"parent\"\nversion = \"1\"\nsource = {registry = \"https://pypi.org/simple\"}\ndependencies = [{name = \"one\"}]\n";
+        let script = "# /// script\n# requires-python = \">=3.9\"\n# dependencies = [\"parent==1\", \"attrs>=20\"]\n# ///\nimport one\n";
+        write_pylock(root, "job.py.lock", lock).await;
+        tokio::fs::write(root.join("job.py"), script).await.unwrap();
+        let project = load_python_locks(root, "one", "1", UUID).await.unwrap();
+        let wheel =
+            ".socket/vendor/pypi/11111111-1111-4111-8111-111111111111/one-1-py3-none-any.whl";
+        let records = wire_python_locks(&project, root, "one", "1", wheel, &"a".repeat(64))
+            .await
+            .unwrap();
+        let entry: VendorEntry = serde_json::from_value(serde_json::json!({
+            "ecosystem": "pypi",
+            "basePurl": "pkg:pypi/one@1",
+            "uuid": UUID,
+            "artifact": { "path": wheel, "sha256": "a".repeat(64) },
+            "wiring": serde_json::to_value(&records).unwrap(),
+            "flavor": "python-lock",
+        }))
+        .unwrap();
+        let wired_script = std::fs::read_to_string(root.join("job.py")).unwrap();
+        let wired_lock = std::fs::read_to_string(root.join("job.py.lock")).unwrap();
+        (entry, wired_script, wired_lock)
+    }
+
+    /// #1287: `uv remove --script job.py parent` drops the transitive
+    /// `one`'s unit (and its parent's) but keeps the script's `[tool.uv]`
+    /// override + source and the lock's `[manifest] overrides`, which still
+    /// name the uuid. The vanished unit is removed, not drift, and the
+    /// surviving wiring reverts, so the script and lock end up as uv writes
+    /// them for the project without `one`.
+    #[tokio::test]
+    async fn script_revert_after_uv_remove_of_the_transitive_parent() {
+        let temp = tempfile::tempdir().unwrap();
+        let root = temp.path();
+        let (entry, wired_script, wired_lock) = vendor_transitive_script_pair(root).await;
+        let removed_script = wired_script.replace("\"parent==1\", ", "");
+        let mut removed_lock = wired_lock.replace(", {name = \"parent\", specifier = \"==1\"}", "");
+        let one = removed_lock.find("\n[[package]]\nname = \"one\"").unwrap();
+        removed_lock.truncate(one);
+        tokio::fs::write(root.join("job.py"), &removed_script)
+            .await
+            .unwrap();
+        write_pylock(root, "job.py.lock", &removed_lock).await;
+
+        let outcome = revert_python_locks(&entry, root, false).await;
+        assert!(outcome.success, "{outcome:?}");
+        assert!(!outcome.drift_skipped(), "{:?}", outcome.warnings);
+        assert!(outcome.lock_entry_removed(), "{:?}", outcome.warnings);
+        let script = std::fs::read_to_string(root.join("job.py")).unwrap();
+        let lock = std::fs::read_to_string(root.join("job.py.lock")).unwrap();
+        assert_eq!(
+            script,
+            "# /// script\n# requires-python = \">=3.9\"\n# dependencies = [\"attrs>=20\"]\n# ///\nimport one\n"
+        );
+        assert!(
+            !lock.contains(UUID) && !lock.contains("overrides"),
+            "{lock}"
+        );
+        assert!(lock.contains("name = \"attrs\""), "{lock}");
+    }
+
     /// #1214: after `uv remove --script` drops the vendored dependency from a
     /// PEP 723 script and its lock, the revert has nothing left to restore.
     /// It must warn `vendor_lock_entry_removed` and finish (the caller then